diff --git a/packages/core/realtime/use-realtime-sync.test.ts b/packages/core/realtime/use-realtime-sync.test.ts new file mode 100644 index 0000000000..1afb3873f3 --- /dev/null +++ b/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 { + 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(messagesKey, [userMessage()]); + qc.setQueryData(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(pendingKey)).toEqual({}); + expect(qc.getQueryData(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(messagesKey, [userMessage(), assistant]); + qc.setQueryData(pendingKey, { + task_id: taskId, + status: "running", + }); + + applyChatDoneToCache(qc, donePayload()); + + expect(qc.getQueryData(messagesKey)).toEqual([ + userMessage(), + assistant, + ]); + expect(qc.getQueryData(pendingKey)).toEqual({}); + }); + + it("falls back to invalidation-only when older servers omit message fields", () => { + const qc = createQueryClient(); + qc.setQueryData(messagesKey, [userMessage()]); + qc.setQueryData(pendingKey, { + task_id: taskId, + status: "running", + }); + + applyChatDoneToCache( + qc, + donePayload({ message_id: undefined, content: undefined }), + ); + + expect(qc.getQueryData(messagesKey)).toEqual([ + userMessage(), + ]); + expect(qc.getQueryData(pendingKey)).toEqual({}); + }); +}); diff --git a/packages/core/realtime/use-realtime-sync.ts b/packages/core/realtime/use-realtime-sync.ts index 8270c91942..9893be4409 100644 --- a/packages/core/realtime/use-realtime-sync.ts +++ b/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( + 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>; } @@ -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( - 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();