diff --git a/ee/lib/api/code_suggestions.rb b/ee/lib/api/code_suggestions.rb
index c48ce5e035acf3d33b1c288303a84337dd852eeb..382d9719207e5c984eff892216927ca4031fafa4 100644
--- a/ee/lib/api/code_suggestions.rb
+++ b/ee/lib/api/code_suggestions.rb
@@ -6,7 +6,7 @@ class CodeSuggestions < ::API::Base
feature_category :code_suggestions
- helpers ::API::Helpers::GlobalIds
+ helpers ::API::Helpers::CloudConnector
# a limit used for overall body size when forwarding request to ai-assist, overall size should not be bigger than
# summary of limits on accepted parameters
@@ -31,17 +31,14 @@ def user_allowed?
def model_gateway_headers(headers, gateway_token)
telemetry_headers = headers.select { |k| /\Ax-gitlab-cs-/i.match?(k) }
- instance_id, user_id = global_instance_and_user_id_for(current_user)
{
- 'X-Gitlab-Instance-Id' => instance_id,
- 'X-Gitlab-Global-User-Id' => user_id,
'X-Gitlab-Host-Name' => Gitlab.config.gitlab.host,
- 'X-Gitlab-Realm' => gitlab_realm,
'X-Gitlab-Authentication-Type' => 'oidc',
'Authorization' => "Bearer #{gateway_token}",
'Content-Type' => 'application/json',
'User-Agent' => headers["User-Agent"] # Forward the User-Agent on to the model gateway
- }.merge(telemetry_headers).merge(saas_headers).transform_values { |v| Array(v) }
+ }.merge(telemetry_headers).merge(saas_headers).merge(cloud_connector_headers(current_user))
+ .transform_values { |v| Array(v) }
end
def saas_headers
@@ -75,9 +72,7 @@ def gitlab_realm
# See https://gitlab.com/groups/gitlab-org/-/epics/11114
return Gitlab::Ai::AccessToken::GITLAB_REALM_SELF_MANAGED if proxied?
- return Gitlab::Ai::AccessToken::GITLAB_REALM_SAAS if Gitlab.org_or_com?
-
- Gitlab::Ai::AccessToken::GITLAB_REALM_SELF_MANAGED
+ super
end
end
diff --git a/ee/lib/api/helpers/cloud_connector.rb b/ee/lib/api/helpers/cloud_connector.rb
new file mode 100644
index 0000000000000000000000000000000000000000..0656872a2b2f82554c686dbc2c26456bc5415553
--- /dev/null
+++ b/ee/lib/api/helpers/cloud_connector.rb
@@ -0,0 +1,25 @@
+# frozen_string_literal: true
+
+module API
+ module Helpers
+ module CloudConnector
+ include ::API::Helpers::GlobalIds
+
+ def cloud_connector_headers(user)
+ instance_id, user_id = global_instance_and_user_id_for(user)
+
+ {
+ 'X-Gitlab-Instance-Id' => instance_id,
+ 'X-Gitlab-Global-User-Id' => user_id,
+ 'X-Gitlab-Realm' => gitlab_realm
+ }
+ end
+
+ def gitlab_realm
+ return Gitlab::Ai::AccessToken::GITLAB_REALM_SAAS if Gitlab.org_or_com? # rubocop:disable Gitlab/AvoidGitlabInstanceChecks -- To align with ee/lib/api/code_suggestions.rb.
+
+ Gitlab::Ai::AccessToken::GITLAB_REALM_SELF_MANAGED
+ end
+ end
+ end
+end
diff --git a/ee/lib/api/internal/ai/x_ray/scan.rb b/ee/lib/api/internal/ai/x_ray/scan.rb
index 8cdb37a6fc86a02ec34253ec950204e53338a366..d99263afcae66310b49bd0c6c4511f8647296fdc 100644
--- a/ee/lib/api/internal/ai/x_ray/scan.rb
+++ b/ee/lib/api/internal/ai/x_ray/scan.rb
@@ -10,6 +10,7 @@ class Scan < ::API::Base
feature_category :code_suggestions
helpers ::API::Ci::Helpers::Runner
+ helpers ::API::Helpers::CloudConnector
before do
authenticate_job!
@@ -67,12 +68,6 @@ def saas_headers
}
end
- def gitlab_realm
- return Gitlab::Ai::AccessToken::GITLAB_REALM_SAAS if Gitlab.org_or_com?
-
- Gitlab::Ai::AccessToken::GITLAB_REALM_SELF_MANAGED
- end
-
def ai_access_token
::Ai::ServiceAccessToken.active.last
end
diff --git a/ee/lib/gitlab/llm/ai_gateway/client.rb b/ee/lib/gitlab/llm/ai_gateway/client.rb
index c835a24c41e916188d05de5005255c2a44e3b4a7..5b4d94aecb9f615c927d29b136986695ff865c85 100644
--- a/ee/lib/gitlab/llm/ai_gateway/client.rb
+++ b/ee/lib/gitlab/llm/ai_gateway/client.rb
@@ -6,8 +6,8 @@ module AiGateway
class Client
include ::Gitlab::Llm::Concerns::ExponentialBackoff
include ::Gitlab::Llm::Concerns::EventTracking
- include ::API::Helpers::GlobalIds
include ::Gitlab::Utils::StrongMemoize
+ include ::API::Helpers::CloudConnector
DEFAULT_PROVIDER = 'anthropic'
DEFAULT_MODEL = 'claude-2.0'
@@ -84,23 +84,13 @@ def enabled?
end
def request_headers
- instance_id, user_id = global_instance_and_user_id_for(user)
{
- 'X-Gitlab-Instance-Id' => instance_id,
- 'X-Gitlab-Global-User-Id' => user_id,
'X-Gitlab-Host-Name' => Gitlab.config.gitlab.host,
- 'X-Gitlab-Realm' => gitlab_realm,
'X-Gitlab-Authentication-Type' => 'oidc',
'Authorization' => "Bearer #{access_token}",
'Content-Type' => 'application/json',
'X-Request-ID' => Labkit::Correlation::CorrelationId.current_or_new_id
- }
- end
-
- def gitlab_realm
- return Gitlab::Ai::AccessToken::GITLAB_REALM_SAAS if Gitlab.org_or_com? # rubocop:disable Gitlab/AvoidGitlabInstanceChecks -- To align with ee/lib/api/code_suggestions.rb.
-
- Gitlab::Ai::AccessToken::GITLAB_REALM_SELF_MANAGED
+ }.merge(cloud_connector_headers(user))
end
def access_token
diff --git a/ee/spec/lib/api/helpers/cloud_connector_spec.rb b/ee/spec/lib/api/helpers/cloud_connector_spec.rb
new file mode 100644
index 0000000000000000000000000000000000000000..c625426ca79eb343d59fa90df57a223770b69b52
--- /dev/null
+++ b/ee/spec/lib/api/helpers/cloud_connector_spec.rb
@@ -0,0 +1,39 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe API::Helpers::CloudConnector, feature_category: :cloud_connector do
+ subject(:helper) do
+ Class.new do
+ include API::Helpers::CloudConnector
+ end.new
+ end
+
+ let_it_be(:user) { build(:user, id: 1) }
+
+ describe '#cloud_connector_headers' do
+ it 'generates a hash with the required fields based on the user' do
+ expect(helper.cloud_connector_headers(user)).to match(
+ {
+ 'X-Gitlab-Instance-Id' => an_instance_of(String),
+ 'X-Gitlab-Global-User-Id' => an_instance_of(String),
+ 'X-Gitlab-Realm' => Gitlab::Ai::AccessToken::GITLAB_REALM_SELF_MANAGED
+ }
+ )
+ end
+ end
+
+ describe '#gitlab_realm' do
+ context 'when the current instance is saas', :saas do
+ it 'returns Gitlab::Ai::AccessToken::GITLAB_REALM_SAAS' do
+ expect(helper.gitlab_realm).to eq(Gitlab::Ai::AccessToken::GITLAB_REALM_SAAS)
+ end
+ end
+
+ context 'when the current instance is not saas' do
+ it 'returns Gitlab::Ai::AccessToken::GITLAB_REALM_SELF_MANAGED' do
+ expect(helper.gitlab_realm).to eq(Gitlab::Ai::AccessToken::GITLAB_REALM_SELF_MANAGED)
+ end
+ end
+ end
+end