Merge branch '470043-simulate-warden-in-sesssionless-sidekiq' into 'master'

Simulate warden session for Duo Chat

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



Merged-by: Jessie Young <jessieyoung@gitlab.com>
Approved-by: Michał Zając <mzajac@gitlab.com>
Approved-by: Jessie Young <jessieyoung@gitlab.com>
Co-authored-by: Mark Chao <mchao@gitlab.com>

ee/app/workers/llm/completion_worker.rb

@@ -31,7 +31,7 @@ def deserialize_message(message_hash, options)
       end
 
       def perform_for(message, options = {})
-        # We want to set it even if it is nil, so session will be set and policy check won't be skipped
+        # set SESSION_ID_HASH_KEY to ensure inside Sidekiq `Gitlab::Session.current` is not nil
         with_ip_address_state.set(
           Gitlab::SidekiqMiddleware::SetSession::Server::SESSION_ID_HASH_KEY => ::Gitlab::Session.session_id_for_worker
         ).perform_async(serialize_message(message), options)
@@ -41,6 +41,14 @@ def perform_for(message, options = {})
     def perform(prompt_message_hash, options = {})
       ai_prompt_message = self.class.deserialize_message(prompt_message_hash, options)
 
+      # set warden to ensure SsoEnforcer#in_context_of_user_web_activity? returns true
+      session = Gitlab::Session.current
+      if Feature.enabled?(:duo_chat_set_warden, ai_prompt_message.user) &&
+          session && !session.key?('warden.user.user.key') && ai_prompt_message.user
+
+        session['warden.user.user.key'] = User.serialize_into_session(ai_prompt_message.user)
+      end
+
       Gitlab::Llm::Tracking.event_for_ai_message(
         self.class.to_s, "perform_completion_worker", ai_message: ai_prompt_message
       )

ee/config/feature_flags/gitlab_com_derisk/duo_chat_set_warden.yml

+---
+name: duo_chat_set_warden
+feature_issue_url: https://gitlab.com/gitlab-org/gitlab/-/issues/470043
+introduced_by_url: https://gitlab.com/gitlab-org/gitlab/-/merge_requests/158022
+rollout_issue_url: https://gitlab.com/gitlab-org/gitlab/-/issues/470646
+milestone: '17.2'
+group: group::duo chat
+type: gitlab_com_derisk
+default_enabled: false

ee/spec/workers/llm/completion_worker_spec.rb

@@ -48,6 +48,52 @@
         client: 'web'
       )
     end
+
+    context 'when warden.user.user.key is nil' do
+      it 'simulates it' do
+        expect_next_instance_of(
+          Llm::Internal::CompletionService,
+          an_object_having_attributes(
+            user: user,
+            resource: resource,
+            request_id: 'uuid',
+            ai_action: ai_action_name
+          ),
+          options
+        ) do |instance|
+          expect(instance).to receive(:execute)
+        end
+
+        Gitlab::Session.with_session(ActionController::TestSession.new) do
+          subject
+
+          expect(Gitlab::Session.current['warden.user.user.key']).to match_array([[user.id], instance_of(String)])
+        end
+      end
+
+      it 'does nothing if duo_chat_set_warden is disabled' do
+        stub_feature_flags(duo_chat_set_warden: false)
+
+        expect_next_instance_of(
+          Llm::Internal::CompletionService,
+          an_object_having_attributes(
+            user: user,
+            resource: resource,
+            request_id: 'uuid',
+            ai_action: ai_action_name
+          ),
+          options
+        ) do |instance|
+          expect(instance).to receive(:execute)
+        end
+
+        Gitlab::Session.with_session(ActionController::TestSession.new) do
+          subject
+
+          expect(Gitlab::Session.current).not_to have_key('warden.user.user.key')
+        end
+      end
+    end
   end
 
   describe '.perform_for' do