Merge branch 'fix-duplicate-chat-message-save' into 'master'

Fix Duo Chat messages being saved twice

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



Merged-by: Gosia Ksionek <mksionek@gitlab.com>
Approved-by: Gosia Ksionek <mksionek@gitlab.com>
Reviewed-by: Mark Chao <mchao@gitlab.com>
Co-authored-by: Alejandro Rodríguez <alejandro@gitlab.com>

ee/lib/gitlab/llm/completions/chat.rb

@@ -70,7 +70,8 @@ def execute
 
           # Send full message to custom clientSubscriptionId at the end of streaming.
           if response_options[:client_subscription_id]
-            ::Gitlab::Llm::ResponseService.new(context, response_options).execute(response: response_modifier)
+            ::Gitlab::Llm::ResponseService.new(context, response_options)
+              .execute(response: response_modifier, save_message: false)
           end
 
           response_handler.execute(response: response_modifier)

ee/lib/gitlab/llm/graphql_subscription_response_service.rb

@@ -3,11 +3,12 @@
 module Gitlab
   module Llm
     class GraphqlSubscriptionResponseService < BaseService
-      def initialize(user, resource, response_modifier, options:)
+      def initialize(user, resource, response_modifier, options:, save_message: true)
         @user = user
         @resource = resource
         @response_modifier = response_modifier
         @options = options
+        @save_message = save_message
         @logger = Gitlab::Llm::Logger.build
       end
 
@@ -39,10 +40,11 @@ def execute
 
       private
 
-      attr_reader :user, :resource, :response_modifier, :options, :logger
+      attr_reader :user, :resource, :response_modifier, :options, :save_message, :logger
 
       def save_message?
-        response_message.is_a?(ChatMessage) &&
+        save_message &&
+          response_message.is_a?(ChatMessage) &&
           !response_message.type &&
           !response_message.chunk_id
       end

ee/lib/gitlab/llm/response_service.rb

@@ -9,10 +9,11 @@ def initialize(context, basic_options)
         @basic_options = basic_options
       end
 
-      def execute(response:, options: {})
+      def execute(response:, options: {}, save_message: true)
         ::Gitlab::Llm::GraphqlSubscriptionResponseService
           .new(user, resource, response,
-            options: basic_options.merge(options))
+            options: basic_options.merge(options),
+            save_message: save_message)
           .execute
       end
 

ee/spec/lib/gitlab/llm/completions/chat_spec.rb

@@ -140,7 +140,7 @@
           an_instance_of(Gitlab::Llm::Chain::GitlabContext), { request_id: 'uuid', ai_action: :chat,
 client_subscription_id: 'someid' }
         ).and_return(stream_response_handler).twice
-        expect(stream_response_handler).to receive(:execute)
+        expect(stream_response_handler).to receive(:execute).with(response: anything, save_message: false)
         expect(categorize_service).to receive(:execute)
         expect(::Llm::ExecuteMethodService).to receive(:new)
           .with(user, user, :categorize_question, categorize_service_params)

ee/spec/lib/gitlab/llm/graphql_subscription_response_service_spec.rb

@@ -75,12 +75,26 @@
   end
 
   describe '#execute' do
-    let(:service) { described_class.new(user, resource, response_modifier, options: options) }
+    let(:save_message) { true }
+    let(:service) do
+      described_class.new(user, resource, response_modifier, options: options, save_message: save_message)
+    end
+
     let_it_be(:resource) { project }
 
     subject { service.execute }
 
     context 'when message is chat' do
+      shared_examples 'not saving the message' do
+        it 'does not save the message' do
+          expect_next_instance_of(::Gitlab::Llm::AiMessage) do |instance|
+            expect(instance).not_to receive(:save!)
+          end
+
+          subject
+        end
+      end
+
       let(:ai_action) { 'chat' }
 
       it 'saves the message' do
@@ -89,28 +103,22 @@
         subject
       end
 
+      context 'when save_message is false' do
+        let(:save_message) { false }
+
+        it_behaves_like 'not saving the message'
+      end
+
       context 'when message is stream chunk' do
         let(:options) { super().merge(chunk_id: 1) }
 
-        it 'does not save the message' do
-          expect_next_instance_of(::Gitlab::Llm::AiMessage) do |instance|
-            expect(instance).not_to receive(:save!)
-          end
-
-          subject
-        end
+        it_behaves_like 'not saving the message'
       end
 
       context 'when message has special type' do
         let(:options) { super().merge(type: 'tool') }
 
-        it 'does not save the message' do
-          expect_next_instance_of(::Gitlab::Llm::AiMessage) do |instance|
-            expect(instance).not_to receive(:save!)
-          end
-
-          subject
-        end
+        it_behaves_like 'not saving the message'
       end
     end
 

ee/spec/lib/gitlab/llm/response_service_spec.rb

@@ -11,15 +11,18 @@
 
   let(:basic_options) { { cache_request: true } }
   let(:options) { { cache_request: false } }
+  let(:save_message) { false }
   let(:graphql_subscription_double) { instance_double(::Gitlab::Llm::GraphqlSubscriptionResponseService) }
 
   describe '#execute' do
     it 'calls GraphQL subscription service with the right params' do
       expect(graphql_subscription_double).to receive(:execute)
       expect(::Gitlab::Llm::GraphqlSubscriptionResponseService).to receive(:new)
-        .with(user, issue, 'response', options: { cache_request: false }).and_return(graphql_subscription_double)
+        .with(user, issue, 'response', options: options, save_message: save_message)
+        .and_return(graphql_subscription_double)
 
-      described_class.new(context, basic_options).execute(response: 'response', options: options)
+      described_class.new(context, basic_options)
+        .execute(response: 'response', options: options, save_message: save_message)
     end
   end
 end