From 7c36c68eb54b606d6d6c9bebd3c944cf9ee96ac2 Mon Sep 17 00:00:00 2001
From: Matthias Kaeppler <mkaeppler@gitlab.com>
Date: Wed, 30 Oct 2024 15:27:01 +0100
Subject: [PATCH] Support namespaces in X-Gitlab-Duo-Seat-Count
We currently only populate this header based on
instance-wide add-on purchases and seat counts,
which does not work for gitlab.com
Changelog: fixed
EE: true
---
.../gitlab_subscriptions/add_on_purchase.rb | 6 +-
ee/lib/gitlab/ai_gateway.rb | 2 +-
ee/lib/gitlab/cloud_connector.rb | 13 ++-
ee/spec/lib/gitlab/ai_gateway_spec.rb | 22 ++--
ee/spec/lib/gitlab/cloud_connector_spec.rb | 10 +-
.../add_on_purchase_spec.rb | 106 ++++++++++++------
6 files changed, 108 insertions(+), 51 deletions(-)
diff --git a/ee/app/models/gitlab_subscriptions/add_on_purchase.rb b/ee/app/models/gitlab_subscriptions/add_on_purchase.rb
index 43d8ca174471e..8cba472df1ecc 100644
--- a/ee/app/models/gitlab_subscriptions/add_on_purchase.rb
+++ b/ee/app/models/gitlab_subscriptions/add_on_purchase.rb
@@ -75,8 +75,10 @@ def self.uniq_namespace_ids
pluck(:namespace_id).compact.uniq
end
- def self.maximum_duo_seat_count
- active.for_duo_pro_or_duo_enterprise.pluck(:quantity).max || 0
+ def self.maximum_duo_seat_count(namespace_ids: [])
+ scope = active.for_duo_pro_or_duo_enterprise
+ scope = scope.by_namespace(namespace_ids) if namespace_ids.any?
+ scope.pluck(:quantity).max || 0
end
def already_assigned?(user)
diff --git a/ee/lib/gitlab/ai_gateway.rb b/ee/lib/gitlab/ai_gateway.rb
index a63ca445c317b..f299270502062 100644
--- a/ee/lib/gitlab/ai_gateway.rb
+++ b/ee/lib/gitlab/ai_gateway.rb
@@ -63,7 +63,7 @@ def self.headers(user:, service:, agent: nil, lsp_version: nil)
# Forward the request time to the model gateway to calculate latency
'X-Gitlab-Rails-Send-Start' => Time.now.to_f.to_s,
'x-gitlab-enabled-feature-flags' => enabled_feature_flags.uniq.join(',')
- }.merge(Gitlab::CloudConnector.ai_headers(user))
+ }.merge(Gitlab::CloudConnector.ai_headers(user, namespace_ids: allowed_by_namespace_ids))
.tap do |result|
result['User-Agent'] = agent if agent # Forward the User-Agent on to the model gateway
diff --git a/ee/lib/gitlab/cloud_connector.rb b/ee/lib/gitlab/cloud_connector.rb
index d4937f6e51a59..c288f1d83e61f 100644
--- a/ee/lib/gitlab/cloud_connector.rb
+++ b/ee/lib/gitlab/cloud_connector.rb
@@ -22,9 +22,18 @@ def headers(user)
end
end
- def ai_headers(user)
+ ###
+ # Returns required HTTP header fields when making AI requests through Cloud Connector.
+ #
+ # user - User making the request, may be null.
+ # namespace_ids - Namespaces for which to return the maximum allowed Duo seat count.
+ # This should only be set when the request is made on gitlab.com.
+ def ai_headers(user, namespace_ids: [])
+ effective_seat_count = GitlabSubscriptions::AddOnPurchase.maximum_duo_seat_count(
+ namespace_ids: namespace_ids
+ )
headers(user).merge(
- 'X-Gitlab-Duo-Seat-Count' => GitlabSubscriptions::AddOnPurchase.maximum_duo_seat_count.to_s
+ 'X-Gitlab-Duo-Seat-Count' => effective_seat_count.to_s
)
end
end
diff --git a/ee/spec/lib/gitlab/ai_gateway_spec.rb b/ee/spec/lib/gitlab/ai_gateway_spec.rb
index 91e2819675a52..e39948ee39515 100644
--- a/ee/spec/lib/gitlab/ai_gateway_spec.rb
+++ b/ee/spec/lib/gitlab/ai_gateway_spec.rb
@@ -86,6 +86,17 @@
let(:service) { instance_double(CloudConnector::BaseAvailableServiceData, name: service_name) }
let(:agent) { nil }
let(:lsp_version) { nil }
+ let(:cloud_connector_headers) do
+ {
+ 'X-Gitlab-Host-Name' => 'hostname',
+ 'X-Gitlab-Instance-Id' => 'ABCDEF',
+ 'X-Gitlab-Global-User-Id' => '123ABC',
+ 'X-Gitlab-Realm' => 'self-managed',
+ 'X-Gitlab-Version' => '17.1.0',
+ 'X-Gitlab-Duo-Seat-Count' => "50"
+ }
+ end
+
let(:expected_headers) do
{
'X-Gitlab-Authentication-Type' => 'oidc',
@@ -94,14 +105,8 @@
'Content-Type' => 'application/json',
'X-Request-ID' => an_instance_of(String),
'X-Gitlab-Rails-Send-Start' => an_instance_of(String),
- 'X-Gitlab-Global-User-Id' => an_instance_of(String),
- 'X-Gitlab-Host-Name' => Gitlab.config.gitlab.host,
- 'X-Gitlab-Instance-Id' => an_instance_of(String),
- 'X-Gitlab-Realm' => Gitlab::CloudConnector::GITLAB_REALM_SELF_MANAGED,
- 'X-Gitlab-Version' => Gitlab.version_info.to_s,
- 'X-Gitlab-Duo-Seat-Count' => "0",
'x-gitlab-enabled-feature-flags' => an_instance_of(String)
- }
+ }.merge(cloud_connector_headers)
end
subject(:headers) { described_class.headers(user: user, service: service, agent: agent, lsp_version: lsp_version) }
@@ -109,6 +114,9 @@
before do
allow(service).to receive(:access_token).with(user).and_return(token)
allow(user).to receive(:allowed_to_use?).with(service_name).and_yield(enabled_by_namespace_ids)
+ allow(Gitlab::CloudConnector).to(
+ receive(:ai_headers).with(user, namespace_ids: enabled_by_namespace_ids).and_return(cloud_connector_headers)
+ )
end
it { is_expected.to match(expected_headers) }
diff --git a/ee/spec/lib/gitlab/cloud_connector_spec.rb b/ee/spec/lib/gitlab/cloud_connector_spec.rb
index a3984800ed425..1aae2c79314ed 100644
--- a/ee/spec/lib/gitlab/cloud_connector_spec.rb
+++ b/ee/spec/lib/gitlab/cloud_connector_spec.rb
@@ -53,15 +53,19 @@
super().merge('X-Gitlab-Duo-Seat-Count' => '0')
end
+ let(:namespace_ids) { [1, 42] }
+
it_behaves_like 'building HTTP headers'
- subject(:headers) { described_class.ai_headers(user) }
+ subject(:headers) { described_class.ai_headers(user, namespace_ids: namespace_ids) }
context 'when Duo seats have been purchased' do
let(:user) { nil }
- it 'set the header to the correct number of seats' do
- expect(GitlabSubscriptions::AddOnPurchase).to receive(:maximum_duo_seat_count).and_return(5)
+ it 'sets the seat count header to the correct number of seats' do
+ expect(GitlabSubscriptions::AddOnPurchase).to(
+ receive(:maximum_duo_seat_count).with(namespace_ids: namespace_ids).and_return(5)
+ )
expect(headers).to include('X-Gitlab-Duo-Seat-Count' => '5')
end
diff --git a/ee/spec/models/gitlab_subscriptions/add_on_purchase_spec.rb b/ee/spec/models/gitlab_subscriptions/add_on_purchase_spec.rb
index 0645a9f77a88a..e525572a074ce 100644
--- a/ee/spec/models/gitlab_subscriptions/add_on_purchase_spec.rb
+++ b/ee/spec/models/gitlab_subscriptions/add_on_purchase_spec.rb
@@ -598,60 +598,94 @@
end
describe '.maximum_duo_seat_count' do
- context 'when there is no duo add-on' do
- let(:expected_maximum_duo_seat_count) { 0 }
+ shared_examples 'seat count calculation' do
+ subject(:seat_count) { described_class.maximum_duo_seat_count(namespace_ids: namespace_ids) }
- it 'returns the default of 0' do
- expect(described_class.maximum_duo_seat_count).to eq(expected_maximum_duo_seat_count)
+ context 'when there is no duo add-on' do
+ let(:expected_maximum_duo_seat_count) { 0 }
+
+ it 'returns the default of 0' do
+ expect(seat_count).to eq(expected_maximum_duo_seat_count)
+ end
end
- end
- context 'when there is a duo pro add-on' do
- let(:expected_maximum_duo_seat_count) { 10 }
+ context 'when there is a duo pro add-on' do
+ let(:expected_maximum_duo_seat_count) { 10 }
- it 'returns the number of seats purchased' do
- create(:gitlab_subscription_add_on_purchase, :active, :gitlab_duo_pro,
- quantity: expected_maximum_duo_seat_count)
+ it 'returns the number of seats purchased' do
+ create(:gitlab_subscription_add_on_purchase, :active, :gitlab_duo_pro,
+ options.merge(quantity: expected_maximum_duo_seat_count))
- expect(described_class.maximum_duo_seat_count).to eq(expected_maximum_duo_seat_count)
+ expect(seat_count).to eq(expected_maximum_duo_seat_count)
+ end
end
- end
- context 'when there is a duo enterprise add-on' do
- let(:expected_maximum_duo_seat_count) { 20 }
+ context 'when there is a duo enterprise add-on' do
+ let(:expected_maximum_duo_seat_count) { 20 }
- it 'returns the number of seats purchased' do
- create(:gitlab_subscription_add_on_purchase, :active, :duo_enterprise,
- quantity: expected_maximum_duo_seat_count)
+ it 'returns the number of seats purchased' do
+ create(:gitlab_subscription_add_on_purchase, :active, :duo_enterprise,
+ options.merge(quantity: expected_maximum_duo_seat_count))
- expect(described_class.maximum_duo_seat_count).to eq(expected_maximum_duo_seat_count)
+ expect(seat_count).to eq(expected_maximum_duo_seat_count)
+ end
+ end
+
+ context 'when there is both a duo pro add-on and a duo enterprise trial add-on' do
+ let(:expected_duo_enterprise_seat_count) { 15 }
+ let(:expected_duo_pro_seat_count) { 5 }
+
+ it 'returns the maximum number of seats purchased for the add-on with the most seats' do
+ create(:gitlab_subscription_add_on_purchase, :active, :duo_enterprise, :trial,
+ options.merge(quantity: expected_duo_enterprise_seat_count))
+ create(:gitlab_subscription_add_on_purchase, :active, :gitlab_duo_pro,
+ options.merge(quantity: expected_duo_pro_seat_count))
+
+ expect(seat_count).to eq(expected_duo_enterprise_seat_count)
+ end
end
- end
- context 'when there is both a duo pro add-on and a duo enterprise trial add-on' do
- let(:expected_duo_enterprise_seat_count) { 15 }
- let(:expected_duo_pro_seat_count) { 5 }
+ context 'when there is both a duo pro trial add-on and a duo enterprise add-on' do
+ let(:expected_duo_enterprise_seat_count) { 10 }
+ let(:expected_duo_pro_seat_count) { 40 }
- it 'returns the maximum number of seats purchased for the add-on with the most seats' do
- create(:gitlab_subscription_add_on_purchase, :active, :duo_enterprise, :trial,
- quantity: expected_duo_enterprise_seat_count)
- create(:gitlab_subscription_add_on_purchase, :active, :gitlab_duo_pro, quantity: expected_duo_pro_seat_count)
+ it 'returns the maximum number of seats purchased for the add-on with the most seats' do
+ create(:gitlab_subscription_add_on_purchase, :active, :duo_enterprise,
+ options.merge(quantity: expected_duo_enterprise_seat_count))
+ create(:gitlab_subscription_add_on_purchase, :active, :gitlab_duo_pro, :trial,
+ options.merge(quantity: expected_duo_pro_seat_count))
- expect(described_class.maximum_duo_seat_count).to eq(expected_duo_enterprise_seat_count)
+ expect(seat_count).to eq(expected_duo_pro_seat_count)
+ end
end
end
- context 'when there is both a duo pro trial add-on and a duo enterprise add-on' do
- let(:expected_duo_enterprise_seat_count) { 10 }
- let(:expected_duo_pro_seat_count) { 40 }
+ context 'when no namespace IDs are given' do
+ let(:namespace_ids) { [] }
+ let(:options) { { namespace: nil } }
+
+ include_examples 'seat count calculation'
+ end
+
+ context 'when namespace IDs are given' do
+ let_it_be(:group) { create(:group) }
+
+ let(:namespace_ids) { [group.id] }
+ let(:options) { { namespace: group } }
+
+ include_examples 'seat count calculation'
- it 'returns the maximum number of seats purchased for the add-on with the most seats' do
- create(:gitlab_subscription_add_on_purchase, :active, :duo_enterprise,
- quantity: expected_duo_enterprise_seat_count)
- create(:gitlab_subscription_add_on_purchase, :active, :gitlab_duo_pro, :trial,
- quantity: expected_duo_pro_seat_count)
+ context 'with more than one group namespace' do
+ let_it_be(:other_group) { create(:group) }
- expect(described_class.maximum_duo_seat_count).to eq(expected_duo_pro_seat_count)
+ it 'returns the highest seat count of both' do
+ create(:gitlab_subscription_add_on_purchase, :active, :duo_enterprise,
+ quantity: 20, namespace: group)
+ create(:gitlab_subscription_add_on_purchase, :active, :gitlab_duo_pro,
+ quantity: 10, namespace: other_group)
+
+ expect(described_class.maximum_duo_seat_count(namespace_ids: [group.id, other_group.id])).to eq(20)
+ end
end
end
end
--
GitLab