test(chat): cover chat done cache handoff

Co-authored-by: multica-agent <github@multica.ai>

packages/core/realtime/use-realtime-sync.test.ts

@@ -0,0 +1,117 @@
+import { QueryClient } from "@tanstack/react-query";
+import { describe, expect, it, vi } from "vitest";
+import { chatKeys } from "../chat/queries";
+import type { ChatDonePayload, ChatMessage, ChatPendingTask } from "../types";
+import { applyChatDoneToCache } from "./use-realtime-sync";
+
+const sessionId = "session-1";
+const taskId = "task-1";
+const messagesKey = chatKeys.messages(sessionId);
+const pendingKey = chatKeys.pendingTask(sessionId);
+
+function createQueryClient() {
+  return new QueryClient({
+    defaultOptions: {
+      queries: { retry: false },
+    },
+  });
+}
+
+function userMessage(): ChatMessage {
+  return {
+    id: "msg-user",
+    chat_session_id: sessionId,
+    role: "user",
+    content: "hello",
+    task_id: null,
+    created_at: "2026-05-13T05:00:00Z",
+  };
+}
+
+function donePayload(overrides: Partial<ChatDonePayload> = {}): ChatDonePayload {
+  return {
+    chat_session_id: sessionId,
+    task_id: taskId,
+    message_id: "msg-assistant",
+    content: "done",
+    elapsed_ms: 1234,
+    created_at: "2026-05-13T05:00:02Z",
+    ...overrides,
+  };
+}
+
+describe("applyChatDoneToCache", () => {
+  it("writes the assistant message before clearing pending task", () => {
+    const qc = createQueryClient();
+    qc.setQueryData<ChatMessage[]>(messagesKey, [userMessage()]);
+    qc.setQueryData<ChatPendingTask>(pendingKey, {
+      task_id: taskId,
+      status: "running",
+    });
+
+    const setQueryData = vi.spyOn(qc, "setQueryData");
+
+    applyChatDoneToCache(qc, donePayload());
+
+    expect(setQueryData.mock.calls[0]?.[0]).toEqual(messagesKey);
+    expect(setQueryData.mock.calls[1]?.[0]).toEqual(pendingKey);
+    expect(qc.getQueryData<ChatPendingTask>(pendingKey)).toEqual({});
+    expect(qc.getQueryData<ChatMessage[]>(messagesKey)).toEqual([
+      userMessage(),
+      {
+        id: "msg-assistant",
+        chat_session_id: sessionId,
+        role: "assistant",
+        content: "done",
+        task_id: taskId,
+        created_at: "2026-05-13T05:00:02Z",
+        elapsed_ms: 1234,
+      },
+    ]);
+  });
+
+  it("does not duplicate a replayed chat done event", () => {
+    const qc = createQueryClient();
+    const assistant: ChatMessage = {
+      id: "msg-assistant",
+      chat_session_id: sessionId,
+      role: "assistant",
+      content: "done",
+      task_id: taskId,
+      created_at: "2026-05-13T05:00:02Z",
+      elapsed_ms: 1234,
+    };
+    qc.setQueryData<ChatMessage[]>(messagesKey, [userMessage(), assistant]);
+    qc.setQueryData<ChatPendingTask>(pendingKey, {
+      task_id: taskId,
+      status: "running",
+    });
+
+    applyChatDoneToCache(qc, donePayload());
+
+    expect(qc.getQueryData<ChatMessage[]>(messagesKey)).toEqual([
+      userMessage(),
+      assistant,
+    ]);
+    expect(qc.getQueryData<ChatPendingTask>(pendingKey)).toEqual({});
+  });
+
+  it("falls back to invalidation-only when older servers omit message fields", () => {
+    const qc = createQueryClient();
+    qc.setQueryData<ChatMessage[]>(messagesKey, [userMessage()]);
+    qc.setQueryData<ChatPendingTask>(pendingKey, {
+      task_id: taskId,
+      status: "running",
+    });
+
+    applyChatDoneToCache(
+      qc,
+      donePayload({ message_id: undefined, content: undefined }),
+    );
+
+    expect(qc.getQueryData<ChatMessage[]>(messagesKey)).toEqual([
+      userMessage(),
+    ]);
+    expect(qc.getQueryData<ChatPendingTask>(pendingKey)).toEqual({});
+  });
+});

