Merge branch '452548-post-fix' into 'master'

Pass session_id to nested sidekiq calls

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



Merged-by: Smriti Garg <sgarg@gitlab.com>
Approved-by: Smriti Garg <sgarg@gitlab.com>
Reviewed-by: Mark Chao <mchao@gitlab.com>
Reviewed-by: Bogdan Denkovych <bdenkovych@gitlab.com>
Co-authored-by: Mark Chao <mchao@gitlab.com>

ee/app/workers/llm/completion_worker.rb

@@ -31,9 +31,10 @@ def deserialize_message(message_hash, options)
       end
 
       def perform_for(message, options = {})
-        session_id = ::Gitlab::Session.current&.id&.private_id
         # We want to set it even if it is nil, so session will be set and policy check won't be skipped
-        with_ip_address_state.set(set_session_id: session_id).perform_async(serialize_message(message), options)
+        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)
       end
     end
 

ee/lib/gitlab/sidekiq_middleware/set_session/server.rb

@@ -4,12 +4,15 @@ module Gitlab
   module SidekiqMiddleware
     module SetSession
       class Server
+        SESSION_ID_HASH_KEY = 'set_session_id'
+
         def call(_worker, job, _queue)
-          if job.key?('set_session_id')
-            session_id = job['set_session_id']
-            session = ActiveSession.sessions_from_ids([job['set_session_id']]).first if session_id
+          if job.key?(SESSION_ID_HASH_KEY)
+            session_id = job[SESSION_ID_HASH_KEY]
+            session = ActiveSession.sessions_from_ids([session_id]).first if session_id
             session ||= {}
             session = session.with_indifferent_access
+            session[SESSION_ID_HASH_KEY] = session_id # Allows nested sidekiq job to be scheduled with session id
 
             ::Gitlab::Session.with_session(session) do
               yield

ee/spec/lib/gitlab/sidekiq_middleware/set_session/server_spec.rb

@@ -21,14 +21,14 @@ def call!(&block)
 
       it 'sets the session and yields' do
         expect(ActiveSession).to receive(:sessions_from_ids).with([session_id]).and_return([session])
-        expect(Gitlab::Session).to receive(:with_session).with(session).and_yield
+        expect(Gitlab::Session).to receive(:with_session).with(session.merge("set_session_id" => session_id)).and_yield
 
         call!
       end
 
       it 'when session is not found' do
         expect(ActiveSession).to receive(:sessions_from_ids).with([session_id]).and_return([])
-        expect(Gitlab::Session).to receive(:with_session).with({}).and_yield
+        expect(Gitlab::Session).to receive(:with_session).with({ "set_session_id" => session_id }).and_yield
 
         call!
       end
@@ -38,7 +38,7 @@ def call!(&block)
 
         it 'sets empty session and yields' do
           expect(ActiveSession).not_to receive(:sessions_from_ids)
-          expect(Gitlab::Session).to receive(:with_session).with({}).and_yield
+          expect(Gitlab::Session).to receive(:with_session).with({ "set_session_id" => session_id }).and_yield
 
           call!
         end

ee/spec/workers/llm/completion_worker_spec.rb

@@ -71,43 +71,21 @@
       )
     end
 
-    context 'when Session is present' do
-      let(:rack_session) { Rack::Session::SessionId.new('6919a6f1bb119dd7396fadc38fd18d0d') }
-      let(:session) { instance_double(ActionDispatch::Request::Session, id: rack_session) }
-
-      it 'sets set_session_id' do
-        ::Gitlab::Session.with_session(session) do
-          described_class.perform_for(prompt_message, options)
-        end
-
-        job = described_class.jobs.first
-
-        expect(job).to include(
-          'ip_address_state' => ip_address,
-          'set_session_id' => rack_session.private_id,
-          'args' => [
-            hash_including("ai_action" => ai_action_name.to_s),
-            options
-          ]
-        )
-      end
-    end
+    it 'sets set_session_id' do
+      allow(::Gitlab::Session).to receive(:session_id_for_worker).and_return('abc')
 
-    context 'when sessionless' do
-      it 'sets set_session_id to nil' do
-        described_class.perform_for(prompt_message, options)
+      described_class.perform_for(prompt_message, options)
 
-        job = described_class.jobs.first
+      job = described_class.jobs.first
 
-        expect(job).to include(
-          'ip_address_state' => ip_address,
-          'set_session_id' => nil,
-          'args' => [
-            hash_including("ai_action" => ai_action_name.to_s),
-            options
-          ]
-        )
-      end
+      expect(job).to include(
+        'ip_address_state' => ip_address,
+        'set_session_id' => 'abc',
+        'args' => [
+          hash_including("ai_action" => ai_action_name.to_s),
+          options
+        ]
+      )
     end
   end
 end

lib/gitlab/session.rb

@@ -17,6 +17,23 @@ def current
         Thread.current[STORE_KEY]
       end
 
+      def session_id_for_worker
+        session = self.current
+
+        return unless session
+
+        if session.is_a?(ActionDispatch::Request::Session)
+          session.id.private_id
+        elsif session.respond_to?(:[]) # Hash-like
+          session['set_session_id']
+        else
+          raise("Unsupported session class: #{session.class}")
+        end
+      rescue StandardError => e
+        Gitlab::ErrorTracking.track_and_raise_for_dev_exception(e)
+        nil
+      end
+
       protected
 
       def current=(value)

spec/lib/gitlab/session_spec.rb

 # frozen_string_literal: true
 
-require 'fast_spec_helper'
+require 'spec_helper'
 
 RSpec.describe Gitlab::Session do
   it 'uses the current thread as a data store' do
@@ -24,4 +24,51 @@
       expect(described_class.current).to eq nil
     end
   end
+
+  describe '.session_id_for_worker' do
+    context 'when session is ActionDispatch::Request::Session' do
+      let(:rack_session) { Rack::Session::SessionId.new('6919a6f1bb119dd7396fadc38fd18d0d') }
+      let(:session) do
+        ActionDispatch::Request::Session.allocate.tap do |session|
+          allow(session).to receive(:id).and_return(rack_session)
+        end
+      end
+
+      it 'returns rack session private id' do
+        described_class.with_session(session) do
+          expect(described_class.session_id_for_worker).to eq(rack_session.private_id)
+        end
+      end
+    end
+
+    context 'when session behaves like Hash' do
+      let(:session) { { set_session_id: 'abc' }.with_indifferent_access }
+
+      it 'returns session id in Hash' do
+        described_class.with_session(session) do
+          expect(described_class.session_id_for_worker).to eq('abc')
+        end
+      end
+    end
+
+    context 'when sessionless' do
+      let(:session) { nil }
+
+      it 'returns nil' do
+        described_class.with_session(session) do
+          expect(described_class.session_id_for_worker).to eq(nil)
+        end
+      end
+    end
+
+    context 'when unknown type' do
+      let(:session) { Object.new }
+
+      it 'raises error' do
+        described_class.with_session(session) do
+          expect { described_class.session_id_for_worker }.to raise_error("Unsupported session class: Object")
+        end
+      end
+    end
+  end
 end