Separate self-hosted and cloud-hosted code suggestions features

ee/app/graphql/resolvers/ai/code_suggestions_access_resolver.rb

@@ -9,7 +9,7 @@ def resolve
         return false unless current_user
 
         Feature.enabled?(:ai_duo_code_suggestions_switch, type: :ops) &&
-          CloudConnector::AvailableServices.find_by_name(:code_suggestions).allowed_for?(current_user)
+          current_user.can?(:access_code_suggestions)
       end
     end
   end

ee/app/models/ai/feature_setting.rb

@@ -28,6 +28,10 @@ class FeatureSetting < ApplicationRecord
       duo_chat: 2
     }
 
+    def self.code_suggestions_self_hosted?
+      exists?(feature: [:code_generations, :code_completions], provider: :self_hosted)
+    end
+
     def provider_title
       title = PROVIDER_TITLES[provider]
       return title unless self_hosted?

ee/app/policies/ee/global_policy.rb

@@ -61,8 +61,12 @@ module GlobalPolicy
 
       condition(:code_suggestions_enabled_for_user) do
         next false unless @user
+        next true if CloudConnector::AvailableServices.find_by_name(:code_suggestions).allowed_for?(@user)
 
-        CloudConnector::AvailableServices.find_by_name(:code_suggestions).allowed_for?(@user)
+        next false unless ::Ai::FeatureSetting.code_suggestions_self_hosted?
+
+        self_hosted_models = CloudConnector::AvailableServices.find_by_name(:self_hosted_models)
+        self_hosted_models.free_access? || self_hosted_models.allowed_for?(@user)
       end
 
       condition(:duo_chat_enabled) do

ee/config/cloud_connector/access_data.yml

@@ -64,3 +64,11 @@ services: # Cloud connector features (i.e. code_suggestions, duo_chat...)
       duo_enterprise:
         unit_primitives:
           - resolve_vulnerability
+  self_hosted_models:
+    backend: 'gitlab-ai-gateway'
+    cut_off_date: 2024-08-31 00:00:00 UTC
+    bundled_with:
+      duo_enterprise:
+        unit_primitives:
+          - code_suggestions
+          - duo_chat

ee/lib/api/code_suggestions.rb

@@ -76,10 +76,6 @@ def saas_headers
           end
         end
         post do
-          token = ::CloudConnector::AvailableServices.find_by_name(:code_suggestions).access_token(current_user)
-
-          unauthorized! if token.nil?
-
           check_rate_limit!(:code_suggestions_api_endpoint, scope: current_user) do
             Gitlab::InternalEvents.track_event(
               'code_suggestions_rate_limit_exceeded',
@@ -95,6 +91,12 @@ def saas_headers
             unsafe_passthrough_params: params.except(:private_token)
           ).task
 
+          service = CloudConnector::AvailableServices.find_by_name(task.feature_name)
+          unauthorized! unless service.free_access? || service.allowed_for?(current_user)
+
+          token = service.access_token(current_user)
+          unauthorized! if token.nil?
+
           body = task.body
           file_too_large! if body.size > MAX_BODY_SIZE
 

ee/lib/code_suggestions/tasks/base.rb

@@ -26,6 +26,10 @@ def body
         body_params.to_json
       end
 
+      def feature_name
+        :code_suggestions
+      end
+
       private
 
       attr_reader :params, :unsafe_passthrough_params

ee/lib/code_suggestions/tasks/self_hosted_code_completion.rb

@@ -17,6 +17,11 @@ def endpoint_name
         'completions'
       end
 
+      override :service_name
+      def feature_name
+        :self_hosted_models
+      end
+
       private
 
       attr_reader :feature_setting

ee/lib/code_suggestions/tasks/self_hosted_code_generation.rb

@@ -17,6 +17,11 @@ def endpoint_name
         'generations'
       end
 
+      override :service_name
+      def feature_name
+        :self_hosted_models
+      end
+
       private
 
       attr_reader :feature_setting

