From c2e88921949e45b61ba6725d499f7f01e5031e68 Mon Sep 17 00:00:00 2001
From: Ryan Yu <69830650+Ryannn41@users.noreply.github.com>
Date: Mon, 29 Jun 2026 15:34:36 +0800
Subject: [PATCH] fix(chat): refresh message caches on reconnect (MUL-3831)
 (#4677)

Invalidate per-session chat message caches (messages, messages-page,
pending-task, task-messages) on websocket reconnect / WS instance change so
a chat that missed chat/task events while disconnected recovers without a
full reload, matching the existing per-issue recovery pattern.

Co-authored-by: Ryan <1141524679@qq.com>
---
 packages/core/chat/queries.ts                 | 12 +++++---
 .../use-realtime-sync-ws-instance.test.tsx    | 28 +++++++++++++++++--
 packages/core/realtime/use-realtime-sync.ts   |  8 ++++++
 3 files changed, 41 insertions(+), 7 deletions(-)

diff --git a/packages/core/chat/queries.ts b/packages/core/chat/queries.ts
index f151fbdad..f8eb8f52b 100644
--- a/packages/core/chat/queries.ts
+++ b/packages/core/chat/queries.ts
@@ -14,13 +14,17 @@ export const chatKeys = {
   /** Full sessions list (active + archived); the dropdown splits locally. */
   sessions: (wsId: string) => [...chatKeys.all(wsId), "sessions"] as const,
   session: (wsId: string, id: string) => [...chatKeys.all(wsId), "session", id] as const,
-  messages: (sessionId: string) => ["chat", "messages", sessionId] as const,
-  messagesPage: (sessionId: string) => ["chat", "messages-page", sessionId] as const,
-  pendingTask: (sessionId: string) => ["chat", "pending-task", sessionId] as const,
+  messagesAll: () => ["chat", "messages"] as const,
+  messages: (sessionId: string) => [...chatKeys.messagesAll(), sessionId] as const,
+  messagesPageAll: () => ["chat", "messages-page"] as const,
+  messagesPage: (sessionId: string) => [...chatKeys.messagesPageAll(), sessionId] as const,
+  pendingTaskAll: () => ["chat", "pending-task"] as const,
+  pendingTask: (sessionId: string) => [...chatKeys.pendingTaskAll(), sessionId] as const,
   /** Aggregate of in-flight chat tasks for the current user — FAB reads this. */
   pendingTasks: (wsId: string) => [...chatKeys.all(wsId), "pending-tasks"] as const,
   /** Per-task execution messages — shared with issue agent cards. */
-  taskMessages: (taskId: string) => ["task-messages", taskId] as const,
+  taskMessagesAll: () => ["task-messages"] as const,
+  taskMessages: (taskId: string) => [...chatKeys.taskMessagesAll(), taskId] as const,
 };
 
 const UUID_PATTERN = /^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$/i;
diff --git a/packages/core/realtime/use-realtime-sync-ws-instance.test.tsx b/packages/core/realtime/use-realtime-sync-ws-instance.test.tsx
index 94b451a92..b40d0d8f9 100644
--- a/packages/core/realtime/use-realtime-sync-ws-instance.test.tsx
+++ b/packages/core/realtime/use-realtime-sync-ws-instance.test.tsx
@@ -102,9 +102,9 @@ describe("useRealtimeSync — ws instance change", () => {
     rerender({ ws: ws2 });
 
     // Should have called invalidateQueries for all workspace-scoped keys
-    // (15 workspace-scoped + 6 per-issue prefixes + 1 workspaceKeys.list()
-    // + 1 cross-workspace inbox unread summary = 23 calls)
-    expect(invalidateSpy).toHaveBeenCalledTimes(23);
+    // (15 workspace-scoped + 6 per-issue prefixes + 4 per-chat prefixes
+    // + 1 workspaceKeys.list() + 1 cross-workspace inbox unread summary = 27 calls)
+    expect(invalidateSpy).toHaveBeenCalledTimes(27);
   });
 
   it("does not re-invalidate when rerendered with the same ws instance", () => {
@@ -164,4 +164,26 @@ describe("useRealtimeSync — ws instance change", () => {
     expect(calls).toContainEqual(["issues", "attachments"]);
     expect(calls).toContainEqual(["issues", "tasks"]);
   });
+
+  it("invalidates per-chat-session caches (no wsId in key) on ws instance change", () => {
+    // These keys are not under the ["chat", wsId] prefix, so they need their
+    // own recovery invalidation when reconnecting after missed chat/task events.
+    const ws1 = createMockWs();
+    const { rerender } = renderHook(
+      ({ ws }) => useRealtimeSync(ws, stores),
+      { initialProps: { ws: ws1 as WSClient | null }, wrapper: createWrapper(qc) },
+    );
+
+    invalidateSpy.mockClear();
+    rerender({ ws: null });
+
+    const ws2 = createMockWs();
+    rerender({ ws: ws2 });
+
+    const calls = invalidateSpy.mock.calls.map((call: [{ queryKey?: unknown }, ...unknown[]]) => call[0].queryKey);
+    expect(calls).toContainEqual(["chat", "messages"]);
+    expect(calls).toContainEqual(["chat", "messages-page"]);
+    expect(calls).toContainEqual(["chat", "pending-task"]);
+    expect(calls).toContainEqual(["task-messages"]);
+  });
 });
diff --git a/packages/core/realtime/use-realtime-sync.ts b/packages/core/realtime/use-realtime-sync.ts
index a292032fb..a4f07c094 100644
--- a/packages/core/realtime/use-realtime-sync.ts
+++ b/packages/core/realtime/use-realtime-sync.ts
@@ -340,6 +340,14 @@ function invalidateWorkspaceScopedQueries(qc: QueryClient): void {
   qc.invalidateQueries({ queryKey: issueKeys.usageAll() });
   qc.invalidateQueries({ queryKey: issueKeys.attachmentsAll() });
   qc.invalidateQueries({ queryKey: issueKeys.tasksAll() });
+  // Per-chat-session caches are also keyed without wsId, so the
+  // chatKeys.all(wsId) prefix above only reaches session lists / aggregates.
+  // Message streams rely on WS invalidation with staleTime: Infinity; recover
+  // sessions that missed chat/task events while the socket was disconnected.
+  qc.invalidateQueries({ queryKey: chatKeys.messagesAll() });
+  qc.invalidateQueries({ queryKey: chatKeys.messagesPageAll() });
+  qc.invalidateQueries({ queryKey: chatKeys.pendingTaskAll() });
+  qc.invalidateQueries({ queryKey: chatKeys.taskMessagesAll() });
   qc.invalidateQueries({ queryKey: workspaceKeys.list() });
 }