Add probe for checks in air-gapped instances

Add probe for checks in air-gapped instances
that self-host AI Gateway

Changelog: changed
EE: true

doc/user/gitlab_duo/turn_on_off.md

@@ -90,7 +90,10 @@ To run a health check:
 
 ### Health check tests
 
-The health check performs the following tests to verify if your instance meets the requirements to use GitLab Duo.
+To verify if your instance meets the requirements to use GitLab Duo, the health check performs tests
+for online and offline environments.
+
+#### For online environments
 
 | Test | Description |
 |-----------------|-------------|
@@ -98,6 +101,14 @@ The health check performs the following tests to verify if your instance meets t
 | Synchronization | Tests whether your subscription: <br>- Has been activated with an activation code and can be synchronized with `customers.gitlab.com`.<br>- Has correct access credentials.<br>- Has been synchronized recently. If it hasn't or the access credentials are missing or expired, you can [manually synchronize](../../subscriptions/self_managed/index.md#manually-synchronize-subscription-data) your subscription data. |
 | System exchange | Tests whether Code Suggestions can be used in your instance. If the system exchange assessment fails, users might not be able to use GitLab Duo features. |
 
+#### For offline environments
+
+| Test | Description |
+|-----------------|-------------|
+| Network | Tests whether: <br>- The environment variable `AI_GATEWAY_URL` has been set to a valid URL.<br> - Your instance can connect to the URL specified by `AI_GATEWAY_URL`.<br><br>If your instance cannot connect to the URL, ensure that your firewall or proxy server settings [allow connection](#configure-gitlab-duo-on-a-self-managed-instance). |
+| License | Tests whether your license has the ability to access Code Suggestions feature. |
+| System exchange | Tests whether Code Suggestions can be used in your instance. If the system exchange assessment fails, users might not be able to use GitLab Duo features. |
+
 ## Turn off GitLab Duo features
 
 You can turn off GitLab Duo for a group, project, or instance.

ee/app/services/cloud_connector/status_checks/probes/end_to_end_probe.rb

@@ -20,7 +20,7 @@ def initialize(user)
 
         override :success_message
         def success_message
-          _('Authentication with GitLab Cloud services succeeded.')
+          _('Authentication with the AI gateway services succeeded.')
         end
 
         def check_user_exists
@@ -35,7 +35,7 @@ def validate_code_completion_availability
         end
 
         def failure_text(error)
-          format(_('Authentication with GitLab Cloud services failed: %{error}'), error: error)
+          format(_('Authentication with the AI gateway services failed: %{error}'), error: error)
         end
       end
     end

ee/app/services/cloud_connector/status_checks/probes/host_probe.rb

@@ -10,16 +10,32 @@ class HostProbe < BaseProbe
 
         attr_reader :host, :port
 
-        validate :validate_connection
+        validate :validate_connection, if: :prerequisites_for_valid_url_met?
 
         def initialize(service_url)
-          uri = URI.parse(service_url)
+          @service_url = service_url
+
+          return if @service_url.blank?
+
+          uri = URI.parse(@service_url)
           @host = uri.host
           @port = uri.port
         end
 
         private
 
+        def prerequisites_for_valid_url_met?
+          return true if @host.present? && @port.present?
+
+          if @service_url.present?
+            errors.add(:base, format(_('%{service_url} is not a valid URL.'), service_url: @service_url))
+          else
+            errors.add(:base, _('Cannot validate connection to host because the URL is empty.'))
+          end
+
+          false
+        end
+
         override :success_message
         def success_message
           format(_('%{host} reachable.'), host: @host)

ee/app/services/cloud_connector/status_checks/probes/self_hosted/ai_gateway_url_presence_probe.rb

+# frozen_string_literal: true
+
+module CloudConnector
+  module StatusChecks
+    module Probes
+      module SelfHosted
+        class AiGatewayUrlPresenceProbe < BaseProbe
+          extend ::Gitlab::Utils::Override
+
+          ENV_VARIABLE_NAME = 'AI_GATEWAY_URL'
+
+          validate :check_ai_gateway_url_presence
+
+          private
+
+          def self_hosted_url
+            ::Gitlab::AiGateway.self_hosted_url
+          end
+
+          def check_ai_gateway_url_presence
+            return if self_hosted_url.present?
+
+            errors.add(:base, failure_message)
+          end
+
+          override :success_message
+          def success_message
+            format(
+              _("Environment variable %{env_variable_name} is set to %{url}."),
+              env_variable_name: ENV_VARIABLE_NAME,
+              url: self_hosted_url
+            )
+          end
+
+          def failure_message
+            format(
+              _("Environment variable %{env_variable_name} is not set."),
+              env_variable_name: ENV_VARIABLE_NAME
+            )
+          end
+        end
+      end
+    end
+  end
+end

ee/app/services/cloud_connector/status_checks/probes/self_hosted/code_suggestions_license_probe.rb

+# frozen_string_literal: true
+
+module CloudConnector
+  module StatusChecks
+    module Probes
+      module SelfHosted
+        class CodeSuggestionsLicenseProbe < BaseProbe
+          extend ::Gitlab::Utils::Override
+
+          validate :check_user_exists
+          validate :validate_code_suggestions_availability
+
+          after_validation :collect_instance_details, :collect_license_details
+
+          def initialize(user)
+            @user = user
+          end
+
+          private
+
+          attr_reader :user
+
+          def check_user_exists
+            errors.add(:base, 'User not provided.') unless user
+          end
+
+          override :success_message
+          def success_message
+            _('License includes access to Code Suggestions.')
+          end
+
+          def validate_code_suggestions_availability
+            return unless user
+            return if Ability.allowed?(user, :access_code_suggestions)
+
+            if ::License.feature_available?(:code_suggestions)
+              text = _(
+                'License includes access to Code Suggestions, but you lack the necessary ' \
+                  'permissions to use this feature.'
+              )
+
+              errors.add(:base, text)
+
+              return
+            end
+
+            errors.add(:base, _('License does not provide access to Code Suggestions.'))
+          end
+
+          def collect_instance_details
+            details.add(:instance_id, Gitlab::GlobalAnonymousId.instance_id)
+            details.add(:gitlab_version, Gitlab::VERSION)
+          end
+
+          def collect_license_details
+            return unless license
+
+            details.add(:license, license.license.as_json)
+          end
+
+          def license
+            @license ||= License.current
+          end
+        end
+      end
+    end
+  end
+end

ee/app/services/cloud_connector/status_checks/status_service.rb

@@ -13,7 +13,7 @@ class StatusService
 
       def initialize(user:, probes: nil)
         @user = user
-        @probes = probes || build_default_probes
+        @probes = probes || selected_probes
       end
 
       def execute
@@ -28,7 +28,17 @@ def execute
 
       private
 
-      def build_default_probes
+      def selected_probes
+        # An air-gapped instance, which requires that they run their own self-hosted AI Gateway,
+        # requires a different set of probes to be executed.
+        if ::Gitlab::Ai::SelfHosted::AiGateway.required?
+          ::Gitlab::Ai::SelfHosted::AiGateway.probes(@user)
+        else
+          default_probes
+        end
+      end
+
+      def default_probes
         [
           CloudConnector::StatusChecks::Probes::LicenseProbe.new,
           CloudConnector::StatusChecks::Probes::HostProbe.new(CUSTOMERS_DOT_URL),

ee/lib/gitlab/ai/self_hosted/ai_gateway.rb

+# frozen_string_literal: true
+
+module Gitlab
+  module Ai
+    module SelfHosted
+      module AiGateway
+        extend self
+
+        # An instance having an offline cloud license is
+        # supposed to be an air-gapped instance.
+        # Air-gapped instances cannot connect to GitLab's default CloudConnector
+        # and are hence required to self-host their own AI Gateway (and the models)
+        def required?
+          ::Feature.enabled?(:ai_custom_model) && # rubocop:disable Gitlab/FeatureFlagWithoutActor -- The feature flag is global
+            ::License.current&.offline_cloud_license?
+        end
+
+        def probes(user)
+          [
+            ::CloudConnector::StatusChecks::Probes::SelfHosted::AiGatewayUrlPresenceProbe.new,
+            ::CloudConnector::StatusChecks::Probes::HostProbe.new(::Gitlab::AiGateway.self_hosted_url),
+            ::CloudConnector::StatusChecks::Probes::SelfHosted::CodeSuggestionsLicenseProbe.new(user),
+            ::CloudConnector::StatusChecks::Probes::EndToEndProbe.new(user)
+          ]
+        end
+      end
+    end
+  end
+end

ee/spec/lib/gitlab/ai/self_hosted/ai_gateway_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Gitlab::Ai::SelfHosted::AiGateway, feature_category: :"self-hosted_models" do
+  describe '.required?' do
+    context 'when the license is not an offline cloud license' do
+      it 'returns false' do
+        expect(described_class.required?).to be(false)
+      end
+    end
+
+    context 'when the license is an offline cloud license' do
+      before do
+        allow(::License).to receive_message_chain(:current, :offline_cloud_license?).and_return(true)
+      end
+
+      it 'returns true' do
+        expect(described_class.required?).to be(true)
+      end
+
+      context 'when the feature flag :ai_custom_model is disabled' do
+        it 'returns false' do
+          stub_feature_flags(ai_custom_model: false)
+          expect(described_class.required?).to be(false)
+        end
+      end
+    end
+  end
+
+  describe '.probes' do
+    let(:user) { build(:user) }
+
+    it 'returns an array with all expected probe instances' do
+      probes = described_class.probes(user)
+
+      expect(probes).to contain_exactly(
+        an_instance_of(::CloudConnector::StatusChecks::Probes::SelfHosted::AiGatewayUrlPresenceProbe),
+        an_instance_of(::CloudConnector::StatusChecks::Probes::HostProbe),
+        an_instance_of(::CloudConnector::StatusChecks::Probes::SelfHosted::CodeSuggestionsLicenseProbe),
+        an_instance_of(::CloudConnector::StatusChecks::Probes::EndToEndProbe)
+      )
+    end
+  end
+end

ee/spec/services/cloud_connector/status_checks/probes/end_to_end_probe_spec.rb

@@ -29,7 +29,7 @@
         result = probe.execute
 
         expect(result.success).to be true
-        expect(result.message).to match('Authentication with GitLab Cloud services succeeded')
+        expect(result.message).to match('Authentication with the AI gateway services succeeded')
       end
     end
 
@@ -46,7 +46,7 @@
         result = probe.execute
 
         expect(result.success).to be false
-        expect(result.message).to match("Authentication with GitLab Cloud services failed: #{error_message}")
+        expect(result.message).to match("Authentication with the AI gateway services failed: #{error_message}")
       end
     end
   end

ee/spec/services/cloud_connector/status_checks/probes/host_probe_spec.rb

@@ -8,6 +8,30 @@
 
     let(:uri) { 'https://example.com' }
 
+    context 'when the host is nil' do
+      let(:uri) { nil }
+
+      it 'returns a failure result' do
+        result = probe.execute
+
+        expect(result).to be_a(CloudConnector::StatusChecks::Probes::ProbeResult)
+        expect(result.success?).to be false
+        expect(result.message).to match("Cannot validate connection to host because the URL is empty.")
+      end
+    end
+
+    context 'when the host is not a valid URL' do
+      let(:uri) { 'not_a_valid_url' }
+
+      it 'returns a failure result' do
+        result = probe.execute
+
+        expect(result).to be_a(CloudConnector::StatusChecks::Probes::ProbeResult)
+        expect(result.success?).to be false
+        expect(result.message).to match("not_a_valid_url is not a valid URL.")
+      end
+    end
+
     context 'when the host is reachable' do
       before do
         allow(TCPSocket).to receive(:new).and_return(instance_double(TCPSocket, close: nil))

ee/spec/services/cloud_connector/status_checks/probes/self_hosted/ai_gateway_url_presence_probe_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe CloudConnector::StatusChecks::Probes::SelfHosted::AiGatewayUrlPresenceProbe, feature_category: :cloud_connector do
+  let(:probe) { described_class.new }
+
+  describe '#execute' do
+    context 'when AI_GATEWAY_URL is set' do
+      before do
+        stub_env('AI_GATEWAY_URL', 'https://ai-gateway.mycompany.com')
+      end
+
+      it 'returns a successful result' do
+        result = probe.execute
+        expect(result).to be_a(CloudConnector::StatusChecks::Probes::ProbeResult)
+        expect(result.success?).to be(true)
+        expect(result.message).to eq('Environment variable AI_GATEWAY_URL is set to https://ai-gateway.mycompany.com.')
+      end
+    end
+
+    context 'when AI_GATEWAY_URL is not set' do
+      before do
+        stub_env('AI_GATEWAY_URL', nil)
+      end
+
+      it 'returns a failed result' do
+        result = probe.execute
+        expect(result).to be_a(CloudConnector::StatusChecks::Probes::ProbeResult)
+        expect(result.success?).to be(false)
+        expect(result.message).to eq('Environment variable AI_GATEWAY_URL is not set.')
+      end
+    end
+  end
+end

ee/spec/services/cloud_connector/status_checks/probes/self_hosted/code_suggestions_license_probe_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe CloudConnector::StatusChecks::Probes::SelfHosted::CodeSuggestionsLicenseProbe, feature_category: :cloud_connector do
+  let(:probe) { described_class.new(user) }
+  let(:user) { build(:user) }
+
+  describe '#execute' do
+    context 'when user has access to code suggestions' do
+      before do
+        allow(Ability).to receive(:allowed?).with(user, :access_code_suggestions).and_return(true)
+      end
+
+      it 'returns a success result' do
+        result = probe.execute
+
+        expect(result.success).to be true
+        expect(result.message).to match('License includes access to Code Suggestions.')
+      end
+    end
+
+    context 'when user does not have access to code suggestions' do
+      before do
+        stub_licensed_features(code_suggestions: true)
+        allow(Ability).to receive(:allowed?).with(user, :access_code_suggestions).and_return(false)
+      end
+
+      it 'returns a failure result' do
+        result = probe.execute
+
+        expect(result.success).to be false
+        expect(result.message).to match(
+          'License includes access to Code Suggestions, but you lack the necessary ' \
+            'permissions to use this feature.'
+        )
+      end
+    end
+
+    context 'when license does not provide access to code suggestions' do
+      before do
+        stub_licensed_features(code_suggestions: false)
+      end
+
+      it 'returns a failure result' do
+        result = probe.execute
+
+        expect(result.success).to be false
+        expect(result.message).to match('License does not provide access to Code Suggestions.')
+      end
+    end
+
+    context 'on collecting details' do
+      let(:license) { build(:license, cloud: false) }
+
+      before do
+        allow(License).to receive(:current).and_return(license)
+      end
+
+      it 'collects the instance details' do
+        result = probe.execute
+
+        expect(result.details[:instance_id]).to eq(Gitlab::GlobalAnonymousId.instance_id)
+        expect(result.details[:gitlab_version]).to eq(Gitlab::VERSION)
+      end
+
+      it 'collects the license details' do
+        result = probe.execute
+
+        expect(result.details[:license]).to eq(License.current.license.as_json)
+      end
+    end
+  end
+end

ee/spec/services/cloud_connector/status_checks/status_service_spec.rb

@@ -11,9 +11,9 @@
   subject(:service) { described_class.new(user: user, probes: probes) }
 
   describe '#initialize' do
-    context 'when no probes are passed' do
-      subject(:service) { described_class.new(user: user) }
+    subject(:service) { described_class.new(user: user) }
 
+    context 'when no probes are passed' do
       it 'created default probes' do
         service_probes = service.probes
 
@@ -26,6 +26,26 @@
         expect(service_probes[5]).to be_an_instance_of(CloudConnector::StatusChecks::Probes::EndToEndProbe)
       end
     end
+
+    context 'when self-hosted AI Gateway is required' do
+      before do
+        allow(::Gitlab::Ai::SelfHosted::AiGateway).to receive(:required?).and_return(true)
+      end
+
+      it 'uses a different set of probes' do
+        service_probes = service.probes
+
+        expect(service_probes.count).to eq(4)
+        expect(service_probes[0]).to be_an_instance_of(
+          CloudConnector::StatusChecks::Probes::SelfHosted::AiGatewayUrlPresenceProbe
+        )
+        expect(service_probes[1]).to be_an_instance_of(CloudConnector::StatusChecks::Probes::HostProbe)
+        expect(service_probes[2]).to be_an_instance_of(
+          CloudConnector::StatusChecks::Probes::SelfHosted::CodeSuggestionsLicenseProbe
+        )
+        expect(service_probes[3]).to be_an_instance_of(CloudConnector::StatusChecks::Probes::EndToEndProbe)
+      end
+    end
   end
 
   describe '#execute' do

locale/gitlab.pot

@@ -1254,6 +1254,9 @@ msgid_plural "%{selectedProjectsCount} projects"
 msgstr[0] ""
 msgstr[1] ""
+msgid "%{service_url} is not a valid URL."
+msgstr ""
+
 msgid "%{size} B"
 msgstr ""
@@ -8028,10 +8031,10 @@ msgstr ""
 msgid "Authentication via WebAuthn device failed."
 msgstr ""
-msgid "Authentication with GitLab Cloud services failed: %{error}"
+msgid "Authentication with the AI gateway services failed: %{error}"
 msgstr ""
-msgid "Authentication with GitLab Cloud services succeeded."
+msgid "Authentication with the AI gateway services succeeded."
 msgstr ""
 msgid "Author"
@@ -10892,6 +10895,9 @@ msgstr ""
 msgid "Cannot skip two factor authentication setup"
 msgstr ""
+msgid "Cannot validate connection to host because the URL is empty."
+msgstr ""
+
 msgid "Capacity threshold"
 msgstr ""
@@ -21272,6 +21278,12 @@ msgstr ""
 msgid "Environment scope"
 msgstr ""
+msgid "Environment variable %{env_variable_name} is not set."
+msgstr ""
+
+msgid "Environment variable %{env_variable_name} is set to %{url}."
+msgstr ""
+
 msgid "Environment variables on this GitLab instance are configured to be %{help_link_start}protected%{help_link_end} by default."
 msgstr ""
@@ -32119,6 +32131,15 @@ msgstr ""
 msgid "License Compliance| Used by %{dependencies}"
 msgstr ""
+msgid "License does not provide access to Code Suggestions."
+msgstr ""
+
+msgid "License includes access to Code Suggestions, but you lack the necessary permissions to use this feature."
+msgstr ""
+
+msgid "License includes access to Code Suggestions."
+msgstr ""
+
 msgid "License key"
 msgstr ""