Automatically test connection on admin page

When admin visits "GitLab Duo Pro" admin page, we automatically run a
simple test completion request to check if code suggestions request work
as expected and show this result as a flash message.

ee/app/controllers/admin/code_suggestions_controller.rb

@@ -13,6 +13,14 @@ class CodeSuggestionsController < Admin::ApplicationController
     before_action :ensure_feature_available!
 
     def index
+      error = ::Gitlab::Llm::AiGateway::CodeSuggestionsClient.new(current_user).test_completion
+
+      if error.blank?
+        flash[:notice] = _('Code completion test was successful')
+      else
+        flash[:alert] = format(_('Code completion test failed: %{error}'), error: error)
+      end
+
       @subscription_name = License.current.subscription_name
     end
 

ee/lib/gitlab/llm/ai_gateway/code_suggestions_client.rb

+# frozen_string_literal: true
+
+module Gitlab
+  module Llm
+    module AiGateway
+      class CodeSuggestionsClient
+        include ::Gitlab::Utils::StrongMemoize
+        include ::API::Helpers::CloudConnector
+
+        COMPLETION_CHECK_TIMEOUT = 3.seconds
+
+        def initialize(user)
+          @user = user
+        end
+
+        def test_completion
+          return 'Access token is missing' unless access_token
+
+          response = Gitlab::HTTP.post(
+            task.endpoint,
+            headers: request_headers,
+            body: task.body,
+            timeout: COMPLETION_CHECK_TIMEOUT,
+            allow_local_requests: true
+          )
+
+          return "AI Gateway returned code #{response.code}: #{response.body}" unless response.code == 200
+          return "Response doesn't contain a completion" unless choice?(response)
+
+          nil
+        rescue StandardError => err
+          err.message
+        end
+
+        private
+
+        attr_reader :user
+
+        def request_headers
+          {
+            'X-Gitlab-Authentication-Type' => 'oidc',
+            'Authorization' => "Bearer #{access_token}",
+            'Content-Type' => 'application/json',
+            'X-Request-ID' => Labkit::Correlation::CorrelationId.current_or_new_id
+          }.merge(cloud_connector_headers(user))
+        end
+
+        def access_token
+          ::CloudConnector::AvailableServices.find_by_name(:code_suggestions).access_token(user)
+        end
+        strong_memoize_attr :access_token
+
+        def choice?(response)
+          response['choices']&.first&.dig('text').present?
+        end
+
+        def task
+          params = {
+            current_file: {
+              file_name: 'test.rb',
+              content_above_cursor: 'def hello_world'
+            }
+          }
+          CodeSuggestions::Tasks::CodeCompletion.new(unsafe_passthrough_params: params)
+        end
+        strong_memoize_attr :task
+      end
+    end
+  end
+end

ee/spec/lib/gitlab/llm/ai_gateway/code_suggestions_client_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Gitlab::Llm::AiGateway::CodeSuggestionsClient, feature_category: :code_suggestions do
+  let_it_be(:user) { create(:user) }
+
+  describe "#test_completion" do
+    let_it_be(:token) { create(:service_access_token, :active) }
+    let(:body) { { choices: [{ text: "puts \"Hello World!\"\nend", index: 0, finish_reason: "length" }] } }
+    let(:code) { 200 }
+
+    subject(:result) { described_class.new(user).test_completion }
+
+    shared_examples "error response" do |message|
+      it "returns an error" do
+        expect(result).to eq(message)
+      end
+    end
+
+    before do
+      stub_request(:post, /#{Gitlab::AiGateway.url}/)
+        .to_return(status: code, body: body.to_json, headers: { "Content-Type" => "application/json" })
+      allow(CloudConnector::AvailableServices).to receive_message_chain(:find_by_name,
+        :access_token).and_return(token)
+    end
+
+    it 'returns nil if there is no error' do
+      expect(result).to be_nil
+    end
+
+    context 'when there is not valid token' do
+      let(:token) { nil }
+
+      it_behaves_like 'error response', "Access token is missing"
+    end
+
+    context 'when response does not contain a valid choice' do
+      let(:body) { { choices: [] } }
+
+      it_behaves_like 'error response', "Response doesn't contain a completion"
+    end
+
+    context 'when response code is not 200' do
+      let(:code) { 401 }
+      let(:body) { 'an error' }
+
+      it_behaves_like 'error response', 'AI Gateway returned code 401: "an error"'
+    end
+
+    context 'when request raises an error' do
+      before do
+        stub_request(:post, /#{Gitlab::AiGateway.url}/).to_raise(StandardError.new('an error'))
+      end
+
+      it_behaves_like 'error response', 'an error'
+    end
+  end
+end

ee/spec/requests/admin/code_suggestions_controller_spec.rb

@@ -15,11 +15,36 @@
     end
 
     shared_examples 'renders the activation form' do
-      it 'renders the activation form' do
-        get admin_code_suggestions_path
+      context 'when connection check succeeds' do
+        before do
+          allow_next_instance_of(Gitlab::Llm::AiGateway::CodeSuggestionsClient) do |client|
+            allow(client).to receive(:test_completion).and_return(nil)
+          end
+        end
+
+        it 'renders the activation form' do
+          get admin_code_suggestions_path
+
+          expect(response).to render_template(:index)
+          expect(response.body).to include('js-code-suggestions-page')
+          expect(flash[:notice]).to eq("Code completion test was successful")
+        end
+      end
 
-        expect(response).to render_template(:index)
-        expect(response.body).to include('js-code-suggestions-page')
+      context 'when connection check fails' do
+        before do
+          allow_next_instance_of(Gitlab::Llm::AiGateway::CodeSuggestionsClient) do |client|
+            allow(client).to receive(:test_completion).and_return('an error')
+          end
+        end
+
+        it 'renders the activation form with alert message' do
+          get admin_code_suggestions_path
+
+          expect(response).to render_template(:index)
+          expect(response.body).to include('js-code-suggestions-page')
+          expect(flash[:alert]).to eq("Code completion test failed: an error")
+        end
       end
     end
 

locale/gitlab.pot

@@ -12583,6 +12583,12 @@ msgstr ""
 msgid "Code block"
 msgstr ""
+msgid "Code completion test failed: %{error}"
+msgstr ""
+
+msgid "Code completion test was successful"
+msgstr ""
+
 msgid "Code coverage statistics for %{ref} %{start_date} - %{end_date}"
 msgstr ""