fix(chat): stabilize message identities

src/app/chat/hooks/use-chat-unlock-coordinator.ts

@@ -19,6 +19,7 @@ import {
 } from "@/stores/chat/chat-context";
 import type { ChatPromotionState } from "@/stores/chat/helper/promotion";
 import type { ChatUnlockPaywallRequest } from "@/stores/chat/chat-state";
+import type { UiMessage } from "@/stores/chat/ui-message";
 import { useUserSelector } from "@/stores/user/user-context";
 
 import { getInsufficientCreditsSubscriptionType } from "../chat-screen.helpers";
@@ -60,10 +61,11 @@ export interface UseChatUnlockCoordinatorInput
 }
 
 export interface UseChatUnlockCoordinatorOutput {
-  requestMessageUnlock: (
-    messageId: string,
-    kind: PendingChatUnlockKind,
-  ) => void;
+  requestMessageUnlock: (input: {
+    displayMessageId: string;
+    remoteMessageId?: string;
+    kind: PendingChatUnlockKind;
+  }) => void;
   dialogs: ChatUnlockDialogModel;
 }
 
@@ -81,6 +83,7 @@ export function useChatUnlockCoordinator({
     (state) => ({
       characterId: state.context.characterId,
       historyLoaded: state.context.historyLoaded,
+      messages: state.context.messages,
       isUnlockingHistory: state.matches({ userSession: "unlockingHistory" }),
       lockedHistoryCount: state.context.lockedHistoryCount,
       unlockHistoryError: state.context.unlockHistoryError,
@@ -141,10 +144,15 @@ export function useChatUnlockCoordinator({
 
       const consumed = await consumePendingChatUnlock(chatState.characterId);
       if (cancelled || !consumed) return;
+      const displayMessageId = resolvePendingUnlockDisplayMessageId(
+        consumed,
+        chatState.messages,
+        promotion,
+      );
 
       chatDispatch({
         type: "ChatUnlockMessageRequested",
-        messageId: consumed.displayMessageId,
+        displayMessageId,
         remoteMessageId: consumed.messageId,
         kind: consumed.kind,
         lockType: consumed.lockType,
@@ -165,14 +173,17 @@ export function useChatUnlockCoordinator({
     imageMessageId,
     imageReturnUrl,
     navigator.isAuthenticatedUser,
+    promotion,
+    chatState.messages,
   ]);
 
-  function requestMessageUnlock(
-    messageId: string,
-    kind: PendingChatUnlockKind,
-  ): void {
+  function requestMessageUnlock(input: {
+    displayMessageId: string;
+    remoteMessageId?: string;
+    kind: PendingChatUnlockKind;
+  }): void {
     const request = resolveMessageUnlockRequest(
-      { messageId, kind },
+      input,
       {
         defaultReturnUrl,
         imageMessageId,
@@ -186,7 +197,7 @@ export function useChatUnlockCoordinator({
       onAuthenticated: () => {
         chatDispatch({
           type: "ChatUnlockMessageRequested",
-          messageId: request.displayMessageId,
+          displayMessageId: request.displayMessageId,
           remoteMessageId: request.messageId,
           kind: request.kind,
           lockType: request.lockType,
@@ -258,23 +269,21 @@ export function useChatUnlockCoordinator({
 
 export function resolveMessageUnlockRequest(
   input: {
-    messageId: string;
+    displayMessageId: string;
+    remoteMessageId?: string;
     kind: PendingChatUnlockKind;
   },
   scope: ChatUnlockCoordinatorScope,
 ): CoordinatedMessageUnlockRequest {
   const matchedPromotion =
-    scope.promotion?.message.id === input.messageId
+    scope.promotion?.message.displayId === input.displayMessageId
       ? scope.promotion
       : null;
-  const temporaryPromotionMessageId = matchedPromotion
-    ? `promotion:${matchedPromotion.session.clientLockId}`
-    : null;
+  const remoteMessageId =
+    input.remoteMessageId ?? matchedPromotion?.message.remoteId;
   const target = {
-    displayMessageId: input.messageId,
-    ...(input.messageId !== temporaryPromotionMessageId
-      ? { messageId: input.messageId }
-      : {}),
+    displayMessageId: input.displayMessageId,
+    ...(remoteMessageId ? { messageId: remoteMessageId } : {}),
     kind: input.kind,
     ...(matchedPromotion
       ? {
@@ -300,6 +309,40 @@ export function shouldResumePendingChatUnlock(
   return isCurrentImageUnlock(pending, scope.imageMessageId);
 }
 
+export function resolvePendingUnlockDisplayMessageId(
+  pending: PendingChatUnlock,
+  messages: readonly UiMessage[],
+  promotion: ChatPromotionState | null,
+): string {
+  if (
+    promotion &&
+    (promotion.message.displayId === pending.displayMessageId ||
+      (pending.messageId &&
+        promotion.message.remoteId === pending.messageId))
+  ) {
+    return promotion.message.displayId;
+  }
+
+  const exactMessage = messages.find(
+    (message) => message.displayId === pending.displayMessageId,
+  );
+  if (exactMessage) return exactMessage.displayId;
+  if (!pending.messageId) return pending.displayMessageId;
+
+  const remoteMessage =
+    messages.find(
+      (message) =>
+        message.isFromAI &&
+        message.locked === true &&
+        message.remoteId === pending.messageId,
+    ) ??
+    messages.find(
+      (message) =>
+        message.isFromAI && message.remoteId === pending.messageId,
+    );
+  return remoteMessage?.displayId ?? pending.displayMessageId;
+}
+
 export function resolveChatUnlockReturnUrl(
   target: Pick<
     CoordinatedMessageUnlockRequest,