packages/core/realtime/use-realtime-sync.ts

@@ -1,7 +1,7 @@
 "use client";
 
 import { useEffect, useRef } from "react";
-import { useQueryClient } from "@tanstack/react-query";
+import { useQueryClient, type QueryClient } from "@tanstack/react-query";
 import type { WSClient } from "../api/ws-client";
 import type { StoreApi, UseBoundStore } from "zustand";
 import type { AuthState } from "../auth/store";
@@ -71,6 +71,42 @@ const chatWsLogger = createLogger("chat.ws");
 
 const logger = createLogger("realtime-sync");
 
+export function applyChatDoneToCache(
+  qc: QueryClient,
+  payload: ChatDonePayload,
+) {
+  const sessionId = payload.chat_session_id;
+  const taskId = payload.task_id;
+  const messageId = payload.message_id;
+  const content = payload.content;
+  if (messageId && content !== undefined) {
+    qc.setQueryData<ChatMessage[] | undefined>(
+      chatKeys.messages(sessionId),
+      (old) => {
+        if (!old) return old; // first fetch will pick it up
+        // Idempotent against reconnect replay.
+        if (old.some((m) => m.id === messageId)) return old;
+        const assistant: ChatMessage = {
+          id: messageId,
+          chat_session_id: sessionId,
+          role: "assistant",
+          content,
+          task_id: taskId,
+          created_at: payload.created_at ?? new Date().toISOString(),
+          elapsed_ms: payload.elapsed_ms ?? null,
+        };
+        return [...old, assistant];
+      },
+    );
+  }
+  // Replacement is in the messages list now; safe to drop pending.
+  qc.setQueryData(chatKeys.pendingTask(sessionId), {});
+  // Authoritative refetch reconciles redaction / migrations / clients
+  // that took the fallback branch above.
+  qc.invalidateQueries({ queryKey: chatKeys.messages(sessionId) });
+  qc.invalidateQueries({ queryKey: chatKeys.pendingTask(sessionId) });
+}
+
 export interface RealtimeSyncStores {
   authStore: UseBoundStore<StoreApi<AuthState>>;
 }
@@ -583,34 +619,7 @@ export function useRealtimeSync(
       // payload (older builds). Older clients hitting a newer server also
       // work: they ignore the extra fields and rely on the invalidate
       // below, which keeps the old behavior alive.
-      const sessionId = payload.chat_session_id;
-      const taskId = payload.task_id;
-      if (payload.message_id && payload.content !== undefined) {
-        qc.setQueryData<ChatMessage[] | undefined>(
-          chatKeys.messages(sessionId),
-          (old) => {
-            if (!old) return old; // first fetch will pick it up
-            // Idempotent against reconnect replay.
-            if (old.some((m) => m.id === payload.message_id)) return old;
-            const assistant: ChatMessage = {
-              id: payload.message_id!,
-              chat_session_id: sessionId,
-              role: "assistant",
-              content: payload.content ?? "",
-              task_id: taskId,
-              created_at: payload.created_at ?? new Date().toISOString(),
-              elapsed_ms: payload.elapsed_ms ?? null,
-            };
-            return [...old, assistant];
-          },
-        );
-      }
-      // Replacement is in the messages list now; safe to drop pending.
-      qc.setQueryData(chatKeys.pendingTask(sessionId), {});
-      // Authoritative refetch reconciles redaction / migrations / clients
-      // that took the fallback branch above.
-      qc.invalidateQueries({ queryKey: chatKeys.messages(sessionId) });
-      qc.invalidateQueries({ queryKey: chatKeys.pendingTask(sessionId) });
+      applyChatDoneToCache(qc, payload);
       invalidatePendingAggregate();
       // Assistant message just landed → has_unread may have flipped to true.
       invalidateSessionLists();