From c63d7927f91fa71ee90343f64a322d491ad020f8 Mon Sep 17 00:00:00 2001
From: t8y2 <1156263951@qq.com>
Date: Wed, 9 Sep 2026 13:12:34 +0800
Subject: [PATCH] fix(ai): preserve references when copying failed messages

Closes #8552
---
 .../src/lib/__tests__/ai/aiMessageCopy.spec.ts     | 13 +++++++++++++
 apps/desktop/src/lib/ai/aiMessageCopy.ts           | 14 ++++++++------
 2 files changed, 21 insertions(+), 6 deletions(-)

diff --git a/apps/desktop/src/lib/__tests__/ai/aiMessageCopy.spec.ts b/apps/desktop/src/lib/__tests__/ai/aiMessageCopy.spec.ts
index 04bacc6105..0e69b9d2ac 100644
--- a/apps/desktop/src/lib/__tests__/ai/aiMessageCopy.spec.ts
+++ b/apps/desktop/src/lib/__tests__/ai/aiMessageCopy.spec.ts
@@ -8,6 +8,19 @@ describe("AI message copy payload", () => {
     expect(resolveAiMessageCopyText({ role: "user", content }, false)).toBe(content);
   });
 
+  it("keeps table and SQL file references when copying a user message", () => {
+    const message = {
+      role: "user" as const,
+      content: "Explain this query",
+      mentions: [
+        { kind: "table" as const, raw: "@public.orders" },
+        { kind: "sqlFile" as const, raw: "@setup.sql" },
+      ],
+    };
+
+    expect(resolveAiMessageCopyText(message, false)).toBe("@public.orders @setup.sql\n\nExplain this query");
+  });
+
   it("copies only the assistant message body", () => {
     const message = {
       role: "assistant" as const,
diff --git a/apps/desktop/src/lib/ai/aiMessageCopy.ts b/apps/desktop/src/lib/ai/aiMessageCopy.ts
index efe8a36d53..46b23875d3 100644
--- a/apps/desktop/src/lib/ai/aiMessageCopy.ts
+++ b/apps/desktop/src/lib/ai/aiMessageCopy.ts
@@ -1,14 +1,16 @@
 export interface AiMessageCopyCandidate {
   role: "user" | "assistant";
   content: string;
+  mentions?: readonly AiMessageCopyMention[];
+}
+
+export interface AiMessageCopyMention {
+  kind: string;
+  raw: string;
 }
 
-/**
- * Keeps message-level copy payloads scoped to the original message body.
- * Reasoning, tool activity, mentions, and rendered HTML live outside this
- * contract and must not be appended by the caller.
- */
 export function resolveAiMessageCopyText(message: AiMessageCopyCandidate, streaming: boolean): string | null {
   if (!message.content || streaming) return null;
-  return message.content;
+  const references = message.mentions?.filter((mention) => mention.kind === "table" || mention.kind === "sqlFile").map((mention) => mention.raw) ?? [];
+  return references.length > 0 ? `${references.join(" ")}\n\n${message.content}` : message.content;
 }