diff --git a/ee/app/models/ee/user.rb b/ee/app/models/ee/user.rb
index fd4deadf105dd52dbc1d92590dd8fe7970efc9a9..4dbf68f214aa544f3a32f85ceae374968c2b79e0 100644
--- a/ee/app/models/ee/user.rb
+++ b/ee/app/models/ee/user.rb
@@ -22,6 +22,7 @@ module User
GROUP_WITH_AI_CHAT_ENABLED_CACHE_PERIOD = 1.hour
GROUP_WITH_AI_CHAT_ENABLED_CACHE_KEY = 'group_with_ai_chat_enabled'
+ GROUP_IDS_WITH_AI_CHAT_ENABLED_CACHE_KEY = 'group_ids_with_ai_chat_enabled'
DUO_PRO_ADD_ON_CACHE_KEY = 'user-%{user_id}-code-suggestions-add-on-cache'
@@ -256,8 +257,9 @@ def use_separate_indices?
def clear_group_with_ai_available_cache(ids)
cache_keys_ai_features = Array.wrap(ids).map { |id| ["users", id, GROUP_WITH_AI_ENABLED_CACHE_KEY] }
cache_keys_ai_chat = Array.wrap(ids).map { |id| ["users", id, GROUP_WITH_AI_CHAT_ENABLED_CACHE_KEY] }
+ cache_keys_ai_chat_group_ids = Array.wrap(ids).map { |id| ["users", id, GROUP_IDS_WITH_AI_CHAT_ENABLED_CACHE_KEY] }
- cache_keys = cache_keys_ai_features + cache_keys_ai_chat
+ cache_keys = cache_keys_ai_features + cache_keys_ai_chat + cache_keys_ai_chat_group_ids
::Gitlab::Instrumentation::RedisClusterValidator.allow_cross_slot_commands do
Rails.cache.delete_multi(cache_keys)
end
@@ -654,6 +656,16 @@ def any_group_with_ai_available?
end
end
+ def ai_chat_enabled_namespace_ids
+ return [] unless ::Feature.enabled?(:ai_chat_enabled_namespace_ids)
+
+ Rails.cache.fetch(['users', id, GROUP_IDS_WITH_AI_CHAT_ENABLED_CACHE_KEY], expires_in: GROUP_WITH_AI_CHAT_ENABLED_CACHE_PERIOD) do
+ groups = member_namespaces.with_ai_supported_plan(:ai_chat)
+ groups = ::Feature.enabled?(:duo_chat_ga) ? groups : groups.namespace_settings_with_ai_features_enabled
+ groups.pluck(Arel.sql('DISTINCT traversal_ids[1]'))
+ end
+ end
+
def any_group_with_ai_chat_available?
Rails.cache.fetch(['users', id, GROUP_WITH_AI_CHAT_ENABLED_CACHE_KEY], expires_in: GROUP_WITH_AI_CHAT_ENABLED_CACHE_PERIOD) do
groups = member_namespaces.with_ai_supported_plan(:ai_chat)
diff --git a/ee/app/services/llm/chat_service.rb b/ee/app/services/llm/chat_service.rb
index 36fdff9f34bc55e930cf018b5b4a809e27b24515..bab386404e5e623129b35a65646959fc2f5bfe30 100644
--- a/ee/app/services/llm/chat_service.rb
+++ b/ee/app/services/llm/chat_service.rb
@@ -21,7 +21,13 @@ def perform
@options = options.merge(agent_version_id: agent_version.id)
end
- track_internal_event('request_duo_chat_response', user: user, project: project, namespace: namespace)
+ track_internal_event(
+ 'request_duo_chat_response',
+ user: user,
+ project: project,
+ namespace: namespace,
+ feature_enabled_by_namespace_ids: user.ai_chat_enabled_namespace_ids
+ )
prompt_message.save!
GraphqlTriggers.ai_completion_response(prompt_message)
diff --git a/ee/config/feature_flags/gitlab_com_derisk/ai_chat_enabled_namespace_ids.yml b/ee/config/feature_flags/gitlab_com_derisk/ai_chat_enabled_namespace_ids.yml
new file mode 100644
index 0000000000000000000000000000000000000000..73738fd4b428428b2f9d1d59c6560cfbb47b1f19
--- /dev/null
+++ b/ee/config/feature_flags/gitlab_com_derisk/ai_chat_enabled_namespace_ids.yml
@@ -0,0 +1,9 @@
+---
+name: ai_chat_enabled_namespace_ids
+feature_issue_url:
+introduced_by_url: https://gitlab.com/gitlab-org/gitlab/-/merge_requests/149221
+rollout_issue_url: https://gitlab.com/gitlab-org/gitlab/-/issues/456920
+milestone: '17.0'
+group: group::ai framework
+type: gitlab_com_derisk
+default_enabled: false
diff --git a/ee/spec/models/ee/user_spec.rb b/ee/spec/models/ee/user_spec.rb
index 2eb1a0e5ea90de9cf85aff47bb1a09a168c6d012..00a75ef8eb314b7c6efab1668fea13d75bf3ef89 100644
--- a/ee/spec/models/ee/user_spec.rb
+++ b/ee/spec/models/ee/user_spec.rb
@@ -3907,6 +3907,109 @@
end
end
+ describe '#ai_chat_enabled_namespace_ids', :saas, :use_clean_rails_redis_caching do
+ using RSpec::Parameterized::TableSyntax
+
+ let_it_be(:user) { create(:user) }
+ let_it_be_with_reload(:ultimate_group) { create(:group_with_plan, plan: :ultimate_plan, name: 'ultimate_group') }
+ let_it_be_with_reload(:ultimate_group_id) { ultimate_group.id }
+ let_it_be_with_reload(:bronze_group) { create(:group_with_plan, plan: :bronze_plan, name: 'bronze_group') }
+ let_it_be_with_reload(:free_group) { create(:group_with_plan, plan: :free_plan, name: 'free_group') }
+ let_it_be_with_reload(:group_without_plan) { create(:group, name: 'group_without_plan') }
+ let_it_be_with_reload(:trial_group) { create(:group_with_plan, plan: :ultimate_plan, trial_ends_on: 1.day.from_now, name: 'trial_group') }
+ let_it_be_with_reload(:trial_group_id) { trial_group.id }
+ let_it_be_with_reload(:ultimate_sub_group) { create(:group, parent: ultimate_group, name: 'ultimate_sub_group') }
+ let_it_be_with_reload(:bronze_sub_group) { create(:group, parent: bronze_group, name: 'bronze_sub_group') }
+
+ subject(:group_with_ai_chat_enabled) { user.ai_chat_enabled_namespace_ids }
+
+ where(:group, :result) do
+ ref(:bronze_group) | []
+ ref(:free_group) | []
+ ref(:group_without_plan) | []
+ ref(:ultimate_group) | [ref(:ultimate_group_id)]
+ ref(:trial_group) | [ref(:trial_group_id)]
+ end
+
+ with_them do
+ context 'when member of the root group' do
+ before do
+ group.add_guest(user)
+ end
+
+ context 'when ai features are enabled' do
+ include_context 'with ai features enabled for group'
+
+ it { is_expected.to eq(result) }
+
+ it 'caches the result' do
+ group_with_ai_chat_enabled
+
+ expect(Rails.cache.fetch(['users', user.id, 'group_ids_with_ai_chat_enabled'])).to eq(result)
+ end
+ end
+
+ context 'when duo_chat_ga feature flag is disaabled and ai features are not enabled' do
+ before do
+ stub_feature_flags(duo_chat_ga: false)
+ end
+
+ it { is_expected.to eq([]) }
+ end
+ end
+ end
+
+ context 'when member of a sub-group only' do
+ include_context 'with ai features enabled for group'
+
+ context 'with eligible group' do
+ let(:group) { ultimate_group }
+
+ before do
+ ultimate_sub_group.add_guest(user)
+ end
+
+ it { is_expected.to eq([group.id]) }
+ end
+
+ context 'with not eligible group' do
+ let(:group) { bronze_group }
+
+ before do
+ bronze_sub_group.add_guest(user)
+ end
+
+ it { is_expected.to eq([]) }
+ end
+ end
+
+ context 'when member of a project only' do
+ include_context 'with ai features enabled for group'
+
+ context 'with eligible group' do
+ let(:group) { ultimate_group }
+ let_it_be(:project) { create(:project, group: ultimate_group) }
+
+ before do
+ project.add_guest(user)
+ end
+
+ it { is_expected.to eq([group.id]) }
+ end
+
+ context 'with not eligible group' do
+ let(:group) { bronze_group }
+ let_it_be(:project) { create(:project, group: bronze_group) }
+
+ before do
+ project.add_guest(user)
+ end
+
+ it { is_expected.to eq([]) }
+ end
+ end
+ end
+
describe '.clear_group_with_ai_available_cache', :use_clean_rails_redis_caching do
let_it_be(:user) { create(:user) }
let_it_be(:other_user) { create(:user) }
@@ -3916,18 +4019,21 @@
user.any_group_with_ai_available?
other_user.any_group_with_ai_available?
yet_another_user.any_group_with_ai_chat_available?
+ yet_another_user.ai_chat_enabled_namespace_ids
end
it 'clears cache from users with the given ids', :aggregate_failures do
expect(Rails.cache.fetch(['users', user.id, 'group_with_ai_enabled'])).to eq(false)
expect(Rails.cache.fetch(['users', other_user.id, 'group_with_ai_enabled'])).to eq(false)
expect(Rails.cache.fetch(['users', yet_another_user.id, 'group_with_ai_chat_enabled'])).to eq(false)
+ expect(Rails.cache.fetch(['users', yet_another_user.id, 'group_ids_with_ai_chat_enabled'])).to eq([])
described_class.clear_group_with_ai_available_cache([user.id, yet_another_user.id])
expect(Rails.cache.fetch(['users', user.id, 'group_with_ai_enabled'])).to be_nil
expect(Rails.cache.fetch(['users', other_user.id, 'group_with_ai_enabled'])).to eq(false)
expect(Rails.cache.fetch(['users', yet_another_user.id, 'group_with_ai_chat_enabled'])).to be_nil
+ expect(Rails.cache.fetch(['users', yet_another_user.id, 'group_ids_with_ai_chat_enabled'])).to be_nil
end
it 'clears cache when given a single id', :aggregate_failures do
diff --git a/ee/spec/services/llm/chat_service_spec.rb b/ee/spec/services/llm/chat_service_spec.rb
index 01fe15d86836e88e2e9e147d388672d2cfc662f7..a6836eede2e68853477ecc2d5a4d984d52d0e548 100644
--- a/ee/spec/services/llm/chat_service_spec.rb
+++ b/ee/spec/services/llm/chat_service_spec.rb
@@ -49,7 +49,9 @@
it_behaves_like 'schedules completion worker'
it_behaves_like 'llm service caches user request'
it_behaves_like 'service emitting message for user prompt'
- it_behaves_like 'track internal event for Duo Chat'
+ it_behaves_like 'track internal event for Duo Chat' do
+ let(:feature_enabled_by_namespace_ids) { [] }
+ end
end
context 'when resource is a user' do
@@ -62,6 +64,7 @@
it_behaves_like 'service emitting message for user prompt'
it_behaves_like 'track internal event for Duo Chat' do
let(:project) { nil }
+ let(:feature_enabled_by_namespace_ids) { [] }
end
end
end
@@ -159,7 +162,9 @@
it_behaves_like 'schedules completion worker'
it_behaves_like 'llm service caches user request'
it_behaves_like 'service emitting message for user prompt'
- it_behaves_like 'track internal event for Duo Chat'
+ it_behaves_like 'track internal event for Duo Chat' do
+ let(:feature_enabled_by_namespace_ids) { [group.id] }
+ end
end
context 'when resource is a user' do
@@ -172,6 +177,7 @@
it_behaves_like 'service emitting message for user prompt'
it_behaves_like 'track internal event for Duo Chat' do
let(:project) { nil }
+ let(:feature_enabled_by_namespace_ids) { [group.id] }
end
end
end