diff --git a/frontend/src/components/chat/chat-session-pane.tsx b/frontend/src/components/chat/chat-session-pane.tsx index f1f1a9a5f..a87b4bf22 100644 --- a/frontend/src/components/chat/chat-session-pane.tsx +++ b/frontend/src/components/chat/chat-session-pane.tsx @@ -429,6 +429,8 @@ export function ChatSessionPane({ const displayedError = lastError ?? persistedError useAdoptServerTranscript({ + chatId: chat?.id, + workspaceId, status, serverMessages: uiMessages, liveMessages: messages, diff --git a/frontend/src/hooks/use-chat.test.tsx b/frontend/src/hooks/use-chat.test.tsx new file mode 100644 index 000000000..9a055d068 --- /dev/null +++ b/frontend/src/hooks/use-chat.test.tsx @@ -0,0 +1,553 @@ +import { QueryClient, QueryClientProvider } from "@tanstack/react-query" +import { act, renderHook } from "@testing-library/react" +import type { ChatStatus, UIMessage } from "ai" +import type { ReactNode } from "react" +import { + decideServerTranscriptAdoption, + useAdoptServerTranscript, +} from "@/hooks/use-chat" + +function textMessage( + id: string, + role: UIMessage["role"], + text: string +): UIMessage { + return { id, role, parts: [{ type: "text", text }] } +} + +function approvalMessage(id: string): UIMessage { + return { + id, + role: "assistant", + parts: [ + { + type: "data-approval-request", + data: [], + } as UIMessage["parts"][number], + ], + } +} + +function toolMessage(id: string, toolCallId: string): UIMessage { + return { + id, + role: "assistant", + parts: [ + { + type: "dynamic-tool", + toolName: "lookup", + toolCallId, + state: "input-available", + input: {}, + }, + ], + } +} + +function createWrapper(queryClient: QueryClient) { + return function QueryWrapper({ children }: { children: ReactNode }) { + return ( + {children} + ) + } +} + +async function advanceTimersBy(ms: number): Promise { + await act(async () => { + jest.advanceTimersByTime(ms) + await Promise.resolve() + }) +} + +describe("useAdoptServerTranscript", () => { + beforeEach(() => { + jest.useFakeTimers() + }) + + afterEach(() => { + jest.clearAllTimers() + jest.useRealTimers() + }) + + it("retains final assistant text missing from an equal-length snapshot and schedules a refetch", async () => { + const queryClient = new QueryClient() + const invalidateQueries = jest.spyOn(queryClient, "invalidateQueries") + const setMessages = jest.fn() + const liveMessages = [ + textMessage("live-user", "user", "Question"), + textMessage("live-assistant", "assistant", "Final streamed answer"), + ] + const serverMessages = [ + textMessage("server-user", "user", "Question"), + textMessage("server-assistant", "assistant", "Previous answer"), + ] + + const { unmount } = renderHook( + () => + useAdoptServerTranscript({ + chatId: "chat-1", + workspaceId: "workspace-1", + status: "ready", + serverMessages, + liveMessages, + setMessages, + }), + { wrapper: createWrapper(queryClient) } + ) + + expect(setMessages).not.toHaveBeenCalled() + expect(jest.getTimerCount()).toBe(3) + + await advanceTimersBy(1_000) + + expect(invalidateQueries).toHaveBeenCalledTimes(1) + expect(invalidateQueries).toHaveBeenCalledWith({ + queryKey: ["chat", "chat-1", "workspace-1", "vercel"], + }) + expect(setMessages).not.toHaveBeenCalled() + unmount() + }) + + it("adopts a later split-row snapshot and cancels remaining retries", async () => { + const queryClient = new QueryClient() + const invalidateQueries = jest.spyOn(queryClient, "invalidateQueries") + const setMessages = jest.fn() + const liveMessages = [ + textMessage("live-user", "user", "Question"), + textMessage("live-assistant", "assistant", "Final streamed answer"), + ] + const staleMessages = [ + textMessage("stale-user", "user", "Question"), + textMessage("stale-assistant", "assistant", "Previous answer"), + ] + const caughtUpMessages = [ + textMessage("server-user", "user", "Question"), + textMessage("server-assistant-1", "assistant", "Final streamed "), + textMessage("server-assistant-2", "assistant", "answer"), + ] + + const { rerender, unmount } = renderHook( + ({ serverMessages }: { serverMessages: UIMessage[] }) => + useAdoptServerTranscript({ + chatId: "chat-1", + workspaceId: "workspace-1", + status: "ready", + serverMessages, + liveMessages, + setMessages, + }), + { + initialProps: { serverMessages: staleMessages }, + wrapper: createWrapper(queryClient), + } + ) + + await advanceTimersBy(1_000) + expect(invalidateQueries).toHaveBeenCalledTimes(1) + + rerender({ serverMessages: caughtUpMessages }) + + expect(setMessages).toHaveBeenCalledWith(caughtUpMessages) + expect(jest.getTimerCount()).toBe(0) + + await advanceTimersBy(8_000) + expect(invalidateQueries).toHaveBeenCalledTimes(1) + unmount() + }) + + it("adopts a resolved-approval snapshot after its approval card is dropped", () => { + const queryClient = new QueryClient() + const setMessages = jest.fn() + const liveMessages = [ + textMessage("live-user", "user", "Run the action"), + textMessage("live-assistant", "assistant", "Approval required"), + approvalMessage("live-approval"), + ] + const serverMessages = [ + textMessage("server-user", "user", "Run the action"), + textMessage("server-assistant", "assistant", "Approval required"), + ] + + const { unmount } = renderHook( + () => + useAdoptServerTranscript({ + chatId: "chat-1", + workspaceId: "workspace-1", + status: "ready", + serverMessages, + liveMessages, + setMessages, + }), + { wrapper: createWrapper(queryClient) } + ) + + expect(setMessages).toHaveBeenCalledWith(serverMessages) + expect(jest.getTimerCount()).toBe(0) + unmount() + }) + + it("adopts a count-covered snapshot after the bounded retries are exhausted", async () => { + const queryClient = new QueryClient() + const invalidateQueries = jest.spyOn(queryClient, "invalidateQueries") + const setMessages = jest.fn() + const queryKey = ["chat", "chat-1", "workspace-1", "vercel"] + const liveMessages = [ + textMessage("live-user", "user", "Question"), + textMessage("live-assistant", "assistant", "Stream-only wording"), + ] + const serverMessages = [ + textMessage("server-user", "user", "Question"), + textMessage("server-assistant", "assistant", "Canonical wording"), + ] + + const { unmount } = renderHook( + () => + useAdoptServerTranscript({ + chatId: "chat-1", + workspaceId: "workspace-1", + status: "ready", + serverMessages, + liveMessages, + setMessages, + }), + { wrapper: createWrapper(queryClient) } + ) + + // Each retry only counts toward guard exhaustion when the query delivered + // fresh data, so simulate a completed refetch before each timer fires. + await advanceTimersBy(1) + act(() => { + queryClient.setQueryData(queryKey, serverMessages) + }) + await advanceTimersBy(999) + expect(setMessages).not.toHaveBeenCalled() + + act(() => { + queryClient.setQueryData(queryKey, serverMessages) + }) + await advanceTimersBy(2_000) + expect(setMessages).not.toHaveBeenCalled() + + act(() => { + queryClient.setQueryData(queryKey, serverMessages) + }) + await advanceTimersBy(5_000) + + expect(invalidateQueries).toHaveBeenCalledTimes(3) + expect(setMessages).toHaveBeenCalledWith(serverMessages) + // No further retry invalidations after adoption (remaining timers belong + // to the query cache's own GC, not the retry episode). + await advanceTimersBy(10_000) + expect(invalidateQueries).toHaveBeenCalledTimes(3) + unmount() + }) + + it("does not exhaust the content guard when refetches deliver no fresh data", async () => { + const queryClient = new QueryClient() + const invalidateQueries = jest.spyOn(queryClient, "invalidateQueries") + const setMessages = jest.fn() + const liveMessages = [ + textMessage("live-user", "user", "Question"), + textMessage("live-assistant", "assistant", "Stream-only wording"), + ] + const serverMessages = [ + textMessage("server-user", "user", "Question"), + textMessage("server-assistant", "assistant", "Canonical wording"), + ] + + const { unmount } = renderHook( + () => + useAdoptServerTranscript({ + chatId: "chat-1", + workspaceId: "workspace-1", + status: "ready", + serverMessages, + liveMessages, + setMessages, + }), + { wrapper: createWrapper(queryClient) } + ) + + // No query data ever lands (transient outage): invalidations run but must + // not advance the guard toward stale adoption. + await advanceTimersBy(1_000) + await advanceTimersBy(2_000) + await advanceTimersBy(5_000) + await advanceTimersBy(10_000) + + expect(invalidateQueries).toHaveBeenCalledTimes(3) + expect(setMessages).not.toHaveBeenCalled() + unmount() + }) + + it("uses only the count guard for a textless tool-only final assistant message", () => { + const liveMessages = [ + textMessage("live-user", "user", "Look this up"), + toolMessage("live-tool", "live-call"), + ] + const countCoveredServerMessages = [ + textMessage("server-user", "user", "Look this up"), + toolMessage("server-tool", "server-call"), + ] + + expect( + decideServerTranscriptAdoption({ + serverMessages: countCoveredServerMessages, + liveMessages, + }) + ).toBe("adopt") + expect( + decideServerTranscriptAdoption({ + serverMessages: countCoveredServerMessages.slice(0, 1), + liveMessages, + }) + ).toBe("reject-count") + }) + + it("does not adopt when only the final retry delivered fresh data", async () => { + const queryClient = new QueryClient() + const invalidateQueries = jest.spyOn(queryClient, "invalidateQueries") + const setMessages = jest.fn() + const queryKey = ["chat", "chat-1", "workspace-1", "vercel"] + const liveMessages = [ + textMessage("live-user", "user", "Question"), + textMessage("live-assistant", "assistant", "Stream-only wording"), + ] + const serverMessages = [ + textMessage("server-user", "user", "Question"), + textMessage("server-assistant", "assistant", "Canonical wording"), + ] + + const { rerender, unmount } = renderHook( + ({ server }: { server: UIMessage[] }) => + useAdoptServerTranscript({ + chatId: "chat-1", + workspaceId: "workspace-1", + status: "ready", + serverMessages: server, + liveMessages, + setMessages, + }), + { + initialProps: { server: serverMessages }, + wrapper: createWrapper(queryClient), + } + ) + + // The first two refetches fail (no fresh data); only the third delivers. + // One successful refetch must not exhaust a guard that promises three — + // counting the timer slot index would retroactively credit the failures. + await advanceTimersBy(1_000) + await advanceTimersBy(2_000) + act(() => { + queryClient.setQueryData(queryKey, serverMessages) + }) + await advanceTimersBy(5_000) + + expect(invalidateQueries).toHaveBeenCalledTimes(3) + expect(setMessages).not.toHaveBeenCalled() + + // The failed series retires the episode: a later snapshot change starts a + // fresh retry series instead of leaving eventual adoption unreachable. + rerender({ + server: [ + textMessage("server-user", "user", "Question"), + textMessage("server-assistant", "assistant", "Rewritten wording"), + ], + }) + await advanceTimersBy(1_000) + expect(invalidateQueries).toHaveBeenCalledTimes(4) + expect(setMessages).not.toHaveBeenCalled() + unmount() + }) + + it("rejects a stale snapshot whose earlier turn repeats the final answer text", () => { + const liveMessages = [ + textMessage("live-user-1", "user", "Question"), + textMessage("live-assistant-1", "assistant", "Same answer"), + textMessage("live-user-2", "user", "Question again"), + textMessage("live-assistant-2", "assistant", "Same answer"), + ] + // Stale finalize-race snapshot: the earlier turn already contains the + // identical text and DB row segmentation makes the snapshot as long as the + // live transcript, but the final turn's rows are still hidden. Coverage + // must be scoped to the final turn or this would be adopted and the final + // bubble dropped. + const serverMessages = [ + textMessage("server-user-1", "user", "Question"), + textMessage("server-assistant-1a", "assistant", "Same "), + textMessage("server-assistant-1b", "assistant", "answer"), + textMessage("server-user-2", "user", "Question again"), + ] + + expect( + decideServerTranscriptAdoption({ serverMessages, liveMessages }) + ).toBe("reject-content") + }) + + it("rejects a stale snapshot that omits the whole final turn including its prompt", () => { + const liveMessages = [ + textMessage("live-user-1", "user", "Question"), + textMessage("live-assistant-1", "assistant", "Same answer"), + textMessage("live-user-2", "user", "Question again"), + textMessage("live-assistant-2", "assistant", "Same answer"), + ] + // The mid-turn DB filter hides every row of the active run, so the stale + // snapshot lacks even the final turn's user prompt. Old-turn segmentation + // satisfies the count guard and the previous answer repeats the final + // text, so only the prompt anchor can prove the turn is missing. + const serverMessages = [ + textMessage("server-user-1", "user", "Question"), + textMessage("server-assistant-1a", "assistant", "Same "), + textMessage("server-assistant-1b", "assistant", "ans"), + textMessage("server-assistant-1c", "assistant", "wer"), + ] + + expect( + decideServerTranscriptAdoption({ serverMessages, liveMessages }) + ).toBe("reject-content") + }) + + it("rejects a stale snapshot when a repeated prompt hides the missing final turn", () => { + // The user asked the same question twice and got the same answer. The + // stale snapshot hides the whole final turn, but an earlier identical + // prompt would anchor-match and its identical answer would cover the + // probe. Only the occurrence count proves the final turn is missing. + const liveMessages = [ + textMessage("live-user-1", "user", "Question"), + textMessage("live-assistant-1", "assistant", "Same answer"), + textMessage("live-user-2", "user", "Question"), + textMessage("live-assistant-2", "assistant", "Same answer"), + ] + const serverMessages = [ + textMessage("server-user-1", "user", "Question"), + textMessage("server-assistant-1a", "assistant", "Same "), + textMessage("server-assistant-1b", "assistant", "ans"), + textMessage("server-assistant-1c", "assistant", "wer"), + ] + + expect( + decideServerTranscriptAdoption({ serverMessages, liveMessages }) + ).toBe("reject-content") + }) + + it("uses only the count guard when the live transcript ends with a user message", () => { + // The finalize race protects the final streamed assistant turn; when the + // last live message is a user prompt there is no such turn, so a previous + // turn's assistant text must not be used as the coverage probe. + const liveMessages = [ + textMessage("live-user-1", "user", "Question"), + textMessage("live-assistant-1", "assistant", "Answer"), + textMessage("live-user-2", "user", "New question"), + ] + const serverMessages = [ + textMessage("server-user-1", "user", "Question"), + textMessage("server-assistant-1", "assistant", "Different text"), + textMessage("server-user-2", "user", "New question"), + ] + + expect( + decideServerTranscriptAdoption({ serverMessages, liveMessages }) + ).toBe("adopt") + }) + + it("adopts cancelled-turn marker rows that retain the streamed partial text", () => { + const queryClient = new QueryClient() + const setMessages = jest.fn() + const liveMessages = [ + textMessage("live-user", "user", "Long request"), + textMessage("live-assistant", "assistant", "Partial response"), + ] + const serverMessages = [ + textMessage("server-user", "user", "Long request"), + textMessage("server-assistant", "assistant", "Partial response"), + textMessage("server-cancelled", "assistant", "[Turn cancelled]"), + ] + + const { unmount } = renderHook( + () => + useAdoptServerTranscript({ + chatId: "chat-1", + workspaceId: "workspace-1", + status: "ready", + serverMessages, + liveMessages, + setMessages, + }), + { wrapper: createWrapper(queryClient) } + ) + + expect(setMessages).toHaveBeenCalledWith(serverMessages) + expect(jest.getTimerCount()).toBe(0) + unmount() + }) + + it("cancels retries when a new turn starts streaming", async () => { + const queryClient = new QueryClient() + const invalidateQueries = jest.spyOn(queryClient, "invalidateQueries") + const setMessages = jest.fn() + const liveMessages = [ + textMessage("live-user", "user", "Question"), + textMessage("live-assistant", "assistant", "Final streamed answer"), + ] + const serverMessages = [ + textMessage("server-user", "user", "Question"), + textMessage("server-assistant", "assistant", "Previous answer"), + ] + + const { rerender, unmount } = renderHook( + ({ status }: { status: ChatStatus }) => + useAdoptServerTranscript({ + chatId: "chat-1", + workspaceId: "workspace-1", + status, + serverMessages, + liveMessages, + setMessages, + }), + { + initialProps: { status: "ready" as ChatStatus }, + wrapper: createWrapper(queryClient), + } + ) + + expect(jest.getTimerCount()).toBe(3) + rerender({ status: "streaming" }) + expect(jest.getTimerCount()).toBe(0) + + await advanceTimersBy(8_000) + expect(invalidateQueries).not.toHaveBeenCalled() + expect(setMessages).not.toHaveBeenCalled() + unmount() + }) + + it("cancels retries when the hook unmounts", () => { + const queryClient = new QueryClient() + const setMessages = jest.fn() + const liveMessages = [ + textMessage("live-user", "user", "Question"), + textMessage("live-assistant", "assistant", "Final streamed answer"), + ] + const serverMessages = [ + textMessage("server-user", "user", "Question"), + textMessage("server-assistant", "assistant", "Previous answer"), + ] + + const { unmount } = renderHook( + () => + useAdoptServerTranscript({ + chatId: "chat-1", + workspaceId: "workspace-1", + status: "ready", + serverMessages, + liveMessages, + setMessages, + }), + { wrapper: createWrapper(queryClient) } + ) + + expect(jest.getTimerCount()).toBe(3) + unmount() + expect(jest.getTimerCount()).toBe(0) + }) +}) diff --git a/frontend/src/hooks/use-chat.ts b/frontend/src/hooks/use-chat.ts index 18bc275c5..856e9a15e 100644 --- a/frontend/src/hooks/use-chat.ts +++ b/frontend/src/hooks/use-chat.ts @@ -700,6 +700,184 @@ function isApprovalCardMessage(m: UIMessage): boolean { ) } +/** Delays between a rejected transcript adoption and its bounded refetches. */ +export const ADOPT_SERVER_TRANSCRIPT_RETRY_DELAYS_MS = [ + 1_000, 3_000, 8_000, +] as const + +/** The reason a server transcript should be adopted or retained. */ +export type ServerTranscriptAdoptionDecision = + | "already-current" + | "adopt" + | "reject-content" + | "reject-count" + +/** Index of the last user-role message, or -1 when there is none. */ +function lastUserMessageIndex(messages: UIMessage[]): number { + for (let index = messages.length - 1; index >= 0; index -= 1) { + if (messages[index].role === "user") { + return index + } + } + return -1 +} + +/** Concatenated text parts of a message ("" when it has none). */ +function messageText(message: UIMessage): string { + let text = "" + for (const part of message.parts) { + if (part.type === "text") { + text += part.text + } + } + return text +} + +/** + * Concatenate the text parts of the final non-approval assistant message. + * + * Only the final turn qualifies: the scan stops at the last user message, so + * an assistant answer from a previous turn is never used as the coverage + * probe. `null` means there is no comparable text, so callers must fall back + * to the count guard instead of comparing tool or data parts that serialize + * differently between the stream and the database. + */ +export function getFinalLiveAssistantText( + liveMessages: UIMessage[] +): string | null { + const turnStart = lastUserMessageIndex(liveMessages) + 1 + for (let index = liveMessages.length - 1; index >= turnStart; index -= 1) { + const message = liveMessages[index] + if (message.role !== "assistant" || isApprovalCardMessage(message)) { + continue + } + + let hasTextPart = false + let text = "" + for (const part of message.parts) { + if (part.type === "text") { + hasTextPart = true + text += part.text + } + } + return hasTextPart ? text : null + } + return null +} + +/** + * Whether the server snapshot contains the live transcript's final turn. + * + * The mid-turn DB filter hides EVERY row tagged with the active run — + * including the turn's user prompt — so a stale snapshot omits the whole + * final turn, and neither whole-transcript text search nor "after the + * server's last user row" reliably identifies the current turn. Coverage is + * therefore anchored on the live final turn's user prompt: the last server + * user row with identical text. A missing anchor means the snapshot predates + * the final turn. Given an anchor, the server text after it must contain the + * live final assistant text (concatenated across messages so one live bubble + * can match several database-backed rows). + * + * A promptless final turn (no live user message, or one without text) and a + * textless final assistant message are deliberately treated as covered and + * left to the count guard — tool, data, and approval parts serialize too + * differently between the stream and the database to compare. + */ +export function serverTranscriptCoversLiveFinalAssistantText( + serverMessages: UIMessage[], + liveMessages: UIMessage[] +): boolean { + const liveUserIndex = lastUserMessageIndex(liveMessages) + if (liveUserIndex === -1) { + return true + } + const liveUserText = messageText(liveMessages[liveUserIndex]) + if (liveUserText === "") { + return true + } + + // The anchor prompt may not be unique (a user can repeat a question + // verbatim), so presence alone cannot prove the final turn is in the + // snapshot. Require at least as many occurrences of the prompt as the live + // transcript has: a stale snapshot that hides the final turn is missing its + // occurrence and fails the count even when an earlier identical prompt — + // and an identical earlier answer — would otherwise match. + let liveOccurrences = 0 + for (const message of liveMessages) { + if (message.role === "user" && messageText(message) === liveUserText) { + liveOccurrences += 1 + } + } + let serverOccurrences = 0 + let anchorIndex = -1 + for (const [index, message] of serverMessages.entries()) { + if (message.role === "user" && messageText(message) === liveUserText) { + serverOccurrences += 1 + anchorIndex = index + } + } + if (anchorIndex === -1 || serverOccurrences < liveOccurrences) { + return false + } + + const finalAssistantText = getFinalLiveAssistantText(liveMessages) + if (finalAssistantText === null) { + return true + } + + let serverText = "" + for (const message of serverMessages.slice(anchorIndex + 1)) { + serverText += messageText(message) + } + return serverText.includes(finalAssistantText) +} + +/** + * Decide whether a quiescent live transcript can be replaced by the server. + * + * Count coverage always applies. Content coverage additionally protects the + * final streamed assistant text unless a bounded retry episode is exhausted. + */ +export function decideServerTranscriptAdoption({ + serverMessages, + liveMessages, + allowContentMismatch = false, +}: { + serverMessages: UIMessage[] + liveMessages: UIMessage[] + allowContentMismatch?: boolean +}): ServerTranscriptAdoptionDecision { + if ( + transcriptSignature(liveMessages) === transcriptSignature(serverMessages) + ) { + return "already-current" + } + + // A resolved approval drops its card from the server transcript, so exclude + // approval-card-only messages from the finalize-race length guard. + const liveComparableLength = liveMessages.filter( + (message) => !isApprovalCardMessage(message) + ).length + if (serverMessages.length < liveComparableLength) { + return "reject-count" + } + if ( + !allowContentMismatch && + !serverTranscriptCoversLiveFinalAssistantText(serverMessages, liveMessages) + ) { + return "reject-content" + } + return "adopt" +} + +type TranscriptRetryEpisode = { + completedAttempts: number + key: string + /** `dataUpdatedAt` of the transcript query when the last counted attempt ran. */ + lastDataUpdatedAt: number + timers: Set> +} + /** * Adopt the server transcript wholesale at quiescent boundaries. * @@ -708,42 +886,166 @@ function isApprovalCardMessage(m: UIMessage): boolean { * Message ids differ between the DB serialization and the live stream, so * merging is impossible by design — we replace, never merge. * - * The backend guarantees this is always safe: an approval pause returns 204 on - * resume, DB history already includes the paused partial turn, and any - * continuation stream carries only the suffix. A normal turn can still finish - * before curr_run_id is cleared, though, making the immediate onFinish refetch - * omit the just-finished rows. Keep the longer live transcript until a later - * server snapshot catches up. Never adopt while streaming/submitted either. + * The approval/continuation contract makes wholesale replacement safe once the + * DB snapshot is current: an approval pause returns 204 on resume, DB history + * already includes the paused partial turn, and any continuation stream carries + * only the suffix. A normal turn can still finish before curr_run_id is cleared, + * though, making the immediate onFinish refetch omit the just-finished rows. + * Keep the live transcript until a later server snapshot covers its final + * assistant text and comparable message count. + * Rejected snapshots receive a bounded refetch series; once every retry has + * delivered fresh data, a snapshot that passes the count guard is adopted even + * if its text still differs because the server is canonical at rest. Failed + * refetches do not count toward that bound — an outage must not hand the + * transcript to a stale cached snapshot. Never adopt or refetch from these + * timers while streaming/submitted. */ export function useAdoptServerTranscript({ + chatId, + workspaceId, status, serverMessages, liveMessages, setMessages, }: { + chatId?: string + workspaceId: string status: ChatStatus serverMessages: UIMessage[] liveMessages: UIMessage[] setMessages: (messages: UIMessage[]) => void }) { - const adoptedSignatureRef = useRef(null) + const queryClient = useQueryClient() + const retryEpisodeRef = useRef(null) + const statusRef = useRef(status) + const [retryRevision, setRetryRevision] = useState(0) + statusRef.current = status + + const cancelRetryEpisode = useCallback(function cancelRetryEpisode(): void { + const episode = retryEpisodeRef.current + if (!episode) return + for (const timer of episode.timers) { + clearTimeout(timer) + } + retryEpisodeRef.current = null + }, []) + + const scheduleRetryEpisode = useCallback( + function scheduleRetryEpisode(key: string): void { + cancelRetryEpisode() + if (!chatId) return + + const queryKey = ["chat", chatId, workspaceId, "vercel"] + const episode: TranscriptRetryEpisode = { + completedAttempts: 0, + key, + lastDataUpdatedAt: + queryClient.getQueryState(queryKey)?.dataUpdatedAt ?? 0, + timers: new Set(), + } + retryEpisodeRef.current = episode + + for (const delayMs of ADOPT_SERVER_TRANSCRIPT_RETRY_DELAYS_MS) { + const timer = setTimeout(() => { + episode.timers.delete(timer) + if ( + retryEpisodeRef.current !== episode || + statusRef.current !== "ready" + ) { + return + } + + void queryClient + .invalidateQueries({ queryKey }) + .catch(() => undefined) + .then(() => { + if ( + retryEpisodeRef.current !== episode || + statusRef.current !== "ready" + ) { + return + } + // invalidateQueries resolves even when the refetch fails, so a + // transient outage must not exhaust the content guard and hand + // the transcript to a stale cached snapshot. Count only the + // attempts that actually delivered fresh data — never the timer + // slot index, which would retroactively credit failed slots. + const dataUpdatedAt = + queryClient.getQueryState(queryKey)?.dataUpdatedAt ?? 0 + const delivered = dataUpdatedAt > episode.lastDataUpdatedAt + if (delivered) { + episode.lastDataUpdatedAt = dataUpdatedAt + episode.completedAttempts += 1 + } + const exhausted = + episode.completedAttempts >= + ADOPT_SERVER_TRANSCRIPT_RETRY_DELAYS_MS.length + if (episode.timers.size === 0 && !exhausted) { + // The series ran out of timers without enough delivered + // refetches. Retire the episode so a future server snapshot + // change can start a fresh series — otherwise eventual + // adoption would be permanently unreachable. No revision bump: + // there is nothing to re-decide until new data arrives, so a + // quiet outage does not turn into an endless polling loop. + retryEpisodeRef.current = null + return + } + if (delivered) { + setRetryRevision((revision) => revision + 1) + } + }) + }, delayMs) + episode.timers.add(timer) + } + }, + [cancelRetryEpisode, chatId, queryClient, workspaceId] + ) + + useEffect(() => cancelRetryEpisode, [cancelRetryEpisode]) + useEffect(() => { - if (status !== "ready") return - const serverSignature = transcriptSignature(serverMessages) - if (adoptedSignatureRef.current === serverSignature) return - if (transcriptSignature(liveMessages) === serverSignature) { - adoptedSignatureRef.current = serverSignature + if (status !== "ready") { + cancelRetryEpisode() return } - // A resolved approval drops its card from the server transcript, so exclude - // approval-card-only messages from the finalize-race length guard. - const liveComparableLength = liveMessages.filter( - (m) => !isApprovalCardMessage(m) - ).length - if (serverMessages.length < liveComparableLength) return - adoptedSignatureRef.current = serverSignature + + const liveSignature = transcriptSignature(liveMessages) + const episodeKey = JSON.stringify([chatId, workspaceId, liveSignature]) + const retryEpisode = retryEpisodeRef.current + const allowContentMismatch = + retryEpisode?.key === episodeKey && + retryEpisode.completedAttempts >= + ADOPT_SERVER_TRANSCRIPT_RETRY_DELAYS_MS.length + const decision = decideServerTranscriptAdoption({ + serverMessages, + liveMessages, + allowContentMismatch, + }) + + if (decision === "already-current") { + cancelRetryEpisode() + return + } + if (decision === "reject-content" || decision === "reject-count") { + if (retryEpisode?.key !== episodeKey) { + scheduleRetryEpisode(episodeKey) + } + return + } + + cancelRetryEpisode() setMessages(serverMessages) - }, [status, serverMessages, liveMessages, setMessages]) + }, [ + cancelRetryEpisode, + chatId, + liveMessages, + retryRevision, + scheduleRetryEpisode, + serverMessages, + setMessages, + status, + workspaceId, + ]) } /** diff --git a/frontend/tests/chat-session-pane.test.tsx b/frontend/tests/chat-session-pane.test.tsx index b670d8177..edb038ee7 100644 --- a/frontend/tests/chat-session-pane.test.tsx +++ b/frontend/tests/chat-session-pane.test.tsx @@ -1647,17 +1647,21 @@ describe("ChatSessionPane", () => { }) // Ownership swap at quiescent boundaries: while a turn streams the stream - // owns the transcript; the moment status returns to `ready` the pane adopts - // the server copy WHOLESALE (replace, never merge — DB and stream message ids - // differ by design). The backend guarantees DB history and any continuation - // stream never overlap, so adopting a shorter, longer, or same-length copy is - // uniformly correct. + // owns the transcript; once status returns to `ready` the pane adopts a + // caught-up server copy WHOLESALE (replace, never merge — DB and stream + // message ids differ by design). A finalize-race snapshot can still omit the + // final streamed assistant content, so length alone cannot prove it is safe. describe("adopt-on-ready ownership swap", () => { const userTurn = (id: string, text: string) => ({ id, role: "user" as const, parts: [{ type: "text" as const, text }], }) + const assistantTurn = (id: string, text: string) => ({ + id, + role: "assistant" as const, + parts: [{ type: "text" as const, text }], + }) function mockLiveMessages( messages: UIMessage[], @@ -1736,17 +1740,26 @@ describe("ChatSessionPane", () => { expect(setMessages).not.toHaveBeenCalled() }) - // Same length but different content (ids/text) is a distinct transcript and - // must be adopted too. - it("adopts a same-length server copy with different content", () => { + // DB row segmentation can make a stale snapshot equal-or-longer in count. + // Missing final assistant text still makes it unsafe to adopt immediately. + it("does not adopt a same-length copy missing final assistant content", () => { const setMessages = jest.fn() - mockLiveMessages([userTurn("m1", "hello")], setMessages) + mockLiveMessages( + [ + userTurn("live-user", "Question"), + assistantTurn("live-assistant", "Final streamed answer"), + ], + setMessages + ) - render(renderSubject([userTurn("s1", "server hello")])) + render( + renderSubject([ + userTurn("server-user", "Question"), + assistantTurn("server-assistant", "Previous server answer"), + ]) + ) - expect(setMessages).toHaveBeenCalledTimes(1) - expect(setMessages.mock.calls[0][0]).toHaveLength(1) - expect(setMessages.mock.calls[0][0][0].id).toBe("s1") + expect(setMessages).not.toHaveBeenCalled() }) // The stream owns the current turn: while status is streaming the pane must