Merge branch '499140-glcom-seat-count-header' into 'master'

Support namespaces in X-Gitlab-Duo-Seat-Count

See merge request https://gitlab.com/gitlab-org/gitlab/-/merge_requests/171135



Merged-by: Matthias Käppler <mkaeppler@gitlab.com>
Approved-by: Roy Zwambag <rzwambag@gitlab.com>
Approved-by: Nikola Milojevic <nmilojevic@gitlab.com>

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)

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
 

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

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) }

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

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