diff --git a/ee/app/graphql/mutations/ai/duo_user_feedback.rb b/ee/app/graphql/mutations/ai/duo_user_feedback.rb
index c3deb47e96a6b0170fd4a672916649457cacf24c..04f4cc9763b9c79e823addacd1a838e5fc72aa0d 100644
--- a/ee/app/graphql/mutations/ai/duo_user_feedback.rb
+++ b/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
diff --git a/ee/app/models/ai/conversation/message.rb b/ee/app/models/ai/conversation/message.rb
index c9d9ceee576f83396792541ba99a610255bd9801..d5ab968457ac83a38c43be0b85bf5e354a76b087 100644
--- a/ee/app/models/ai/conversation/message.rb
+++ b/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
diff --git a/ee/lib/gitlab/llm/ai_gateway/completions/categorize_question.rb b/ee/lib/gitlab/llm/ai_gateway/completions/categorize_question.rb
index 81980b563287a496ee0d69c2d6913ded5063e427..7f58aad76b7ee44dadfad8386720e0da2f4e394e 100644
--- a/ee/lib/gitlab/llm/ai_gateway/completions/categorize_question.rb
+++ b/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
diff --git a/ee/spec/models/ai/conversation/message_spec.rb b/ee/spec/models/ai/conversation/message_spec.rb
index 70b73b78596170916a82905cb0d58424a1a20e00..813984b4f72eeb2862abac338a27736fd0d0e5e7 100644
--- a/ee/spec/models/ai/conversation/message_spec.rb
+++ b/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