ee/spec/lib/cloud_connector/self_signed/access_data_reader_spec.rb

@@ -12,6 +12,9 @@
     let_it_be(:duo_chat_bundled_with) { { "duo_pro" => duo_chat_unit_primitives } }
     let_it_be(:backend) { 'gitlab-ai-gateway' }
 
+    let_it_be(:self_hosted_models_cut_off_date) { Time.zone.parse("2024-08-31 00:00:00 UTC").utc }
+    let_it_be(:self_hosted_models_bundled_with) { { "duo_enterprise" => [:code_suggestions, :duo_chat] } }
+
     let_it_be(:anthropic_proxy_bundled_with) do
       {
         "duo_enterprise" => %i[
@@ -61,8 +64,8 @@
           duo_chat: [nil, duo_chat_bundled_with, backend],
           anthropic_proxy: [nil, anthropic_proxy_bundled_with, backend],
           vertex_ai_proxy: [nil, vertex_ai_proxy_bundled_with, backend],
-          resolve_vulnerability: [nil, resolve_vulnerability_bundled_with,
-            backend]
+          resolve_vulnerability: [nil, resolve_vulnerability_bundled_with, backend],
+          self_hosted_models: [self_hosted_models_cut_off_date, self_hosted_models_bundled_with, backend]
         }
       end
     end

ee/spec/lib/code_suggestions/tasks/base_spec.rb

@@ -3,7 +3,7 @@
 require 'spec_helper'
 
 RSpec.describe CodeSuggestions::Tasks::Base, feature_category: :code_suggestions do
-  subject { described_class.new }
+  subject(:task) { described_class.new }
 
   describe '.base_url' do
     it 'returns correct URL' do
@@ -13,7 +13,13 @@
 
   describe '#endpoint' do
     it 'raies NotImplementedError' do
-      expect { subject.endpoint }.to raise_error(NotImplementedError)
+      expect { task.endpoint }.to raise_error(NotImplementedError)
+    end
+  end
+
+  describe '#feature_name' do
+    it 'returns code suggestions feature name' do
+      expect(task.feature_name).to eq(:code_suggestions)
     end
   end
 end

ee/spec/lib/code_suggestions/tasks/self_hosted_code_completion_spec.rb

@@ -48,8 +48,8 @@
   end
 
   subject(:task) do
-    described_class.new(feature_setting: feature_setting, params: params,
-      unsafe_passthrough_params: unsafe_params)
+    described_class.new(feature_setting: feature_setting,
+      params: params, unsafe_passthrough_params: unsafe_params)
   end
 
   describe '#body' do
@@ -117,4 +117,10 @@
       end
     end
   end
+
+  describe '#feature_name' do
+    it 'returns code suggestions feature name' do
+      expect(task.feature_name).to eq(:self_hosted_models)
+    end
+  end
 end

ee/spec/lib/code_suggestions/tasks/self_hosted_code_generation_spec.rb

@@ -55,8 +55,8 @@
   end
 
   subject(:task) do
-    described_class.new(feature_setting: code_generations_feature_setting, params: params,
-      unsafe_passthrough_params: unsafe_params)
+    described_class.new(feature_setting: code_generations_feature_setting,
+      params: params, unsafe_passthrough_params: unsafe_params)
   end
 
   describe '#body' do
@@ -144,4 +144,10 @@
       end
     end
   end
+
+  describe '#feature_name' do
+    it 'returns code suggestions feature name' do
+      expect(task.feature_name).to eq(:self_hosted_models)
+    end
+  end
 end

ee/spec/models/ai/feature_setting_spec.rb

@@ -28,4 +28,24 @@
 
     it { expect(feature_setting.provider_title).to eq('Disabled') }
   end
+
+  describe '#code_suggestions_self_hosted?' do
+    where(:feature, :provider, :code_suggestions_self_hosted) do
+      [
+        [:code_generations, :self_hosted, true],
+        [:code_generations, :vendored, false],
+        [:code_completions, :self_hosted, true],
+        [:code_generations, :vendored, false],
+        [:duo_chat, :self_hosted, false]
+      ]
+    end
+
+    with_them do
+      it 'returns whether code generations or completions are self hosted' do
+        create(:ai_feature_setting, feature: feature, provider: provider)
+
+        expect(described_class.code_suggestions_self_hosted?).to eq(code_suggestions_self_hosted)
+      end
+    end
+  end
 end

ee/spec/policies/global_policy_spec.rb

@@ -600,11 +600,17 @@
 
     let_it_be_with_reload(:current_user) { create(:user) }
 
-    where(:duo_pro_seat_assigned, :code_suggestions_licensed, :code_suggestions_enabled_for_user) do
-      true  | true  | be_allowed(:access_code_suggestions)
-      true  | false | be_disallowed(:access_code_suggestions)
-      false | true  | be_disallowed(:access_code_suggestions)
-      false | false | be_disallowed(:access_code_suggestions)
+    where(:code_suggestions_licensed, :duo_pro_seat_assigned, :self_hosted_enabled, :self_hosted_licensed,
+      :self_hosted_free_access, :code_suggestions_enabled_for_user) do
+      true  | true  | true  | true   | false  | be_allowed(:access_code_suggestions)
+      true  | true  | false | false  | false  | be_allowed(:access_code_suggestions)
+      true  | false | true  | true   | false  | be_allowed(:access_code_suggestions)
+      true  | false | true  | false  | true   | be_allowed(:access_code_suggestions)
+      true  | false | true  | false  | false  | be_disallowed(:access_code_suggestions)
+      true  | false | true  | false  | false  | be_disallowed(:access_code_suggestions)
+      true  | false | false | false  | false  | be_disallowed(:access_code_suggestions)
+      false | true  | true  | true   | false  | be_disallowed(:access_code_suggestions)
+      false | false | false | false  | false  | be_disallowed(:access_code_suggestions)
     end
 
     with_them do
@@ -615,6 +621,13 @@
                                                                           .and_return(code_suggestions_service_data)
         allow(code_suggestions_service_data).to receive(:allowed_for?).with(current_user)
                                                                       .and_return(duo_pro_seat_assigned)
+        allow(::Ai::FeatureSetting).to receive(:code_suggestions_self_hosted?).and_return(:self_hosted_enabled)
+        self_hosted_models_service_data = instance_double(CloudConnector::BaseAvailableServiceData)
+        allow(CloudConnector::AvailableServices).to receive(:find_by_name).with(:self_hosted_models)
+                                                                          .and_return(self_hosted_models_service_data)
+        allow(self_hosted_models_service_data).to receive(:allowed_for?).with(current_user)
+                                                                      .and_return(self_hosted_licensed)
+        allow(self_hosted_models_service_data).to receive(:free_access?).and_return(self_hosted_free_access)
       end
 
       it { is_expected.to code_suggestions_enabled_for_user }

ee/spec/requests/api/code_suggestions_spec.rb

@@ -170,8 +170,10 @@ def is_even(n: int) ->
 
     before do
       allow(Gitlab::ApplicationRateLimiter).to receive(:threshold).and_return(0)
-      allow(CloudConnector::AvailableServices).to receive_message_chain(:find_by_name,
-        :access_token).and_return(token)
+
+      service = instance_double('::CloudConnector::SelfSigned::AvailableServiceData')
+      allow(::CloudConnector::AvailableServices).to receive(:find_by_name).and_return(service)
+      allow(service).to receive_messages({ free_access?: false, allowed_for?: true, access_token: token })
     end
 
     shared_examples 'code completions endpoint' do

ee/spec/requests/api/graphql/user_code_suggestions_available_spec.rb

@@ -34,6 +34,8 @@
     before do
       allow(Ability)
         .to receive(:allowed?).and_call_original
+
+      stub_licensed_features(code_suggestions: true)
     end
 
     context 'when user has access to code suggestions' do