From 9c45c7eb69ab42e81ebe944f48a0fb84e443a2ce Mon Sep 17 00:00:00 2001
From: chenhang <averych@126.com>
Date: Wed, 1 Jul 2026 10:22:11 +0800
Subject: [PATCH] fix(chat): avoid overwriting user messages on unlock

---
 .../chat-machine.transitions.test.ts          | 65 +++++++++++++++++++
 src/stores/chat/chat-unlock-helpers.ts        | 13 +++-
 2 files changed, 77 insertions(+), 1 deletion(-)

diff --git a/src/stores/chat/__tests__/chat-machine.transitions.test.ts b/src/stores/chat/__tests__/chat-machine.transitions.test.ts
index 2de59909..c23cb0a2 100644
--- a/src/stores/chat/__tests__/chat-machine.transitions.test.ts
+++ b/src/stores/chat/__tests__/chat-machine.transitions.test.ts
@@ -529,6 +529,71 @@ describe("chatMachine transitions", () => {
     actor.stop();
   });
 
+  it("does not overwrite a user message when it shares the private message id", async () => {
+    const actor = createActor(
+      createTestChatMachine({
+        historyMessages: [
+          {
+            id: "msg-shared-id",
+            content: "User original question",
+            isFromAI: false,
+            date: "2026-06-29",
+          },
+          {
+            id: "msg-shared-id",
+            content: "",
+            isFromAI: true,
+            date: "2026-06-29",
+            locked: true,
+            lockReason: "private_message",
+            lockedPrivate: true,
+            privateMessageHint: "A private message is waiting.",
+          },
+        ],
+        unlockMessageOutput: {
+          messageId: "msg-shared-id",
+          response: makeUnlockPrivateResponse({
+            content: "This is the unlocked private message.",
+          }),
+        },
+      }),
+    ).start();
+
+    actor.send({ type: "ChatUserLogin", token: "token" });
+    await waitFor(actor, (snapshot) =>
+      snapshot.matches({ userSession: "ready" }),
+    );
+
+    actor.send({
+      type: "ChatUnlockMessageRequested",
+      messageId: "msg-shared-id",
+      kind: "private",
+    });
+
+    await waitFor(actor, (snapshot) =>
+      snapshot.matches({ userSession: "ready" }),
+    );
+
+    expect(actor.getSnapshot().context.messages).toMatchObject([
+      {
+        id: "msg-shared-id",
+        content: "User original question",
+        isFromAI: false,
+      },
+      {
+        id: "msg-shared-id",
+        content: "This is the unlocked private message.",
+        isFromAI: true,
+        locked: false,
+        lockReason: null,
+        lockedPrivate: false,
+        privateMessageHint: null,
+      },
+    ]);
+
+    actor.stop();
+  });
+
   it("requests payment when single message unlock lacks credits", async () => {
     const actor = createActor(
       createTestChatMachine({
diff --git a/src/stores/chat/chat-unlock-helpers.ts b/src/stores/chat/chat-unlock-helpers.ts
index ad34a973..e4477536 100644
--- a/src/stores/chat/chat-unlock-helpers.ts
+++ b/src/stores/chat/chat-unlock-helpers.ts
@@ -16,7 +16,7 @@ export function applySingleUnlockOutput(
   if (!output.response.unlocked) return [...messages];
 
   return messages.map((message) => {
-    if (message.id !== output.messageId) return message;
+    if (!shouldApplySingleUnlock(message, output.messageId)) return message;
 
     const nextContent =
       output.response.content.length > 0
@@ -38,6 +38,17 @@ export function applySingleUnlockOutput(
   });
 }
 
+function shouldApplySingleUnlock(message: UiMessage, messageId: string): boolean {
+  if (message.id !== messageId) return false;
+  if (!message.isFromAI) return false;
+  if (message.locked !== true) return false;
+  return (
+    message.imagePaywalled === true ||
+    message.lockReason === "private_message" ||
+    message.lockReason === "voice_message"
+  );
+}
+
 export function shouldAutoUnlockHistory(messages: readonly UiMessage[]): boolean {
   return countLockedHistoryMessages(messages) === 1;
 }