Merge branch '517049-feedback' into 'master'

Make Duo User Feedback thread-aware

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



Merged-by: Shinya Maeda <shinya@gitlab.com>
Approved-by: Dillon Wheeler <dwheeler@gitlab.com>
Approved-by: Shinya Maeda <shinya@gitlab.com>
Co-authored-by: Mark Chao <mchao@gitlab.com>

ee/app/graphql/mutations/ai/duo_user_feedback.rb

@@ -15,7 +15,8 @@ class DuoUserFeedback < BaseMutation
       def resolve(**args)
         raise_resource_not_available_error! unless current_user
 
-        chat_storage = ::Gitlab::Llm::ChatStorage.new(current_user, args[:agent_version_id]&.model_id)
+        thread = ::Ai::Conversation::Message.find_for_user!(args[:ai_message_id], current_user)&.thread
+        chat_storage = ::Gitlab::Llm::ChatStorage.new(current_user, args[:agent_version_id]&.model_id, thread)
         message = chat_storage.messages.find { |m| m.id == args[:ai_message_id] }
 
         raise_resource_not_available_error! unless message
@@ -25,6 +26,8 @@ def resolve(**args)
         track_snowplow_event(args[:tracking_event], message)
 
         { errors: [] }
+      rescue ActiveRecord::RecordNotFound
+        raise_resource_not_available_error!
       end
 
       private

ee/app/models/ai/conversation/message.rb

@@ -21,6 +21,10 @@ class Message < ApplicationRecord
 
       before_create :populate_organization
 
+      def self.find_for_user!(xid, user)
+        for_message_xid(xid).for_user(user).first!
+      end
+
       def self.recent(limit)
         order(id: :desc).limit(limit).reverse
       end

ee/lib/gitlab/llm/ai_gateway/completions/categorize_question.rb

@@ -78,7 +78,7 @@ def valid?
           end
 
           def messages
-            message = ::Ai::Conversation::Message.for_user(user).for_message_xid(options[:message_id]).first! # rubocop:disable CodeReuse/ActiveRecord -- not sure why first is allowed but not first!
+            message = ::Ai::Conversation::Message.find_for_user!(options[:message_id], user)
             ::Gitlab::Llm::ChatStorage.new(user, nil, message.thread).messages_up_to(options[:message_id])
           end
           strong_memoize_attr :messages

ee/spec/models/ai/conversation/message_spec.rb

@@ -67,6 +67,38 @@
         expect(messages).to contain_exactly(message)
       end
     end
+
+    describe '.find_for_user!' do
+      let_it_be(:user) { create(:user) }
+      let_it_be(:thread) { create(:ai_conversation_thread, user: user) }
+      let_it_be(:message) { create(:ai_conversation_message, thread: thread) }
+
+      context 'when message exists and belongs to the user' do
+        it 'returns the message' do
+          expect(described_class.find_for_user!(message.message_xid, user)).to eq(message)
+        end
+      end
+
+      context 'when message exists but belongs to different user' do
+        let(:other_user) { create(:user) }
+
+        it 'raises ActiveRecord::RecordNotFound' do
+          expect do
+            described_class.find_for_user!(message.message_xid, other_user)
+          end.to raise_error(ActiveRecord::RecordNotFound)
+        end
+      end
+
+      context 'when message_xid does not exist' do
+        let(:non_existent_xid) { SecureRandom.uuid }
+
+        it 'raises ActiveRecord::RecordNotFound' do
+          expect do
+            described_class.find_for_user!(non_existent_xid, user)
+          end.to raise_error(ActiveRecord::RecordNotFound)
+        end
+      end
+    end
   end
 
   describe '.recent' do