diff --git a/.rubocop_todo/rspec/be_eq.yml b/.rubocop_todo/rspec/be_eq.yml
index 7cd0920d3f8028acc178631c1065396781dea4f7..a7a667947ce03c42db686fc266c682b40696771e 100644
--- a/.rubocop_todo/rspec/be_eq.yml
+++ b/.rubocop_todo/rspec/be_eq.yml
@@ -92,7 +92,6 @@ RSpec/BeEq:
- 'ee/spec/lib/arkose/verify_response_spec.rb'
- 'ee/spec/lib/bulk_imports/groups/pipelines/epics_pipeline_spec.rb'
- 'ee/spec/lib/bulk_imports/groups/pipelines/iterations_cadences_pipeline_spec.rb'
- - 'ee/spec/lib/code_suggestions/completions_model_details_spec.rb'
- 'ee/spec/lib/code_suggestions/tasks/base_spec.rb'
- 'ee/spec/lib/code_suggestions/tasks/code_completion_spec.rb'
- 'ee/spec/lib/code_suggestions/tasks/code_generation_spec.rb'
diff --git a/ee/app/models/concerns/ai/user_authorizable.rb b/ee/app/models/concerns/ai/user_authorizable.rb
index cd09509b11b06ddb24474a8564208cfde561e1a7..c4ec64fd079a6a36e0e0b3097a5e61b9984831f3 100644
--- a/ee/app/models/concerns/ai/user_authorizable.rb
+++ b/ee/app/models/concerns/ai/user_authorizable.rb
@@ -134,6 +134,12 @@ def allowed_to_use?(ai_feature, service_name: nil, licensed_feature: :ai_feature
end
end
+ def allowed_by_namespace_ids(*args, **kwargs)
+ allowed_to_use?(*args, **kwargs) { |namespace_ids| return namespace_ids } # rubocop:disable Cop/AvoidReturnFromBlocks -- Early return if namespace ids are yielded
+
+ []
+ end
+
private
def licensed_to_use_in_com?(maturity)
diff --git a/ee/lib/api/code_suggestions.rb b/ee/lib/api/code_suggestions.rb
index 2745c70c7748b336676596b65209b39af44695d0..c3c6f034254ca48655691bc085c98cf655e82d17 100644
--- a/ee/lib/api/code_suggestions.rb
+++ b/ee/lib/api/code_suggestions.rb
@@ -30,8 +30,10 @@ def model_gateway_headers(headers, service)
).merge(saas_headers).transform_values { |v| Array(v) }
end
- def connector_public_headers
- Gitlab::CloudConnector.ai_headers(current_user)
+ def connector_public_headers(service_name)
+ namespace_ids = current_user.allowed_by_namespace_ids(service_name)
+
+ Gitlab::CloudConnector.ai_headers(current_user, namespace_ids: namespace_ids)
.merge(saas_headers)
.merge('X-Gitlab-Authentication-Type' => 'oidc')
end
@@ -181,7 +183,7 @@ def file_too_large_with_origin_header!
# https://gitlab.com/gitlab-org/modelops/applied-ml/code-suggestions/ai-assist/-/issues/429
token: token[:token],
expires_at: token[:expires_at],
- headers: connector_public_headers
+ headers: connector_public_headers(completion_model_details.feature_name)
}.tap do |a|
a[:model_details] = details_hash unless details_hash.blank?
end
diff --git a/ee/lib/code_suggestions/completions_model_details.rb b/ee/lib/code_suggestions/completions_model_details.rb
index 28ebd3098180e7f51f3e2783c5641d6ff2a53253..103e0e2bee412b038005916fa1e984ac95ce1dfe 100644
--- a/ee/lib/code_suggestions/completions_model_details.rb
+++ b/ee/lib/code_suggestions/completions_model_details.rb
@@ -1,11 +1,11 @@
# frozen_string_literal: true
module CodeSuggestions
- class CompletionsModelDetails
+ class CompletionsModelDetails < ModelDetails
FEATURE_SETTING_NAME = 'code_completions'
def initialize(current_user:)
- @current_user = current_user
+ super(current_user: current_user, feature_setting_name: FEATURE_SETTING_NAME)
end
def current_model
@@ -21,14 +21,8 @@ def current_model
{}
end
- def feature_disabled?
- !!feature_setting&.disabled?
- end
-
private
- attr_reader :current_user
-
def codestral_model_details
{
model_provider: CodeSuggestions::Prompts::CodeCompletion::VertexCodestral::MODEL_PROVIDER,
@@ -43,10 +37,6 @@ def fireworks_qwen_2_5_model_details
}
end
- def self_hosted?
- feature_setting&.self_hosted?
- end
-
def use_codestral_for_code_completions?
Feature.enabled?(:use_codestral_for_code_completions, current_user, type: :beta)
end
@@ -54,9 +44,5 @@ def use_codestral_for_code_completions?
def use_fireworks_qwen_for_code_completions?
Feature.enabled?(:fireworks_qwen_code_completion, current_user, type: :beta)
end
-
- def feature_setting
- @feature_setting ||= ::Ai::FeatureSetting.find_by_feature(FEATURE_SETTING_NAME)
- end
end
end
diff --git a/ee/lib/code_suggestions/model_details.rb b/ee/lib/code_suggestions/model_details.rb
new file mode 100644
index 0000000000000000000000000000000000000000..748f18661df6804df0ae36845bf74e5c2d91569c
--- /dev/null
+++ b/ee/lib/code_suggestions/model_details.rb
@@ -0,0 +1,41 @@
+# frozen_string_literal: true
+
+module CodeSuggestions
+ class ModelDetails
+ def initialize(current_user:, feature_setting_name:)
+ @current_user = current_user
+ @feature_setting_name = feature_setting_name
+ end
+
+ def feature_setting
+ @feature_setting ||= ::Ai::FeatureSetting.find_by_feature(feature_setting_name)
+ end
+
+ def base_url
+ feature_setting&.base_url || Gitlab::AiGateway.url
+ end
+
+ def feature_name
+ if self_hosted?
+ :self_hosted_models
+ else
+ :code_suggestions
+ end
+ end
+
+ def feature_disabled?
+ # In case the code suggestions feature is being used via self-hosted models,
+ # it can also be disabled completely. In such cases, this check
+ # can be used to prevent exposing the feature via UI/API.
+ !!feature_setting&.disabled?
+ end
+
+ def self_hosted?
+ feature_setting&.self_hosted?
+ end
+
+ private
+
+ attr_reader :current_user, :feature_setting_name
+ end
+end
diff --git a/ee/lib/code_suggestions/tasks/base.rb b/ee/lib/code_suggestions/tasks/base.rb
index cc16bd2245cbadced099f4de32a4c876e24b71e2..d7b7df777249d9df4e184408ea14649d854a7b6d 100644
--- a/ee/lib/code_suggestions/tasks/base.rb
+++ b/ee/lib/code_suggestions/tasks/base.rb
@@ -5,36 +5,15 @@ module Tasks
class Base
AI_GATEWAY_CONTENT_SIZE = 100_000
+ delegate :base_url, :self_hosted?, :feature_setting, :feature_name, :feature_disabled?, to: :model_details
+
def initialize(params: {}, unsafe_passthrough_params: {}, current_user: nil)
- @feature_setting = ::Ai::FeatureSetting.find_by_feature(feature_setting_name)
+ @model_details = ModelDetails.new(current_user: current_user, feature_setting_name: feature_setting_name)
@params = params
@unsafe_passthrough_params = unsafe_passthrough_params
@current_user = current_user
end
- def base_url
- feature_setting&.base_url || Gitlab::AiGateway.url
- end
-
- def self_hosted?
- feature_setting&.self_hosted?
- end
-
- def feature_disabled?
- # In case the code suggestions feature is being used via self-hosted models,
- # it can also be disabled completely. In such cases, this check
- # can be used to prevent exposing the feature via UI/API.
- !!feature_setting&.disabled?
- end
-
- def feature_name
- if self_hosted?
- :self_hosted_models
- else
- :code_suggestions
- end
- end
-
def endpoint
# TODO: After their migration to AIGW, both generations and
# completions will use the same `/completions` endpoint in v3.
@@ -56,7 +35,7 @@ def body
private
- attr_reader :params, :unsafe_passthrough_params, :feature_setting, :current_user
+ attr_reader :params, :unsafe_passthrough_params, :model_details, :current_user
def endpoint_name
raise NotImplementedError
diff --git a/ee/lib/gitlab/ai_gateway.rb b/ee/lib/gitlab/ai_gateway.rb
index f299270502062662d83225f30d0994c76502deb5..d994d29cc105e9426f6691a2c0ff3db401dc0e47 100644
--- a/ee/lib/gitlab/ai_gateway.rb
+++ b/ee/lib/gitlab/ai_gateway.rb
@@ -50,9 +50,7 @@ def self.enabled_feature_flags
end
def self.headers(user:, service:, agent: nil, lsp_version: nil)
- allowed_by_namespace_ids = []
-
- user&.allowed_to_use?(service.name) { |namespace_ids| allowed_by_namespace_ids = namespace_ids }
+ allowed_by_namespace_ids = user&.allowed_by_namespace_ids(service.name) || []
{
'X-Gitlab-Authentication-Type' => 'oidc',
diff --git a/ee/lib/gitlab/cloud_connector.rb b/ee/lib/gitlab/cloud_connector.rb
index 5ae48bd51733e0015c74a1f4102d10d83ea60eee..79bf389302cecd21123d86ca00fd7bb61a91de1a 100644
--- a/ee/lib/gitlab/cloud_connector.rb
+++ b/ee/lib/gitlab/cloud_connector.rb
@@ -33,7 +33,8 @@ def ai_headers(user, namespace_ids: [])
namespace_ids: namespace_ids
)
headers(user).merge(
- 'X-Gitlab-Duo-Seat-Count' => effective_seat_count.to_s
+ 'X-Gitlab-Duo-Seat-Count' => effective_seat_count.to_s,
+ 'X-Gitlab-Feature-Enabled-By-Namespace-Ids' => namespace_ids.join(',')
)
end
diff --git a/ee/spec/lib/code_suggestions/completions_model_details_spec.rb b/ee/spec/lib/code_suggestions/completions_model_details_spec.rb
index 1d2f9f267581eb9c5b8899c9450c50bd9dab3de5..1820a99136b6f4e3e95b3f43bbf6cb611e48c613 100644
--- a/ee/spec/lib/code_suggestions/completions_model_details_spec.rb
+++ b/ee/spec/lib/code_suggestions/completions_model_details_spec.rb
@@ -64,22 +64,4 @@
end
end
end
-
- describe '#feature_disabled?' do
- subject(:feature_disabled?) { completions_model_details.feature_disabled? }
-
- it 'returns false' do
- expect(feature_disabled?).to eq(false)
- end
-
- context 'when code_completions is self-hosted, but set to disabled' do
- let_it_be(:feature_setting) do
- create(:ai_feature_setting, provider: :disabled, feature: :code_completions)
- end
-
- it 'returns true' do
- expect(feature_disabled?).to eq(true)
- end
- end
- end
end
diff --git a/ee/spec/lib/code_suggestions/model_details_spec.rb b/ee/spec/lib/code_suggestions/model_details_spec.rb
new file mode 100644
index 0000000000000000000000000000000000000000..4ab754b45f06867279a4207d4d2f0a27523d026a
--- /dev/null
+++ b/ee/spec/lib/code_suggestions/model_details_spec.rb
@@ -0,0 +1,56 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe CodeSuggestions::ModelDetails, feature_category: :code_suggestions do
+ let_it_be(:feature_setting_name) { 'code_completions' }
+ let(:user) { create(:user) }
+ let(:completions_model_details) do
+ described_class.new(current_user: user, feature_setting_name: feature_setting_name)
+ end
+
+ describe '#feature_disabled?' do
+ subject(:feature_disabled?) { completions_model_details.feature_disabled? }
+
+ it 'returns false' do
+ expect(feature_disabled?).to be(false)
+ end
+
+ context 'when the feature is self-hosted, but set to disabled' do
+ let_it_be(:feature_setting) do
+ create(:ai_feature_setting, provider: :disabled, feature: feature_setting_name)
+ end
+
+ it 'returns true' do
+ expect(feature_disabled?).to be(true)
+ end
+ end
+ end
+
+ describe '#base_url' do
+ include_context 'when loading 1_settings initializer'
+
+ # Reload settings to ensure a consistent state
+ # for Settings.cloud_connector base_url
+ # and isolate tests to reduce the risk of flaky tests
+ # due to shared state with other specs
+ before do
+ load_settings
+ end
+
+ it 'returns correct URL' do
+ expect(completions_model_details.base_url).to eql('https://cloud.gitlab.com/ai')
+ end
+
+ context 'when the feature is customized' do
+ let_it_be(:feature_setting) { create(:ai_feature_setting, provider: :vendored) }
+
+ it 'takes the base url from feature settings' do
+ url = "http://localhost:5000"
+ expect(::Gitlab::AiGateway).to receive(:cloud_connector_url).and_return(url)
+
+ expect(completions_model_details.base_url).to eq(url)
+ end
+ end
+ end
+end
diff --git a/ee/spec/lib/code_suggestions/tasks/base_spec.rb b/ee/spec/lib/code_suggestions/tasks/base_spec.rb
index b1944750b29bdf0d4b69404ede3b8db92b544bcc..ed4cd0df5cf8f7795db6baeae3471fb619c9750f 100644
--- a/ee/spec/lib/code_suggestions/tasks/base_spec.rb
+++ b/ee/spec/lib/code_suggestions/tasks/base_spec.rb
@@ -11,33 +11,6 @@ def feature_setting_name
end
end
- describe '#base_url' do
- include_context 'when loading 1_settings initializer'
-
- # Reload settings to ensure a consistent state
- # for Settings.cloud_connector base_url
- # and isolate tests to reduce the risk of flaky tests
- # due to shared state with other specs
- before do
- load_settings
- end
-
- it 'returns correct URL' do
- expect(klass.new.base_url).to eql('https://cloud.gitlab.com/ai')
- end
-
- context 'when the feature is customized' do
- let_it_be(:feature_setting) { create(:ai_feature_setting, provider: :vendored) }
-
- it 'takes the base url from feature settings' do
- url = "http://localhost:5000"
- expect(::Gitlab::AiGateway).to receive(:cloud_connector_url).and_return(url)
-
- expect(klass.new.base_url).to eq(url)
- end
- end
- end
-
describe '#endpoint' do
it 'raies NotImplementedError' do
expect { klass.new.endpoint }.to raise_error(NotImplementedError)
diff --git a/ee/spec/lib/gitlab/cloud_connector_spec.rb b/ee/spec/lib/gitlab/cloud_connector_spec.rb
index 578ac598a0a9cddaedd92f0c2be4c0a9606ca17d..2017505b48f0da7d852768c36e6c22a527127205 100644
--- a/ee/spec/lib/gitlab/cloud_connector_spec.rb
+++ b/ee/spec/lib/gitlab/cloud_connector_spec.rb
@@ -50,7 +50,10 @@
describe '.ai_headers' do
let(:expected_headers) do
- super().merge('X-Gitlab-Duo-Seat-Count' => '0')
+ super().merge(
+ 'X-Gitlab-Duo-Seat-Count' => '0',
+ 'X-Gitlab-Feature-Enabled-By-Namespace-Ids' => namespace_ids.join(',')
+ )
end
let(:namespace_ids) { [1, 42] }
diff --git a/ee/spec/models/concerns/ai/user_authorizable_spec.rb b/ee/spec/models/concerns/ai/user_authorizable_spec.rb
index dd5e7f1dcdfa135809f4fa732ee9d7738b02e4ad..1c590fcc844bc8dfa28e0e706c3f6f1bbd9871f8 100644
--- a/ee/spec/models/concerns/ai/user_authorizable_spec.rb
+++ b/ee/spec/models/concerns/ai/user_authorizable_spec.rb
@@ -240,6 +240,30 @@
end
end
+ describe '#allowed_by_namespace_ids' do
+ let(:ai_feature) { :my_feature }
+
+ subject { user.allowed_by_namespace_ids(ai_feature) }
+
+ context "when allowed_to_use? doesn't yield any value" do
+ before do
+ allow(user).to receive(:allowed_to_use?).with(ai_feature)
+ end
+
+ it { is_expected.to eq([]) }
+ end
+
+ context 'when allowed_to_use? yields namespace ids' do
+ let(:namespace_ids) { [1, 2] }
+
+ before do
+ allow(user).to receive(:allowed_to_use?).with(ai_feature).and_yield(namespace_ids)
+ end
+
+ it { is_expected.to eq(namespace_ids) }
+ end
+ end
+
describe '#any_group_with_ai_available?', :saas, :use_clean_rails_redis_caching do
using RSpec::Parameterized::TableSyntax
diff --git a/ee/spec/requests/api/code_suggestions_spec.rb b/ee/spec/requests/api/code_suggestions_spec.rb
index 9d6dc8c210ac6f5316d6172aadc8e75ecf6a2f3f..6748ddce1a3c45d0ceb0ec1664be6aa4d21c66f5 100644
--- a/ee/spec/requests/api/code_suggestions_spec.rb
+++ b/ee/spec/requests/api/code_suggestions_spec.rb
@@ -16,7 +16,7 @@
}
end
- let(:enabled_by_namespace_ids) { [1, 2] }
+ let(:enabled_by_namespace_ids) { [] }
let(:current_user) { nil }
let(:headers) { {} }
let(:access_code_suggestions) { true }
@@ -24,6 +24,9 @@
let(:global_instance_id) { 'instance-ABC' }
let(:global_user_id) { 'user-ABC' }
let(:gitlab_realm) { 'saas' }
+ let(:service_name) { :code_suggestions }
+ let(:service) { instance_double('::CloudConnector::SelfSigned::AvailableServiceData') }
+ let_it_be(:token) { 'generated-jwt' }
before do
allow(Gitlab).to receive(:com?).and_return(is_saas)
@@ -38,6 +41,12 @@
allow(Gitlab::GlobalAnonymousId).to receive(:user_id).and_return(global_user_id)
allow(Gitlab::GlobalAnonymousId).to receive(:instance_id).and_return(global_instance_id)
+
+ allow(::CloudConnector::AvailableServices).to receive(:find_by_name).with(service_name).and_return(service)
+ allow(service).to receive_messages(access_token: token, name: service_name)
+ allow(service).to receive_message_chain(:add_on_purchases, :assigned_to_user, :any?).and_return(true)
+ allow(service).to receive_message_chain(:add_on_purchases, :assigned_to_user, :uniq_namespace_ids)
+ .and_return(enabled_by_namespace_ids)
end
shared_examples 'a response' do |case_name|
@@ -164,7 +173,6 @@ def is_even(n: int) ->
end
let(:file_name) { 'test.py' }
- let(:service_name) { :code_suggestions }
let(:additional_params) { {} }
let(:body) do
{
@@ -224,19 +232,12 @@ def is_even(n: int) ->
}
end
- let(:service) { instance_double('::CloudConnector::SelfSigned::AvailableServiceData') }
-
subject(:post_api) do
post api('/code_suggestions/completions', current_user), headers: headers, params: body.to_json
end
before do
allow(Gitlab::ApplicationRateLimiter).to receive(:threshold).and_return(0)
- allow(::CloudConnector::AvailableServices).to receive(:find_by_name).with(service_name).and_return(service)
- allow(service).to receive_messages(access_token: token, name: service_name)
- allow(service).to receive_message_chain(:add_on_purchases, :assigned_to_user, :any?).and_return(true)
- allow(service).to receive_message_chain(:add_on_purchases, :assigned_to_user, :uniq_namespace_ids)
- .and_return(enabled_by_namespace_ids)
stub_feature_flags(use_codestral_for_code_completions: false)
stub_feature_flags(fireworks_qwen_code_completion: false)
end
@@ -295,7 +296,7 @@ def request
'X-Gitlab-Host-Name' => [Gitlab.config.gitlab.host],
'X-Gitlab-Realm' => [gitlab_realm],
'Authorization' => ["Bearer #{token}"],
- 'X-Gitlab-Feature-Enabled-By-Namespace-Ids' => [enabled_by_namespace_ids.join(',')],
+ 'X-Gitlab-Feature-Enabled-By-Namespace-Ids' => [""],
'Content-Type' => ['application/json'],
'User-Agent' => ['Super Awesome Browser 43.144.12']
)
@@ -348,7 +349,7 @@ def request
'X-Gitlab-Host-Name' => [Gitlab.config.gitlab.host],
'X-Gitlab-Realm' => [gitlab_realm],
'Authorization' => ["Bearer #{token}"],
- 'X-Gitlab-Feature-Enabled-By-Namespace-Ids' => [enabled_by_namespace_ids.join(',')],
+ 'X-Gitlab-Feature-Enabled-By-Namespace-Ids' => [""],
'Content-Type' => ['application/json'],
'User-Agent' => ['Super Awesome Browser 43.144.12']
)
@@ -373,7 +374,7 @@ def request
expect(params['Header']).to include({
'X-Gitlab-Authentication-Type' => ['oidc'],
'Authorization' => ["Bearer #{token}"],
- 'X-Gitlab-Feature-Enabled-By-Namespace-Ids' => [enabled_by_namespace_ids.join(',')],
+ 'X-Gitlab-Feature-Enabled-By-Namespace-Ids' => [""],
'Content-Type' => ['application/json'],
'X-Gitlab-Instance-Id' => [global_instance_id],
'X-Gitlab-Global-User-Id' => [global_user_id],
@@ -970,7 +971,8 @@ def get_user(session):
'X-Gitlab-Realm' => gitlab_realm,
'X-Gitlab-Version' => Gitlab.version_info.to_s,
'X-Gitlab-Authentication-Type' => 'oidc',
- 'X-Gitlab-Duo-Seat-Count' => duo_seat_count
+ 'X-Gitlab-Duo-Seat-Count' => duo_seat_count,
+ 'X-Gitlab-Feature-Enabled-By-Namespace-Ids' => enabled_by_namespace_ids.join(',')
}
end
@@ -996,6 +998,7 @@ def request
context 'when user belongs to a namespace with an active code suggestions purchase' do
let_it_be(:add_on_purchase) { create(:gitlab_subscription_add_on_purchase) }
+ let_it_be(:enabled_by_namespace_ids) { [add_on_purchase.namespace_id] }
let(:duo_seat_count) { '1' }
let(:headers) do