From f59ca3cb0d548799b53bbc35c5d072b73b462307 Mon Sep 17 00:00:00 2001 From: claude code agent 227 Date: Tue, 23 Jun 2026 02:25:52 +0300 Subject: [PATCH] refactor(ai-chat): extract useChatSession hook + lock the id lifecycle with tests MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Addresses the 2nd PR #138 review (test debt + the Variant-B architecture ask). The new→persisted chat id lifecycle (mount key, both adoption paths, the history-load latch, the render-phase reconciler, onTurnFinished) is moved out of the 768-line window into a new useChatSession hook driven by a pure threadSessionReducer (reconcile/adopt), so adopt-vs-switch is one explicit dispatch point and the scattering the review flagged is gone (window: 768→~620). Tests (the blockers): - use-chat-session.test.tsx — hook-level locks incl. the #137 regression (adopts the authoritative streamed id 'A', NOT chats.items[0]='B' — fails on the old heuristic), the error-path fallback (arm/adopt/ambiguous/add+delete), the disarm-on-reconcile lock (a fallback armed then switched away must not be adopted by a late refetch), in-place-adopt-keeps-key vs external-switch-remount, and the waitingForHistory latch. - extractServerChatId (reading message.metadata.chatId) and newlyAddedChatIds extracted as pure helpers with unit tests; threadSessionReducer tested. Cleanups: single canonical #137 explanation in adopt-chat-id.ts (other sites reference it); fallback effect computes the set diff once; invalidate callbacks memoized; redundant invariant tests folded. Behavior preserved — re-verified live (z.ai glm-5.2): new-chat adopt + 2nd turn in the same row, no mid-conversation remount, two-tab race leak-free, switch to an existing chat reseeds full history, reload restores history. Co-Authored-By: Claude Opus 4.8 --- .../ai-chat/components/ai-chat-window.tsx | 196 +++--------------- .../ai-chat/components/chat-thread.tsx | 22 +- .../ai-chat/hooks/use-chat-session.test.tsx | 149 +++++++++++++ .../ai-chat/hooks/use-chat-session.ts | 196 ++++++++++++++++++ .../ai-chat/utils/adopt-chat-id.test.ts | 58 +++++- .../features/ai-chat/utils/adopt-chat-id.ts | 57 ++++- .../ai-chat/utils/thread-identity.test.ts | 56 +++-- .../features/ai-chat/utils/thread-identity.ts | 26 +++ .../src/core/ai-chat/ai-chat.service.ts | 9 +- 9 files changed, 556 insertions(+), 213 deletions(-) create mode 100644 apps/client/src/features/ai-chat/hooks/use-chat-session.test.tsx create mode 100644 apps/client/src/features/ai-chat/hooks/use-chat-session.ts diff --git a/apps/client/src/features/ai-chat/components/ai-chat-window.tsx b/apps/client/src/features/ai-chat/components/ai-chat-window.tsx index a2418110..83200747 100644 --- a/apps/client/src/features/ai-chat/components/ai-chat-window.tsx +++ b/apps/client/src/features/ai-chat/components/ai-chat-window.tsx @@ -6,7 +6,6 @@ import { useRef, useState, } from "react"; -import { generateId } from "ai"; import { type UIMessage } from "@ai-sdk/react"; import { Group, Loader, Tooltip } from "@mantine/core"; import { @@ -42,16 +41,7 @@ import { import ConversationList from "@/features/ai-chat/components/conversation-list.tsx"; import ChatThread from "@/features/ai-chat/components/chat-thread.tsx"; import { buildChatMarkdown } from "@/features/ai-chat/utils/chat-markdown.ts"; -import { - resolveAdoptedChatId, - pickNewlyCreatedChatId, -} from "@/features/ai-chat/utils/adopt-chat-id.ts"; -import { - type ThreadIdentity, - newThread, - switchThread, - adoptThread, -} from "@/features/ai-chat/utils/thread-identity.ts"; +import { useChatSession } from "@/features/ai-chat/hooks/use-chat-session.ts"; import { shouldCollapseOnOutsidePointer, isHeaderClick, @@ -143,37 +133,6 @@ export default function AiChatWindow() { // left partly off-screen). const [geom, setGeom] = useAtom(aiChatWindowGeomAtom); - // Latch: the chat id whose full persisted history has finished loading while - // its thread is mounted. Used so a later BACKGROUND refetch (the post-turn - // messages invalidation) never tears the live thread back down to the loader. - const historyLoadedKeyRef = useRef(null); - - // Error-path fallback for new-chat id adoption. When a brand-new chat's first - // turn errors BEFORE the server's `start` chunk, no authoritative chatId ever - // reaches the client, so the primary metadata adoption cannot run. We then - // ARM this ref with a snapshot of the currently-known chat ids; once the list - // refetch lands with the just-created row, the fallback effect below adopts - // the SINGLE newly-appeared id (unambiguous — unlike the old items[0] guess). - // `null` = not armed. - const pendingNewChatRef = useRef(null); - - // The mounted thread's identity: ONE atomic value tying ChatThread's mount - // key (`thread.key`) to the chat id that mounted thread holds (`thread.chatId`). - // Consolidating these makes the "key vs chat id diverged" state unrepresentable - // — every change goes through an explicit transition (see thread-identity.ts): - // • switchThread — a genuine switch (key := chat id) → remount + reseed; - // • adoptThread — a brand-new chat learns its real id (key UNCHANGED, chat - // id null→realId) → NO remount, so the live useChat store (holding the - // just-finished turn) is preserved and the 2nd turn sends the real chatId. - // A non-null initial activeChatId yields {key: id, chatId: id}; the `""` newKey - // is unused on that (non-null) path. A null initial chat gets a fresh session - // key with no chat id yet. - const [thread, setThread] = useState(() => - activeChatId === null - ? newThread(`new-${generateId()}`) - : switchThread(activeChatId), - ); - const { data: chats } = useAiChatsQuery(); // Roles for the new-chat picker (any member may list them). Only fetched while // the window is open. @@ -214,14 +173,14 @@ export default function AiChatWindow() { ? { id: openPageData.id, title: openPageData.title } : null; + // startNewChat/selectChat only set the public atom; useChatSession's render- + // phase reconciler does the remount AND disarms any pending new-chat fallback. const startNewChat = useCallback((): void => { setActiveChatId(null); setHistoryOpen(false); setDraft(""); // Default the picker back to "Universal assistant" for the fresh chat. setSelectedRoleId(null); - // The user moved on — disarm any pending error-path adoption fallback. - pendingNewChatRef.current = null; }, [setActiveChatId, setDraft, setSelectedRoleId]); const selectChat = useCallback( @@ -232,83 +191,33 @@ export default function AiChatWindow() { // Reset the card-picked role so a stale pick can't leak into the existing // chat's header/assistant-name (which prefers the chat's persisted role). setSelectedRoleId(null); - // The user moved on — disarm any pending error-path adoption fallback. - pendingNewChatRef.current = null; }, [setActiveChatId, setDraft, setSelectedRoleId], ); - // After a turn finishes, refresh the chat list. For a brand-new chat (no id - // yet), the server has just created the row and reported its AUTHORITATIVE id - // via the streamed assistant message metadata (`serverChatId`); we adopt THAT - // exact id so the thread switches from "new" to its OWN persisted chat — never - // guessing the newest chat in the list, which would adopt a SIBLING chat that - // a second tab created in the same moment and leak this tab's later turns into - // it (the two-tab adoption race, #137). - const onTurnFinished = useCallback( - (serverChatId?: string) => { - const adopted = resolveAdoptedChatId(activeChatId, serverChatId); - if (adopted) { - // PRIMARY path. In-place adoption: set the public selection and the - // thread identity to the real id together. `adoptThread` keeps the SAME - // mount key (chat id null→realId), so the render-phase reconciler below - // sees `activeChatId === thread.chatId` and keeps the SAME mounted thread - // (its useChat already holds the just-finished turn, and its chatIdRef - // now mirrors the real id so the SECOND turn sends the correct chatId) - // instead of remounting and re-seeding from not-yet-persisted history. - setActiveChatId(adopted); - setThread((prev) => adoptThread(prev, adopted)); - // Primary adoption won — disarm any previously-armed fallback. - pendingNewChatRef.current = null; - } else if (activeChatId === null) { - // FALLBACK path: a brand-new chat finished with NO server id (the first - // turn errored before the `start` chunk). Arm the bounded list-refetch - // fallback by snapshotting the currently-known chat ids. `chats` is still - // the pre-refetch list here, so the just-created row is NOT yet in it; the - // effect below adopts the single id that newly appears after the refetch. - pendingNewChatRef.current = chats?.items?.map((c) => c.id) ?? []; - } - queryClient.invalidateQueries({ queryKey: AI_CHATS_RQ_KEY }); - // Re-sync the persisted message rows for the active chat so the Markdown - // export and the token counters reflect the turn that just finished. The - // live thread renders from its own useChat store (stable thread.key / store - // id), so refetching these rows never re-seeds or tears down the open - // thread. For a brand-new chat activeChatId is still null here; that chat's - // first row load happens right after id adoption, and every later turn hits - // this invalidation with the adopted id. - if (activeChatId) { - queryClient.invalidateQueries({ - queryKey: AI_CHAT_MESSAGES_RQ_KEY(activeChatId), - }); - } - }, - [activeChatId, chats, queryClient, setActiveChatId], + // The AI-chat thread-identity lifecycle (mount key, both new-chat id adoption + // paths, the history-loaded latch, the render-phase reconciler) lives in this + // hook. See adopt-chat-id.ts for the canonical #137 two-tab race explanation. + // Memoized so the hook's `onTurnFinished` useCallback keeps a stable identity + // across renders (it is read live by useChat's onFinish, so this is cosmetic, + // not load-bearing — but avoids needless re-creation). + const invalidateChatList = useCallback( + () => queryClient.invalidateQueries({ queryKey: AI_CHATS_RQ_KEY }), + [queryClient], ); - - // FALLBACK resolver. Armed only by onTurnFinished when a brand-new chat's - // first turn errored before the `start` chunk (no authoritative id streamed). - // Once the per-user list refetch lands with the just-created row, adopt the - // SINGLE id that newly appeared relative to the pre-refetch snapshot. Adoption - // is IN PLACE (set activeChatId + adoptThread the identity together) exactly - // like the primary path, so the render-phase reconciler does not remount. - useEffect(() => { - const before = pendingNewChatRef.current; - if (before === null || activeChatId !== null) return; // not armed / already adopted - const after = chats?.items?.map((c) => c.id) ?? []; - // Keep waiting until the list membership actually CHANGES (the refetch - // landed). A length compare would miss the change if another session - // concurrently deleted a chat in this same window (length stays equal while - // a new row was added); a set-membership compare is robust to that. - const beforeSet = new Set(before); - if (after.length === before.length && after.every((id) => beforeSet.has(id))) - return; // list not refetched yet — keep waiting - const adopted = pickNewlyCreatedChatId(before, after); // single new id, or null if ambiguous - pendingNewChatRef.current = null; // give up after one resolved refetch either way - if (adopted) { - setActiveChatId(adopted); - setThread((prev) => adoptThread(prev, adopted)); - } - }, [chats, activeChatId, setActiveChatId]); + const invalidateChatMessages = useCallback( + (id: string) => + queryClient.invalidateQueries({ queryKey: AI_CHAT_MESSAGES_RQ_KEY(id) }), + [queryClient], + ); + const { threadKey, waitingForHistory, onTurnFinished } = useChatSession({ + activeChatId, + setActiveChatId, + chats, + messagesLoading, + onInvalidateChatList: invalidateChatList, + onInvalidateChatMessages: invalidateChatMessages, + }); // The active chat object (for its title) and an export gate: only enable the // export button when an existing chat with loaded persisted rows is active. @@ -362,59 +271,6 @@ export default function AiChatWindow() { notifications.show({ message: t("Copied") }); }, [activeChatId, messageRows, activeChat, clipboard, t]); - // NOTE: new-chat id adoption has two paths, neither of which is the old - // "newest chat in the list" heuristic that raced a second tab (#137). - // PRIMARY: the server reports the authoritative created chat id on the streamed - // assistant message metadata, and `onTurnFinished(serverChatId)` adopts THAT id - // in place — see above. FALLBACK (only when the first turn errors before the - // `start` chunk, so no metadata id is streamed): adopt the SINGLE chat that - // newly appeared in the refetched per-user list relative to a pre-refetch - // snapshot taken when the failed turn finished — unambiguous, and it does not - // race a sibling chat the way `items[0]` did (see the fallback effect above). - - // Reconcile the thread identity against the active-chat atom during render when - // they diverge — the React-sanctioned alternative to an effect (it re-renders - // before paint, no extra commit, and converges since the next render finds them - // equal). This reconciliation MUST remain: `activeChatId` is the public - // selection and is ALSO set from OUTSIDE this component (e.g. page-history's - // history-item calls setActiveChatId to open a referenced chat). A divergence - // here is a genuine SWITCH (external atom change OR user switch via - // selectChat/startNewChat), so `switchThread` remounts + reseeds. In-place - // adoption never reaches this branch: it set activeChatId and the thread's - // chatId to the same value, so `activeChatId === thread.chatId` here. - if (activeChatId !== thread.chatId) { - setThread( - activeChatId === null - ? newThread(`new-${generateId()}`) - : switchThread(activeChatId), - ); - } - // Latch the active chat once its full history has loaded and its thread is - // mounted, so a later background refetch (the post-turn messages - // invalidation, which can transiently flip hasNextPage for a chat whose - // message count is an exact multiple of the server page size) does not tear - // the live thread down to a loader and lose its in-progress useChat state. - if ( - activeChatId !== null && - thread.key === activeChatId && - !messagesLoading && - historyLoadedKeyRef.current !== activeChatId - ) { - historyLoadedKeyRef.current = activeChatId; - } - - // Show the history loader only when freshly OPENING an existing chat (the key - // equals the chat id) whose history has not been fully loaded yet. For a live - // in-place thread that adopted its id, the key is still the "new-…" session - // key, so we keep showing the live thread instead of unmounting it behind a - // loader; and once a chat's history has loaded, a later background refetch no - // longer tears the thread back down (see the latch above). - const waitingForHistory = - activeChatId !== null && - messagesLoading && - thread.key === activeChatId && - historyLoadedKeyRef.current !== activeChatId; - // Current context size for the active chat: how much the conversation now // occupies in the model's context window — NOT the cumulative tokens spent. // We read the most recent assistant row that carries a context figure: @@ -739,7 +595,7 @@ export default function AiChatWindow() { ) : ( the generic "AI agent". */ assistantName?: string; - /** Called when a turn finishes; the parent refreshes the chat list and, for - * a new chat, adopts the freshly created chat id. `serverChatId` is the - * authoritative id the server attached to the streamed assistant message - * metadata (see the server's `messageMetadata`); the parent adopts THIS for a - * new chat instead of guessing the newest chat in the list (fixes the two-tab - * adoption race, #137). Undefined on a failed turn that produced no metadata. */ + /** Called when a turn finishes; the parent refreshes the chat list and, for a + * new chat, adopts the freshly created chat id. `serverChatId` is the + * authoritative id the server streamed on the assistant message metadata, or + * undefined on a failed turn — see adopt-chat-id.ts for the full #137 design. */ onTurnFinished: (serverChatId?: string) => void; /** Parent-owned ref that this thread keeps updated with its live useChat * snapshot (full message list + streaming flag), so the header's @@ -251,13 +250,10 @@ export default function ChatThread({ // would be wrong, so on Stop/disconnect/error the queue is left intact for // the user to decide. onFinish: ({ message, isAbort, isDisconnect, isError }) => { - // The server attaches the authoritative chatId to the streamed assistant - // message metadata (see `messageMetadata` in ai-chat.service.ts). Forward it - // so the parent adopts the REAL created chat id for a new chat, rather than - // guessing the newest chat in the list (which races a second tab — #137). - const serverChatId = (message?.metadata as { chatId?: string } | undefined) - ?.chatId; - onTurnFinished(serverChatId); + // Forward the authoritative server chatId (streamed on the assistant + // message metadata) so the parent adopts the REAL created chat id for a new + // chat — see adopt-chat-id.ts for the full #137 design. + onTurnFinished(extractServerChatId(message)); // Show a neutral "stopped" marker for an aborted turn; the red error banner // (via `error`) already covers isError, and a clean finish clears any marker. if (isError) setStopNotice(null); diff --git a/apps/client/src/features/ai-chat/hooks/use-chat-session.test.tsx b/apps/client/src/features/ai-chat/hooks/use-chat-session.test.tsx new file mode 100644 index 00000000..e7b8695c --- /dev/null +++ b/apps/client/src/features/ai-chat/hooks/use-chat-session.test.tsx @@ -0,0 +1,149 @@ +import { describe, it, expect, vi, beforeEach } from "vitest"; +import { renderHook } from "@testing-library/react"; +import { useChatSession } from "./use-chat-session"; + +// Drive the hook the way the window does: the parent owns `activeChatId` and +// passes it back in. `setActiveChatId` is a spy so we can assert the EXACT id the +// hook adopts (the #137 regression: it must be the authoritative streamed id, not +// the newest chat in the list). +function setup(initial: { + activeChatId: string | null; + chats: { items?: { id: string }[] } | undefined; + messagesLoading?: boolean; +}) { + const setActiveChatId = vi.fn(); + const onInvalidateChatList = vi.fn(); + const onInvalidateChatMessages = vi.fn(); + const { result, rerender } = renderHook( + (props: { + activeChatId: string | null; + chats: { items?: { id: string }[] } | undefined; + messagesLoading?: boolean; + }) => + useChatSession({ + activeChatId: props.activeChatId, + setActiveChatId, + chats: props.chats, + messagesLoading: props.messagesLoading ?? false, + onInvalidateChatList, + onInvalidateChatMessages, + }), + { initialProps: initial }, + ); + return { + result, + rerender, + setActiveChatId, + onInvalidateChatList, + onInvalidateChatMessages, + }; +} + +describe("useChatSession", () => { + beforeEach(() => vi.clearAllMocks()); + + it("#137 REGRESSION LOCK: adopts the authoritative streamed id, NOT items[0]", () => { + // Brand-new chat, list already holds a SIBLING chat B as items[0] (a second + // tab just created it). The server streams the real id "A" for THIS chat. + const { result, setActiveChatId } = setup({ + activeChatId: null, + chats: { items: [{ id: "B" }] }, + }); + result.current.onTurnFinished("A"); + // Must adopt the authoritative id, not the newest-in-list guess. + expect(setActiveChatId).toHaveBeenCalledWith("A"); + expect(setActiveChatId).not.toHaveBeenCalledWith("B"); + }); + + it("fallback adopt: arms on a server-id-less finish, adopts the single new id after refetch", () => { + const { result, rerender, setActiveChatId } = setup({ + activeChatId: null, + chats: { items: [{ id: "x" }] }, + }); + // No server id => arm the fallback (no adoption yet). + result.current.onTurnFinished(undefined); + expect(setActiveChatId).not.toHaveBeenCalled(); + // The refetch lands with the new row => adopt it. + rerender({ activeChatId: null, chats: { items: [{ id: "x" }, { id: "new" }] } }); + expect(setActiveChatId).toHaveBeenCalledWith("new"); + }); + + it("fallback ambiguous: two new ids appear => no adoption", () => { + const { result, rerender, setActiveChatId } = setup({ + activeChatId: null, + chats: { items: [{ id: "x" }] }, + }); + result.current.onTurnFinished(undefined); + rerender({ + activeChatId: null, + chats: { items: [{ id: "x" }, { id: "n1" }, { id: "n2" }] }, + }); + expect(setActiveChatId).not.toHaveBeenCalled(); + }); + + it("fallback add+delete in one window: adopts the new id (membership compare)", () => { + const { result, rerender, setActiveChatId } = setup({ + activeChatId: null, + chats: { items: [{ id: "a" }, { id: "b" }] }, + }); + result.current.onTurnFinished(undefined); + // a was deleted, new was added — same length, but membership changed. + rerender({ activeChatId: null, chats: { items: [{ id: "b" }, { id: "new" }] } }); + expect(setActiveChatId).toHaveBeenCalledWith("new"); + }); + + it("disarm on reconcile: a fallback armed then switched away is NOT adopted by a late refetch", () => { + // Arm the error-path fallback on a brand-new chat (snapshot before=["x"]). + const { result, rerender, setActiveChatId } = setup({ + activeChatId: null, + chats: { items: [{ id: "x" }] }, + }); + result.current.onTurnFinished(undefined); + // The user switches to an existing chat C BEFORE the refetch lands; the + // render-phase reconciler must DISARM the pending fallback. + rerender({ activeChatId: "C", chats: { items: [{ id: "x" }] } }); + // ...then starts a fresh new chat again (back to null), without re-arming. + rerender({ activeChatId: null, chats: { items: [{ id: "x" }] } }); + // A late refetch now brings a new row. Because the earlier fallback was + // disarmed on the switch (not left armed with the stale ["x"] snapshot), it + // must NOT be adopted. (Without the disarm this would wrongly adopt "new".) + rerender({ + activeChatId: null, + chats: { items: [{ id: "x" }, { id: "new" }] }, + }); + expect(setActiveChatId).not.toHaveBeenCalledWith("new"); + }); + + it("in-place adopt keeps threadKey stable; an external switch remounts", () => { + const chats = { items: [{ id: "B" }] }; + const { result, rerender } = setup({ activeChatId: null, chats }); + const keyBefore = result.current.threadKey; + // Adopt the streamed id; the PARENT then reflects activeChatId="A" back in. + result.current.onTurnFinished("A"); + rerender({ activeChatId: "A", chats }); + // In-place adoption: SAME mount key (the live useChat store is preserved). + expect(result.current.threadKey).toBe(keyBefore); + + // An EXTERNAL switch (not via adopt) to a different chat must remount: the + // key becomes the chat id. + rerender({ activeChatId: "C", chats }); + expect(result.current.threadKey).toBe("C"); + }); + + it("waitingForHistory gates the loader only while opening an unloaded existing chat", () => { + // Open an existing chat whose history is still loading => loader on. + const { result, rerender } = setup({ + activeChatId: "chat-1", + chats: { items: [{ id: "chat-1" }] }, + messagesLoading: true, + }); + expect(result.current.waitingForHistory).toBe(true); + // Once loading finishes, the latch flips and the loader is off. + rerender({ + activeChatId: "chat-1", + chats: { items: [{ id: "chat-1" }] }, + messagesLoading: false, + }); + expect(result.current.waitingForHistory).toBe(false); + }); +}); diff --git a/apps/client/src/features/ai-chat/hooks/use-chat-session.ts b/apps/client/src/features/ai-chat/hooks/use-chat-session.ts new file mode 100644 index 00000000..65392f62 --- /dev/null +++ b/apps/client/src/features/ai-chat/hooks/use-chat-session.ts @@ -0,0 +1,196 @@ +import { useCallback, useEffect, useReducer, useRef } from "react"; +import { generateId } from "ai"; +import { + resolveAdoptedChatId, + newlyAddedChatIds, +} from "@/features/ai-chat/utils/adopt-chat-id.ts"; +import { + newThread, + switchThread, + threadSessionReducer, +} from "@/features/ai-chat/utils/thread-identity.ts"; + +/** What the window needs from a chat session: the ChatThread mount key, the + * history-loader gate, and the turn-finished callback. */ +export interface ChatSession { + /** ChatThread mount key (was `thread.key`). */ + threadKey: string; + /** Show the history loader instead of the live thread. */ + waitingForHistory: boolean; + /** Call when a turn finishes; `serverChatId` is the authoritative streamed id + * (undefined on a failed turn). Handles new-chat id adoption + invalidations. */ + onTurnFinished: (serverChatId?: string) => void; +} + +/** + * Owns the AI-chat thread-identity lifecycle: the single atomic thread identity, + * both new-chat id adoption paths (primary streamed-metadata + bounded error-path + * fallback), the history-loaded latch, and the render-phase reconciler that keeps + * the thread's mount key in sync with the public `activeChatId` atom. + * + * This is the twice-bugged area for the #137 two-tab adoption race; the canonical + * explanation of the adoption design lives in adopt-chat-id.ts. + */ +export function useChatSession(params: { + activeChatId: string | null; + setActiveChatId: (id: string | null) => void; + chats: { items?: { id: string }[] } | undefined; + messagesLoading: boolean; + /** Wraps queryClient.invalidateQueries(AI_CHATS_RQ_KEY). */ + onInvalidateChatList: () => void; + /** Wraps the per-chat messages invalidation. */ + onInvalidateChatMessages: (chatId: string) => void; +}): ChatSession { + const { + activeChatId, + setActiveChatId, + chats, + messagesLoading, + onInvalidateChatList, + onInvalidateChatMessages, + } = params; + + // The mounted thread's identity: ONE atomic value tying ChatThread's mount key + // (`thread.key`) to the chat id that mounted thread holds (`thread.chatId`). + // Consolidating these makes the "key vs chat id diverged" state unrepresentable + // — every change goes through an explicit transition (see thread-identity.ts): + // `newThread`/`switchThread` to (re)mount, `adoptThread` for in-place adoption. + // Initial: a non-null activeChatId switches to it; a null one gets a fresh + // session key with no chat id yet. + const [thread, dispatch] = useReducer( + threadSessionReducer, + undefined, + () => + activeChatId === null + ? newThread(`new-${generateId()}`) + : switchThread(activeChatId), + ); + + // Error-path fallback for new-chat id adoption. When a brand-new chat's first + // turn errors BEFORE the server's `start` chunk, no authoritative chatId ever + // reaches the client, so the primary metadata adoption cannot run. We then ARM + // this ref with a snapshot of the currently-known chat ids; once the list + // refetch lands with the just-created row, the fallback effect below adopts the + // SINGLE newly-appeared id. `null` = not armed. See adopt-chat-id.ts (#137). + const pendingNewChatRef = useRef(null); + + // Latch: the chat id whose full persisted history has finished loading while + // its thread is mounted. Used so a later BACKGROUND refetch (the post-turn + // messages invalidation) never tears the live thread back down to the loader. + const historyLoadedKeyRef = useRef(null); + + // After a turn finishes, refresh the chat list. For a brand-new chat (no id + // yet) we adopt the server's AUTHORITATIVE streamed id (never the newest in the + // list, which races a second tab — #137; see adopt-chat-id.ts). + const onTurnFinished = useCallback( + (serverChatId?: string) => { + const adopted = resolveAdoptedChatId(activeChatId, serverChatId); + if (adopted) { + // PRIMARY path. In-place adoption: set the public selection and the + // thread identity to the real id together. `adopt` keeps the SAME mount + // key, so the render-phase reconciler sees `activeChatId === thread.chatId` + // and keeps the SAME mounted thread (its useChat already holds the + // just-finished turn) instead of remounting + re-seeding from + // not-yet-persisted history. + setActiveChatId(adopted); + dispatch({ type: "adopt", chatId: adopted }); + // Primary adoption won — disarm any previously-armed fallback. + pendingNewChatRef.current = null; + } else if (activeChatId === null) { + // FALLBACK path: a brand-new chat finished with NO server id (the first + // turn errored before the `start` chunk). Arm the bounded list-refetch + // fallback by snapshotting the currently-known chat ids. `chats` is still + // the pre-refetch list here, so the just-created row is NOT yet in it; the + // effect below adopts the single id that newly appears after the refetch. + pendingNewChatRef.current = chats?.items?.map((c) => c.id) ?? []; + } + onInvalidateChatList(); + // Re-sync the persisted message rows for the active chat so the Markdown + // export and token counters reflect the just-finished turn. The live thread + // renders from its own useChat store (stable thread.key), so this never + // re-seeds or tears down the open thread. For a brand-new chat activeChatId + // is still null here; later turns hit this with the adopted id. + if (activeChatId) { + onInvalidateChatMessages(activeChatId); + } + }, + [ + activeChatId, + chats, + setActiveChatId, + onInvalidateChatList, + onInvalidateChatMessages, + ], + ); + + // FALLBACK resolver. Armed only by onTurnFinished when a brand-new chat's first + // turn errored before the `start` chunk (no authoritative id streamed). Once + // the per-user list refetch lands with the just-created row, adopt the SINGLE + // id that newly appeared relative to the pre-refetch snapshot. Adoption is IN + // PLACE (set activeChatId + `adopt` together) like the primary path, so the + // render-phase reconciler does not remount. + useEffect(() => { + const before = pendingNewChatRef.current; + if (before === null || activeChatId !== null) return; // not armed / already adopted + const after = chats?.items?.map((c) => c.id) ?? []; + const added = newlyAddedChatIds(before, after); + // Keep waiting until a genuinely-new id appears. Set-based, so it is robust + // to an add+delete in the same window (a length compare would miss it), and + // it deliberately keeps waiting through an unrelated deletion (no new id yet) + // until the just-created row actually lands, rather than giving up early. + if (added.size === 0) return; // list not refetched yet — keep waiting + pendingNewChatRef.current = null; // resolved — disarm + if (added.size === 1) { + // single unambiguous new id; >1 = ambiguous → give up + const adopted = [...added][0]; + setActiveChatId(adopted); + dispatch({ type: "adopt", chatId: adopted }); + } + }, [chats, activeChatId, setActiveChatId]); + + // Reconcile the thread identity against the active-chat atom during render when + // they diverge — the React-sanctioned alternative to an effect (re-renders + // before paint, no extra commit, and converges since the next render finds them + // equal). This reconciliation MUST remain: `activeChatId` is the public + // selection and is ALSO set from OUTSIDE this component (e.g. page-history opens + // a referenced chat via setActiveChatId). A divergence here is a genuine SWITCH + // (external atom change OR user switch via selectChat/startNewChat), so + // `reconcile` remounts + reseeds. In-place adoption never reaches this branch: + // it set activeChatId and thread.chatId to the same value. + if (activeChatId !== thread.chatId) { + // A genuine switch makes any pending error-path new-chat fallback moot. + pendingNewChatRef.current = null; + dispatch({ + type: "reconcile", + chatId: activeChatId, + newKey: `new-${generateId()}`, + }); + } + + // Latch the active chat once its full history has loaded and its thread is + // mounted, so a later background refetch (the post-turn messages invalidation, + // which can transiently flip hasNextPage for a chat whose message count is an + // exact multiple of the server page size) does not tear the live thread down to + // a loader and lose its in-progress useChat state. + if ( + activeChatId !== null && + thread.key === activeChatId && + !messagesLoading && + historyLoadedKeyRef.current !== activeChatId + ) { + historyLoadedKeyRef.current = activeChatId; + } + + // Show the history loader only when freshly OPENING an existing chat (the key + // equals the chat id) whose history has not been fully loaded yet. For a live + // in-place thread that adopted its id, the key is still the "new-…" session + // key, so the live thread keeps rendering; and once a chat's history has loaded, + // a later background refetch no longer tears it down (see the latch above). + const waitingForHistory = + activeChatId !== null && + messagesLoading && + thread.key === activeChatId && + historyLoadedKeyRef.current !== activeChatId; + + return { threadKey: thread.key, waitingForHistory, onTurnFinished }; +} diff --git a/apps/client/src/features/ai-chat/utils/adopt-chat-id.test.ts b/apps/client/src/features/ai-chat/utils/adopt-chat-id.test.ts index d5ba12ec..39a40984 100644 --- a/apps/client/src/features/ai-chat/utils/adopt-chat-id.test.ts +++ b/apps/client/src/features/ai-chat/utils/adopt-chat-id.test.ts @@ -1,5 +1,10 @@ import { describe, it, expect } from "vitest"; -import { resolveAdoptedChatId, pickNewlyCreatedChatId } from "./adopt-chat-id"; +import { + resolveAdoptedChatId, + pickNewlyCreatedChatId, + newlyAddedChatIds, + extractServerChatId, +} from "./adopt-chat-id"; describe("resolveAdoptedChatId", () => { it("adopts the server id for a brand-new chat (activeChatId null + id)", () => { @@ -41,3 +46,54 @@ describe("pickNewlyCreatedChatId", () => { expect(pickNewlyCreatedChatId(["a", "b"], ["b", "a"])).toBeNull(); }); }); + +describe("newlyAddedChatIds", () => { + it("returns the single new id", () => { + expect([...newlyAddedChatIds(["a", "b"], ["a", "b", "c"])]).toEqual(["c"]); + }); + + it("returns an empty set when nothing was added", () => { + expect(newlyAddedChatIds(["a", "b"], ["b", "a"]).size).toBe(0); + }); + + it("returns both new ids when two were added", () => { + expect(newlyAddedChatIds(["a"], ["a", "b", "c"])).toEqual( + new Set(["b", "c"]), + ); + }); + + it("keeps only the new id across an add+delete in the same window", () => { + // before [a,b] -> after [b,new]: a was deleted, new was added. + expect([...newlyAddedChatIds(["a", "b"], ["b", "new"])]).toEqual(["new"]); + }); + + it("dedupes a repeated new id to a single entry", () => { + expect(newlyAddedChatIds(["a"], ["a", "new", "new"])).toEqual( + new Set(["new"]), + ); + }); +}); + +describe("extractServerChatId", () => { + it("returns the chatId when present on metadata", () => { + expect(extractServerChatId({ metadata: { chatId: "chat-1" } })).toBe( + "chat-1", + ); + }); + + it("returns undefined when the message has no metadata", () => { + expect(extractServerChatId({})).toBeUndefined(); + }); + + it("returns undefined when metadata lacks chatId", () => { + expect(extractServerChatId({ metadata: { other: 1 } })).toBeUndefined(); + }); + + it("returns undefined for a non-string chatId", () => { + expect(extractServerChatId({ metadata: { chatId: 42 } })).toBeUndefined(); + }); + + it("returns undefined for an undefined message", () => { + expect(extractServerChatId(undefined)).toBeUndefined(); + }); +}); diff --git a/apps/client/src/features/ai-chat/utils/adopt-chat-id.ts b/apps/client/src/features/ai-chat/utils/adopt-chat-id.ts index 16bebc04..57d2388b 100644 --- a/apps/client/src/features/ai-chat/utils/adopt-chat-id.ts +++ b/apps/client/src/features/ai-chat/utils/adopt-chat-id.ts @@ -1,15 +1,33 @@ /** * Pure helpers for adopting a brand-new chat's authoritative server id. * + * ============================ CANONICAL #137 NOTE ============================ + * This docblock is the single authoritative explanation of the new-chat id + * adoption design and the #137 two-tab race it fixes. Other call sites + * (use-chat-session.ts, the server's `chatStreamStartMetadata`) reference here + * rather than restating it. + * + * When a user sends the first turn of a BRAND-NEW chat, the client has no chat + * id yet (`activeChatId === null`). The server creates the row and the client + * must "adopt" that row's real id so the SECOND turn targets the same chat. + * + * The OLD heuristic adopted `items[0]` — the newest chat in the refetched list. + * That races a second tab: if another tab created a chat in the same moment, + * its row could be `items[0]`, so this tab would adopt the SIBLING chat and + * leak its later turns into it (#137). We adopt by IDENTITY instead, two ways: + * * PRIMARY path: the server streams the real chat id on the assistant message - * metadata (see `chatStreamStartMetadata` server-side); `resolveAdoptedChatId` - * turns that into the id to adopt for a new chat. + * metadata's `start` part (see `chatStreamStartMetadata` server-side); + * `extractServerChatId` reads it off the finished message and + * `resolveAdoptedChatId` turns it into the id to adopt for a new chat. This is + * authoritative and immune to the race. * * FALLBACK path (only when a new chat's first turn errors BEFORE the `start` * chunk, so no metadata id ever reached the client): adopt the single chat that * NEWLY appeared in the per-user list relative to a pre-refetch snapshot — - * `pickNewlyCreatedChatId`. This is unambiguous and does not race a second tab - * the way the old "newest chat in the list" guess did (#137). + * `newlyAddedChatIds` / `pickNewlyCreatedChatId`. This is unambiguous and does + * not race a second tab the way the old "newest chat in the list" guess did. + * ============================================================================ */ /** @@ -24,6 +42,32 @@ export function resolveAdoptedChatId( return activeChatId === null && serverChatId ? serverChatId : null; } +/** + * Read the authoritative server chat id off a finished assistant message. The + * server attaches it as `message.metadata.chatId` on the `start` part (see + * `chatStreamStartMetadata`). Returns it only when it is a string; undefined for + * a missing message, missing metadata, or a non-string `chatId`. + */ +export function extractServerChatId( + message: { metadata?: unknown } | undefined, +): string | undefined { + const m = message?.metadata as { chatId?: string } | undefined; + return typeof m?.chatId === "string" ? m.chatId : undefined; +} + +/** + * The deduped set of ids present in `afterIds` but not in `beforeIds`. A + * paginated/flatMapped list can repeat the same id, so dedupe: one genuinely-new + * chat must not read as multiple from a duplicate. + */ +export function newlyAddedChatIds( + beforeIds: readonly string[], + afterIds: readonly string[], +): Set { + const before = new Set(beforeIds); + return new Set(afterIds.filter((id) => !before.has(id))); +} + /** * Return the single id present in `afterIds` but not in `beforeIds`. Returns * null when zero or more-than-one such id exists (ambiguous — do not adopt). @@ -32,9 +76,6 @@ export function pickNewlyCreatedChatId( beforeIds: readonly string[], afterIds: readonly string[], ): string | null { - const before = new Set(beforeIds); - // Dedupe the new ids: a paginated/flatMapped list can repeat the same id, and - // one genuinely-new chat must not read as "ambiguous" (>1) from a duplicate. - const added = new Set(afterIds.filter((id) => !before.has(id))); + const added = newlyAddedChatIds(beforeIds, afterIds); return added.size === 1 ? [...added][0] : null; } diff --git a/apps/client/src/features/ai-chat/utils/thread-identity.test.ts b/apps/client/src/features/ai-chat/utils/thread-identity.test.ts index 3ab65a11..eab1eadb 100644 --- a/apps/client/src/features/ai-chat/utils/thread-identity.test.ts +++ b/apps/client/src/features/ai-chat/utils/thread-identity.test.ts @@ -1,5 +1,10 @@ import { describe, it, expect } from "vitest"; -import { newThread, switchThread, adoptThread } from "./thread-identity"; +import { + newThread, + switchThread, + adoptThread, + threadSessionReducer, +} from "./thread-identity"; describe("newThread", () => { it("uses the supplied key and has no chat id yet", () => { @@ -17,6 +22,10 @@ describe("switchThread", () => { }); describe("adoptThread", () => { + // Key UNCHANGED (no remount) + chatId moved null->realId. The unchanged key is + // what keeps the live useChat store alive; the matching chatId is what makes the + // window's render-phase reconciler (activeChatId !== thread.chatId) treat the + // adopted thread as already-in-sync rather than a switch. it("adopts in place for a new chat: keeps the key, sets the chat id", () => { const prev = newThread("new-abc"); expect(adoptThread(prev, "chat-1")).toEqual({ @@ -32,22 +41,39 @@ describe("adoptThread", () => { }; expect(adoptThread(prev, "chat-2")).toBe(prev); }); +}); - it("INVARIANT: adoption never remounts (key unchanged) while chatId changes", () => { - const prev = newThread("new-abc"); - const next = adoptThread(prev, "chat-1"); - // The mount key is preserved (no remount) ... - expect(next.key).toBe(prev.key); - // ... while the chat id moved from null to the real id. - expect(prev.chatId).toBeNull(); - expect(next.chatId).toBe("chat-1"); +describe("threadSessionReducer", () => { + it("reconcile to an existing id switches (key becomes the id)", () => { + const next = threadSessionReducer(newThread("new-abc"), { + type: "reconcile", + chatId: "chat-1", + newKey: "new-xyz", + }); + expect(next).toEqual({ key: "chat-1", chatId: "chat-1" }); }); - it("INVARIANT: after adoption thread.chatId equals the adopted id, so the window's render-phase reconciler (activeChatId !== thread.chatId) does NOT fire and remount the live thread", () => { - const adoptedId = "chat-1"; - const next = adoptThread(newThread("new-abc"), adoptedId); - // The window sets activeChatId to the same adoptedId; this asserts they match - // so the reconciler treats it as already-in-sync, not a switch. - expect(next.chatId).toBe(adoptedId); + it("reconcile to null starts a fresh new thread with the supplied key", () => { + const next = threadSessionReducer(switchThread("chat-1"), { + type: "reconcile", + chatId: null, + newKey: "new-xyz", + }); + expect(next).toEqual({ key: "new-xyz", chatId: null }); + }); + + it("adopt on a new thread keeps the key and sets the id", () => { + const next = threadSessionReducer(newThread("new-abc"), { + type: "adopt", + chatId: "chat-1", + }); + expect(next).toEqual({ key: "new-abc", chatId: "chat-1" }); + }); + + it("adopt on a persisted thread is a no-op", () => { + const prev = switchThread("chat-1"); + expect(threadSessionReducer(prev, { type: "adopt", chatId: "chat-2" })).toBe( + prev, + ); }); }); diff --git a/apps/client/src/features/ai-chat/utils/thread-identity.ts b/apps/client/src/features/ai-chat/utils/thread-identity.ts index 4b289615..f5408fce 100644 --- a/apps/client/src/features/ai-chat/utils/thread-identity.ts +++ b/apps/client/src/features/ai-chat/utils/thread-identity.ts @@ -45,3 +45,29 @@ export function switchThread(chatId: string): ThreadIdentity { export function adoptThread(prev: ThreadIdentity, chatId: string): ThreadIdentity { return prev.chatId === null ? { key: prev.key, chatId } : prev; } + +/** + * Thread-identity transitions as a reducer action. See `threadSessionReducer`. + */ +export type ThreadSessionAction = + | { type: "reconcile"; chatId: string | null; newKey: string } + | { type: "adopt"; chatId: string }; + +/** + * Single source of truth for thread-identity transitions. `reconcile` handles a + * genuine switch (user OR external atom write) -> remount; `adopt` moves a brand- + * new chat to its real id in place (no remount). + */ +export function threadSessionReducer( + state: ThreadIdentity, + action: ThreadSessionAction, +): ThreadIdentity { + switch (action.type) { + case "reconcile": + return action.chatId === null + ? newThread(action.newKey) + : switchThread(action.chatId); + case "adopt": + return adoptThread(state, action.chatId); + } +} diff --git a/apps/server/src/core/ai-chat/ai-chat.service.ts b/apps/server/src/core/ai-chat/ai-chat.service.ts index b8d925e4..71550f6c 100644 --- a/apps/server/src/core/ai-chat/ai-chat.service.ts +++ b/apps/server/src/core/ai-chat/ai-chat.service.ts @@ -562,12 +562,9 @@ export class AiChatService { } /** - * Compute the `messageMetadata` payload for the streamed assistant UI message. - * The AI SDK invokes `messageMetadata` on each stream part (ai@6); we attach the - * authoritative `chatId` ONLY on the `start` part so it reaches the client at the - * very first chunk (as `message.metadata.chatId`). The client adopts THAT id for a - * new chat instead of guessing the newest chat in its list — fixing the two-tab - * adoption race (#137). Returning undefined for any non-start part adds no metadata. + * Attach the authoritative `chatId` to the streamed assistant message's `start` + * part (as `message.metadata.chatId`) so the client can adopt the real id for a + * new chat. See the client's adopt-chat-id.ts for the full #137 design. */ export function chatStreamStartMetadata( part: { type: string },