From 6dac21f06366cc6e59cd2b28ec9950c21a18a750 Mon Sep 17 00:00:00 2001 From: Tim McCarthy Date: Wed, 28 May 2025 14:09:06 +0200 Subject: [PATCH] Explicitly set organization for AI::Conversation::Thread We don't want it to fallback to the first organization --- app/models/issue.rb | 2 ++ app/models/note.rb | 1 + ee/app/graphql/mutations/ai/action.rb | 4 +++ .../resolvers/ai/chat_messages_resolver.rb | 1 + ee/app/models/ai/conversation/thread.rb | 9 ----- .../duo_code_review_chat_worker.rb | 12 ++++--- ee/lib/gitlab/llm/ai_message.rb | 2 +- ee/lib/gitlab/llm/chat_message.rb | 2 +- ee/lib/gitlab/llm/chat_storage.rb | 7 ++-- ee/lib/gitlab/llm/chat_storage/base.rb | 5 +-- ee/lib/gitlab/llm/chat_storage/postgresql.rb | 8 +++-- ee/spec/lib/gitlab/llm/chat_message_spec.rb | 34 ++++++++++++------- .../llm/chat_storage/postgresql_spec.rb | 8 +++-- ee/spec/lib/gitlab/llm/chat_storage_spec.rb | 6 ++-- ee/spec/models/ai/conversation/thread_spec.rb | 27 +++++---------- .../duo_code_review_chat_worker_spec.rb | 10 ++++-- 16 files changed, 75 insertions(+), 63 deletions(-) diff --git a/app/models/issue.rb b/app/models/issue.rb index af402adc696c59..0371be818f49e6 100644 --- a/app/models/issue.rb +++ b/app/models/issue.rb @@ -140,6 +140,8 @@ def most_recent validate :due_date_after_start_date, if: :validate_due_date? validate :parent_link_confidentiality + delegate :organization, to: :project + alias_attribute :external_author, :service_desk_reply_to pg_full_text_searchable columns: [{ name: 'title', weight: 'A' }, { name: 'description', weight: 'B' }] diff --git a/app/models/note.rb b/app/models/note.rb index 70caa1465b3e25..8e38f985922cdf 100644 --- a/app/models/note.rb +++ b/app/models/note.rb @@ -93,6 +93,7 @@ class Note < ApplicationRecord delegate :gfm_reference, :local_reference, to: :noteable delegate :name, to: :project, prefix: true + delegate :organization, to: :project delegate :title, to: :noteable, allow_nil: true accepts_nested_attributes_for :note_metadata diff --git a/ee/app/graphql/mutations/ai/action.rb b/ee/app/graphql/mutations/ai/action.rb index b75aae79917201..35d908cf58fa2f 100644 --- a/ee/app/graphql/mutations/ai/action.rb +++ b/ee/app/graphql/mutations/ai/action.rb @@ -145,6 +145,10 @@ def handle_chat_arguments(options) raise Gitlab::Graphql::Errors::ArgumentError, e.message end + def base_thread_relation + current_user.ai_conversation_threads.for_organization(Current.organization) + end + def check_feature_flag_enabled!(method) return if Gitlab::Llm::Utils::FlagChecker.flag_enabled_for_feature?(method) diff --git a/ee/app/graphql/resolvers/ai/chat_messages_resolver.rb b/ee/app/graphql/resolvers/ai/chat_messages_resolver.rb index 1437c8228bff62..87f2438fe6606d 100644 --- a/ee/app/graphql/resolvers/ai/chat_messages_resolver.rb +++ b/ee/app/graphql/resolvers/ai/chat_messages_resolver.rb @@ -36,6 +36,7 @@ def resolve(**args) ::Gitlab::Llm::ChatStorage.new( current_user, + Current.organization, agent_version_id, thread, thread_fallback: view_legacy_messages?(args[:conversation_type]) diff --git a/ee/app/models/ai/conversation/thread.rb b/ee/app/models/ai/conversation/thread.rb index 609b1be49987c1..2bce8d6a29b32a 100644 --- a/ee/app/models/ai/conversation/thread.rb +++ b/ee/app/models/ai/conversation/thread.rb @@ -27,8 +27,6 @@ class Thread < ApplicationRecord duo_chat: 4 } - before_create :populate_organization - def self.expired(column, days) raise ArgumentError unless EXPIRATION_COLUMNS.include?(column.to_s) @@ -40,13 +38,6 @@ def to_new_thread! attributes.slice(*%w[user_id organization_id conversation_type]) ) end - - private - - def populate_organization - self.organization_id ||= user.organizations.first&.id || - Organizations::Organization.first.id - end end end end diff --git a/ee/app/workers/merge_requests/duo_code_review_chat_worker.rb b/ee/app/workers/merge_requests/duo_code_review_chat_worker.rb index b86f9e3a548378..5314d2c032b617 100644 --- a/ee/app/workers/merge_requests/duo_code_review_chat_worker.rb +++ b/ee/app/workers/merge_requests/duo_code_review_chat_worker.rb @@ -38,7 +38,10 @@ def perform(note_id) def prepare_prompt_message(note) author = note.author - thread = author.ai_conversation_threads.create!(conversation_type: :duo_code_review) + organization = note.organization + thread = author.ai_conversation_threads.for_organization(organization).create!( + conversation_type: :duo_code_review + ) prompt_message = nil notes = note.discussion.notes @@ -58,7 +61,7 @@ def prepare_prompt_message(note) # We set the MR object as the resource so that it's accessible if Duo Chat decides # to utilize the merge_request_reader tool with identifier type "current". - prompt_message = save_prompt_message(author, role, note.noteable, content, thread) + prompt_message = save_prompt_message(author, role, note.noteable, content, thread, organization) end prompt_message @@ -81,7 +84,7 @@ def parse_response(note, response_body) ::Gitlab::Llm::Utils::CodeSuggestionFormatter.parse(response_body)[:body] end - def save_prompt_message(user, role, resource, content, thread) + def save_prompt_message(user, role, resource, content, thread, organization) prompt_message = ::Gitlab::Llm::ChatMessage .new( ai_action: 'chat', @@ -89,7 +92,8 @@ def save_prompt_message(user, role, resource, content, thread) content: content, role: role, context: ::Gitlab::Llm::AiMessageContext.new(resource: resource), - thread: thread + thread: thread, + organization: organization ) prompt_message.save! diff --git a/ee/lib/gitlab/llm/ai_message.rb b/ee/lib/gitlab/llm/ai_message.rb index a873765a35f6d3..c5a5d6d2e07b6c 100644 --- a/ee/lib/gitlab/llm/ai_message.rb +++ b/ee/lib/gitlab/llm/ai_message.rb @@ -14,7 +14,7 @@ class AiMessage :id, :request_id, :content, :role, :timestamp, :errors, :extras, :user, :ai_action, :client_subscription_id, :type, :chunk_id, :context, :agent_version_id, :referer_url, :platform_origin, :additional_context, - :thread + :thread, :organization ].freeze SLASH_COMMAND_TOOLS = [ diff --git a/ee/lib/gitlab/llm/chat_message.rb b/ee/lib/gitlab/llm/chat_message.rb index 5cefd7eeb12ba2..c5d22802b2319a 100644 --- a/ee/lib/gitlab/llm/chat_message.rb +++ b/ee/lib/gitlab/llm/chat_message.rb @@ -12,7 +12,7 @@ class ChatMessage < AiMessage alias_method :message_xid, :id def save! - storage = ChatStorage.new(user, agent_version_id, thread) + storage = ChatStorage.new(user, organization, agent_version_id, thread) if clear_history? storage.clear! diff --git a/ee/lib/gitlab/llm/chat_storage.rb b/ee/lib/gitlab/llm/chat_storage.rb index 7524b06e2d5431..bbb747d00cb02b 100644 --- a/ee/lib/gitlab/llm/chat_storage.rb +++ b/ee/lib/gitlab/llm/chat_storage.rb @@ -9,8 +9,9 @@ class ChatStorage delegate :messages, :current_thread, :add, :clear!, to: :postgres_storage - def initialize(user, agent_version_id = nil, thread = nil, thread_fallback: true) + def initialize(user, organization, agent_version_id = nil, thread = nil, thread_fallback: true) @user = user + @organization = organization @agent_version_id = agent_version_id @thread = thread @thread_fallback = thread_fallback @@ -42,7 +43,7 @@ def messages_up_to(message_id) private - attr_reader :user, :agent_version_id, :thread + attr_reader :user, :organization, :agent_version_id, :thread def storage_class(type) "Gitlab::Llm::ChatStorage::#{type.camelize}".constantize @@ -57,7 +58,7 @@ def matches_filters?(message, filters) def postgres_storage @postgres_storage ||= storage_class(POSTGRESQL_STORAGE) - .new(user, agent_version_id, thread, thread_fallback: @thread_fallback) + .new(user, organization, agent_version_id, thread, thread_fallback: @thread_fallback) end end end diff --git a/ee/lib/gitlab/llm/chat_storage/base.rb b/ee/lib/gitlab/llm/chat_storage/base.rb index e805edd25dd1ea..8821d80ce19d7d 100644 --- a/ee/lib/gitlab/llm/chat_storage/base.rb +++ b/ee/lib/gitlab/llm/chat_storage/base.rb @@ -15,7 +15,7 @@ class Base # sufficient. MAX_TEXT_LIMIT = 20_000 - def initialize(user, agent_version_id = nil, thread = nil, thread_fallback: true) + def initialize(user, organization, agent_version_id = nil, thread = nil, thread_fallback: true) if thread.nil? log_error( message: 'thread absent', @@ -25,6 +25,7 @@ def initialize(user, agent_version_id = nil, thread = nil, thread_fallback: true end @thread = thread + @organization = organization @agent_version_id = agent_version_id @user = user @thread_fallback = thread_fallback @@ -44,7 +45,7 @@ def clear! private - attr_reader :user, :agent_version_id, :thread + attr_reader :user, :agent_version_id, :thread, :organization end end end diff --git a/ee/lib/gitlab/llm/chat_storage/postgresql.rb b/ee/lib/gitlab/llm/chat_storage/postgresql.rb index 52698598fbaef9..c691190739530d 100644 --- a/ee/lib/gitlab/llm/chat_storage/postgresql.rb +++ b/ee/lib/gitlab/llm/chat_storage/postgresql.rb @@ -52,11 +52,15 @@ def current_thread private def find_default_thread - user.ai_conversation_threads.for_conversation_type(DEFAULT_CONVERSATION_TYPE).last + base_thread_relation.for_conversation_type(DEFAULT_CONVERSATION_TYPE).last end def create_default_thread - user.ai_conversation_threads.create!(conversation_type: DEFAULT_CONVERSATION_TYPE) + base_thread_relation.create!(conversation_type: DEFAULT_CONVERSATION_TYPE) + end + + def base_thread_relation + user.ai_conversation_threads.for_organization(organization) end end end diff --git a/ee/spec/lib/gitlab/llm/chat_message_spec.rb b/ee/spec/lib/gitlab/llm/chat_message_spec.rb index 368fa85ef82112..76c89a601ed05c 100644 --- a/ee/spec/lib/gitlab/llm/chat_message_spec.rb +++ b/ee/spec/lib/gitlab/llm/chat_message_spec.rb @@ -6,7 +6,7 @@ let_it_be(:organization) { create(:organization) } let_it_be(:user) { create(:user, organizations: [organization]) } - subject { build(:ai_chat_message, user: user, agent_version_id: 1) } + subject { build(:ai_chat_message, user: user, organization: organization, agent_version_id: 1) } describe '#conversation_reset?' do it 'returns true for reset message' do @@ -50,7 +50,8 @@ describe '#save!' do it 'saves the message to chat storage' do - expect_next_instance_of(Gitlab::Llm::ChatStorage, subject.user, subject.agent_version_id, nil) do |instance| + expect_next_instance_of(Gitlab::Llm::ChatStorage, subject.user, subject.organization, subject.agent_version_id, + nil) do |instance| expect(instance).to receive(:add).with(subject) end @@ -60,27 +61,29 @@ it 'creates a new thread and assign it to the message' do expect(subject.thread).to be_nil - expect { subject.save! }.to change { user.ai_conversation_threads.count }.by(1) + expect { subject.save! }.to change { user.ai_conversation_threads.for_organization(organization).count }.by(1) expect(subject.thread).to be_an_instance_of(::Ai::Conversation::Thread) end context 'when a thread exists for the user' do - let!(:thread) { create(:ai_conversation_thread, user: user, conversation_type: :duo_chat_legacy) } + let!(:thread) do + create(:ai_conversation_thread, user: user, organization: organization, conversation_type: :duo_chat_legacy) + end it 'fetches the thread and assign it to the message' do expect(subject.thread).to be_nil - expect { subject.save! }.not_to change { user.ai_conversation_threads.count } + expect { subject.save! }.not_to change { user.ai_conversation_threads.for_organization(organization).count } expect(subject.thread).to eq(thread) end it 'unchanges the assigned thread when thread is specified' do message = build(:ai_chat_message, user: user, agent_version_id: 1, thread: thread) - create(:ai_conversation_thread, user: user) + create(:ai_conversation_thread, user: user, organization: organization) - expect { subject.save! }.not_to change { user.ai_conversation_threads.count } + expect { subject.save! }.not_to change { user.ai_conversation_threads.for_organization(organization).count } expect(message.thread).to eq(thread) end @@ -88,9 +91,11 @@ context 'for /reset message' do it 'saves the message to chat storage' do - message = build(:ai_chat_message, user: user, content: '/reset', agent_version_id: 1) + message = build(:ai_chat_message, user: user, organization: organization, content: '/reset', + agent_version_id: 1) - expect_next_instance_of(Gitlab::Llm::ChatStorage, message.user, message.agent_version_id, nil) do |instance| + expect_next_instance_of(Gitlab::Llm::ChatStorage, message.user, message.organization, message.agent_version_id, + nil) do |instance| expect(instance).to receive(:add).with(message) end @@ -100,9 +105,10 @@ context 'for slash commands to clear history' do it "removes all messages from chat storage for message '/new'" do - message = build(:ai_chat_message, user: user, content: '/new', agent_version_id: 1) + message = build(:ai_chat_message, user: user, organization: organization, content: '/new', agent_version_id: 1) - expect_next_instance_of(Gitlab::Llm::ChatStorage, message.user, message.agent_version_id, nil) do |instance| + expect_next_instance_of(Gitlab::Llm::ChatStorage, message.user, message.organization, message.agent_version_id, + nil) do |instance| expect(instance).to receive(:clear!) end @@ -110,9 +116,11 @@ end it "removes all messages from chat storage for message '/clear'" do - message = build(:ai_chat_message, user: user, content: '/clear', agent_version_id: 1) + message = build(:ai_chat_message, user: user, organization: organization, content: '/clear', + agent_version_id: 1) - expect_next_instance_of(Gitlab::Llm::ChatStorage, message.user, message.agent_version_id, nil) do |instance| + expect_next_instance_of(Gitlab::Llm::ChatStorage, message.user, message.organization, message.agent_version_id, + nil) do |instance| expect(instance).to receive(:clear!) end diff --git a/ee/spec/lib/gitlab/llm/chat_storage/postgresql_spec.rb b/ee/spec/lib/gitlab/llm/chat_storage/postgresql_spec.rb index e963cc873a0157..e2e44086e7831f 100644 --- a/ee/spec/lib/gitlab/llm/chat_storage/postgresql_spec.rb +++ b/ee/spec/lib/gitlab/llm/chat_storage/postgresql_spec.rb @@ -42,7 +42,7 @@ let_it_be(:agent_version_id) { 1 } let(:thread) { nil } - subject(:storage) { described_class.new(user, agent_version_id, thread) } + subject(:storage) { described_class.new(user, organization, agent_version_id, thread) } describe '#add' do let(:message) { build(:ai_chat_message, payload) } @@ -145,7 +145,7 @@ end context 'when thread_fallback is false' do - subject(:storage) { described_class.new(user, thread_fallback: false) } + subject(:storage) { described_class.new(user, organization, thread_fallback: false) } it 'does not create a new thread' do expect { current_thread }.not_to change { user.ai_conversation_threads.count } @@ -156,7 +156,9 @@ end context 'when a thread exists for the user' do - let!(:existing_thread) { create(:ai_conversation_thread, user: user, conversation_type: :duo_chat_legacy) } + let!(:existing_thread) do + create(:ai_conversation_thread, user: user, conversation_type: :duo_chat_legacy, organization: organization) + end it 'returns the latest thread' do expect(current_thread).to eq(existing_thread) diff --git a/ee/spec/lib/gitlab/llm/chat_storage_spec.rb b/ee/spec/lib/gitlab/llm/chat_storage_spec.rb index fc748857d2f90b..546b26cfee28bb 100644 --- a/ee/spec/lib/gitlab/llm/chat_storage_spec.rb +++ b/ee/spec/lib/gitlab/llm/chat_storage_spec.rb @@ -27,9 +27,9 @@ let_it_be(:agent_version_id) { 1 } let(:message) { build(:ai_chat_message, payload) } - let(:postgres_storage) { Gitlab::Llm::ChatStorage::Postgresql.new(user, agent_version_id) } + let(:postgres_storage) { Gitlab::Llm::ChatStorage::Postgresql.new(user, organization, agent_version_id) } - subject { described_class.new(user, agent_version_id) } + subject { described_class.new(user, organization, agent_version_id) } describe '#add' do it 'stores the message in PostgreSQL' do @@ -58,7 +58,7 @@ it_behaves_like '#messages' describe '#messages without thread, and thread_fallback false' do - subject { described_class.new(user, thread_fallback: false) } + subject { described_class.new(user, organization, thread_fallback: false) } it 'returns []' do expect do diff --git a/ee/spec/models/ai/conversation/thread_spec.rb b/ee/spec/models/ai/conversation/thread_spec.rb index 4c746cc733638c..c17e64234e1d0f 100644 --- a/ee/spec/models/ai/conversation/thread_spec.rb +++ b/ee/spec/models/ai/conversation/thread_spec.rb @@ -69,28 +69,17 @@ expect(threads).to eq([thread_2, thread_1]) end end - end - - describe 'callbacks' do - describe 'before_create :populate_organization_id' do - let_it_be(:organization) { create(:organization) } - let(:user) { create(:user, organizations: [organization]) } + describe '.for_organization' do + let(:organization_a) { create(:organization) } + let(:organization_b) { create(:organization) } + let(:thread_a) { create(:ai_conversation_thread, organization: organization_a) } + let(:thread_b) { create(:ai_conversation_thread, organization: organization_b) } - it 'assigns organization_id from user organization' do - thread = described_class.create!(user: user, conversation_type: :duo_chat) - - expect(thread.organization_id).to eq(user.organizations.first.id) - end + subject(:threads) { described_class.for_organization(organization_a) } - context 'when user is not assigned to an organization' do - let(:user) { create(:user, organizations: []) } - - it 'assigns organization_id from first found organization' do - thread = described_class.create!(user: user, conversation_type: :duo_chat) - - expect(thread.organization_id).to eq(Organizations::Organization.first.id) - end + it 'returns threads scoped to organization' do + expect(threads).to eq([thread_a]) end end end diff --git a/ee/spec/workers/merge_requests/duo_code_review_chat_worker_spec.rb b/ee/spec/workers/merge_requests/duo_code_review_chat_worker_spec.rb index ac392aa0270e38..435f2ff48b378d 100644 --- a/ee/spec/workers/merge_requests/duo_code_review_chat_worker_spec.rb +++ b/ee/spec/workers/merge_requests/duo_code_review_chat_worker_spec.rb @@ -5,7 +5,8 @@ RSpec.describe MergeRequests::DuoCodeReviewChatWorker, feature_category: :code_review_workflow do subject(:worker) { described_class.new } - let_it_be(:project) { create(:project, :repository) } + let_it_be(:organization) { create(:organization) } + let_it_be(:project) { create(:project, :repository, organization: organization) } let_it_be(:merge_request) { create(:merge_request, source_project: project, target_project: project) } let(:note_type) { :diff_note_on_merge_request } @@ -286,7 +287,9 @@ .and_call_original .ordered - worker.perform(note.id) + expect { worker.perform(note.id) }.to change { + Ai::Conversation::Thread.for_organization(organization).count + }.by(1) expect(Note.last.note).to eq <<~NOTE_CONTENT Comment with suggestions @@ -344,7 +347,8 @@ def chat_message_params(role, resource, content) content: content, role: role, context: an_object_having_attributes(resource: resource), - thread: an_instance_of(::Ai::Conversation::Thread) + thread: an_instance_of(::Ai::Conversation::Thread), + organization: organization } end end -- GitLab