diff --git a/ee/app/workers/llm/completion_worker.rb b/ee/app/workers/llm/completion_worker.rb
index af200dc9e29339063be91e25120a6b31b6b54b37..a48c272a79e5fe70eae9da9c8a38ac74c4a40354 100644
--- a/ee/app/workers/llm/completion_worker.rb
+++ b/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
diff --git a/ee/lib/gitlab/sidekiq_middleware/set_session/server.rb b/ee/lib/gitlab/sidekiq_middleware/set_session/server.rb
index 59e01e35e410c25cbee1042b62d9d5f511b41080..88f8d024031f62e60a19e5ffe16deed426effb82 100644
--- a/ee/lib/gitlab/sidekiq_middleware/set_session/server.rb
+++ b/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
diff --git a/ee/spec/lib/gitlab/sidekiq_middleware/set_session/server_spec.rb b/ee/spec/lib/gitlab/sidekiq_middleware/set_session/server_spec.rb
index ae3f154b7d7bbc4888c767eb3b48a7d08a96e298..b992420d26227589dbda4a69e9641bfa478c74c7 100644
--- a/ee/spec/lib/gitlab/sidekiq_middleware/set_session/server_spec.rb
+++ b/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
diff --git a/ee/spec/workers/llm/completion_worker_spec.rb b/ee/spec/workers/llm/completion_worker_spec.rb
index 086d63d5c6f01c8327761c4761da1c5694d47466..2d09ee65949293c532ebaacf40a36344cfa4f464 100644
--- a/ee/spec/workers/llm/completion_worker_spec.rb
+++ b/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
diff --git a/lib/gitlab/session.rb b/lib/gitlab/session.rb
index 7487ba04a6dcb15f79715ab675fd7e2031f7c99d..cb74ec860ee855b2e5014b6ae4ceaee5a4673185 100644
--- a/lib/gitlab/session.rb
+++ b/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)
diff --git a/spec/lib/gitlab/session_spec.rb b/spec/lib/gitlab/session_spec.rb
index 171288da1d55d2cab3511f01f4ea9d4fd93669d4..3efac84ffbd8c1daed3b5a3679032540b86968c6 100644
--- a/spec/lib/gitlab/session_spec.rb
+++ b/spec/lib/gitlab/session_spec.rb
@@ -1,6 +1,6 @@
# 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