diff --git a/AGENTS.md b/AGENTS.md index 333445a9..48a282ad 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -210,7 +210,7 @@ pnpm --filter @docmost/mcp test # node --test (unit + mock) pnpm --filter @docmost/mcp test:e2e # MCP end-to-end against a live instance ``` -**Database migrations** (Kysely, run from `apps/server`; they auto-run on server startup too): +**Database migrations** (Kysely, run from `apps/server`). **Where they auto-apply:** in **production** (the built image / `start:prod`) pending migrations run automatically on server boot. In **local dev** (the `pnpm dev` stand / `nest start --watch`) they do **NOT** auto-run — after you pull or switch branches you must apply them yourself with `pnpm --filter server migration:latest`, or any endpoint touching a new column/table 500s (e.g. a freshly-added `ai_chats.page_id` blanket-500s all of AI chat until migrated). ```bash pnpm --filter server migration:create --name=my_change # new empty migration pnpm --filter server migration:latest # apply all pending 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 6d543289..5ec80874 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,6 +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 { useChatSession } from "@/features/ai-chat/hooks/use-chat-session.ts"; import { shouldCollapseOnOutsidePointer, isHeaderClick, @@ -101,7 +101,8 @@ function clampGeom(g: { left: number; top: number; width: number; height: number /** * Floating, draggable, resizable, minimizable AI chat window. Replaces the * former right-aside `AiChatPanel`: it owns ALL chat orchestration (active - * chat, new chat, adopt-new-chat, open-page context, token sum) and wraps the + * chat, new chat, in-place id adoption from streamed metadata, open-page + * context, token sum) and wraps the * reused inner components (ConversationList + ChatThread) in window chrome * ported from the GitmostAgent.jsx design. */ @@ -132,30 +133,6 @@ export default function AiChatWindow() { // left partly off-screen). const [geom, setGeom] = useAtom(aiChatWindowGeomAtom); - // Track whether we are awaiting the id of a just-created (new) chat, so we - // can adopt it once the chat list refreshes after the first turn finishes. - const adoptNewChat = useRef(false); - - // 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); - - // Mount key for ChatThread + the chat the currently-mounted thread represents. - // `threadKey` normally tracks the active chat, so selecting a different chat - // (incl. from page history) remounts and re-seeds. The ONE exception is - // in-place adoption of a brand-new chat's server id: the adopt effect moves - // `liveThreadChatId` to the new id TOGETHER with `activeChatId`, so the switch - // check below does not fire and the SAME thread stays mounted (its useChat - // already holds the just-finished turn) instead of being re-seeded from - // not-yet-persisted history. - const [threadKey, setThreadKey] = useState( - () => activeChatId ?? `new-${generateId()}`, - ); - const [liveThreadChatId, setLiveThreadChatId] = useState( - activeChatId, - ); - const { data: chats } = useAiChatsQuery(); // Roles for the new-chat picker (any member may list them). Only fetched while // the window is open. @@ -196,21 +173,42 @@ export default function AiChatWindow() { ? { id: openPageData.id, title: openPageData.title } : null; + // 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. + // The invalidate closures are passed inline: `onTurnFinished` is read live by + // useChat's onFinish (never in an effect dep array), so their identity does not + // matter — no memoization ceremony needed. + const { threadKey, waitingForHistory, onTurnFinished, cancelPendingAdoption } = + useChatSession({ + activeChatId, + setActiveChatId, + chats, + messagesLoading, + onInvalidateChatList: () => + queryClient.invalidateQueries({ queryKey: AI_CHATS_RQ_KEY }), + onInvalidateChatMessages: (id) => + queryClient.invalidateQueries({ queryKey: AI_CHAT_MESSAGES_RQ_KEY(id) }), + }); + + // startNewChat/selectChat set the public atom; the hook's render-phase + // reconciler handles the remount when activeChatId actually CHANGES. But + // pressing "New chat" while already in a new chat leaves activeChatId === null + // (a no-op for the atom), so the reconciler never fires — explicitly disarm any + // armed error-path fallback here so a late refetch can't yank the user into a + // just-failed chat after they chose a fresh one. const startNewChat = useCallback((): void => { - // Cancel any pending adoption so a just-finished new chat can't yank the user - // back here after they explicitly started a fresh one. - adoptNewChat.current = false; + cancelPendingAdoption(); setActiveChatId(null); setHistoryOpen(false); setDraft(""); // Default the picker back to "Universal assistant" for the fresh chat. setSelectedRoleId(null); - }, [setActiveChatId, setDraft, setSelectedRoleId]); + }, [cancelPendingAdoption, setActiveChatId, setDraft, setSelectedRoleId]); const selectChat = useCallback( (chatId: string): void => { - // Cancel any pending adoption so it can't override an explicit selection. - adoptNewChat.current = false; + cancelPendingAdoption(); setActiveChatId(chatId); setHistoryOpen(false); setDraft(""); @@ -218,30 +216,9 @@ export default function AiChatWindow() { // chat's header/assistant-name (which prefers the chat's persisted role). setSelectedRoleId(null); }, - [setActiveChatId, setDraft, setSelectedRoleId], + [cancelPendingAdoption, 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; adopt the newest chat id so the - // thread switches from "new" to the persisted chat (and loads its history on - // later opens). - const onTurnFinished = useCallback(() => { - if (activeChatId === null) adoptNewChat.current = true; - 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 threadKey / 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, queryClient]); - // 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. const activeChat = useMemo( @@ -294,62 +271,6 @@ export default function AiChatWindow() { notifications.show({ message: t("Copied") }); }, [activeChatId, messageRows, activeChat, clipboard, t]); - // When awaiting a new chat's id, adopt the most-recent chat (the list is - // ordered newest-first) once it appears. - useEffect(() => { - if (!adoptNewChat.current) return; - const newest = chats?.items?.[0]; - if (newest) { - adoptNewChat.current = false; - // In-place adoption: move the active chat AND the live-thread marker to the - // new id together, so the threadKey derivation below sees no "switch" and - // keeps the SAME mounted thread (its useChat already holds the finished - // turn) instead of remounting and re-seeding from not-yet-persisted history. - // ASSUMPTION: these two updates (jotai atom + useState) must land in ONE - // render so the render-phase guard never observes the new activeChatId with - // a stale liveThreadChatId (which would wrongly remount). React 18 automatic - // batching inside this effect callback guarantees that; if the store/atom - // mechanism ever changes, gate adoption on an explicit flag instead. - setLiveThreadChatId(newest.id); - setActiveChatId(newest.id); - } - }, [chats, setActiveChatId]); - - // Adjust the derived thread state during render when the active chat genuinely - // changes — the React-sanctioned alternative to an effect (it re-renders before - // paint, no extra commit, and converges since the next render finds them equal). - // In-place adoption of a new chat's id never reaches here because the adopt - // effect moves liveThreadChatId in lockstep with activeChatId. - if (activeChatId !== liveThreadChatId) { - setLiveThreadChatId(activeChatId); - setThreadKey(activeChatId ?? `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 && - threadKey === 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 && - threadKey === 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: diff --git a/apps/client/src/features/ai-chat/components/chat-thread.tsx b/apps/client/src/features/ai-chat/components/chat-thread.tsx index 9ee92377..b5dc6d48 100644 --- a/apps/client/src/features/ai-chat/components/chat-thread.tsx +++ b/apps/client/src/features/ai-chat/components/chat-thread.tsx @@ -22,6 +22,7 @@ import { IAiRole, } from "@/features/ai-chat/types/ai-chat.types.ts"; import { describeChatError } from "@/features/ai-chat/utils/error-message.ts"; +import { extractServerChatId } from "@/features/ai-chat/utils/adopt-chat-id.ts"; import { dequeue, enqueueMessage, @@ -57,9 +58,11 @@ interface ChatThreadProps { /** Display name for the assistant label / typing line (the role name); * forwarded to MessageList. Absent => 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. */ - onTurnFinished: () => void; + /** 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 * "Copy chat" export can include the in-progress, not-yet-persisted @@ -246,8 +249,11 @@ export default function ChatThread({ // sending after the user hit Stop — or blindly retrying after a failure — // would be wrong, so on Stop/disconnect/error the queue is left intact for // the user to decide. - onFinish: ({ isAbort, isDisconnect, isError }) => { - onTurnFinished(); + onFinish: ({ message, isAbort, isDisconnect, isError }) => { + // 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); @@ -259,9 +265,11 @@ export default function ChatThread({ }, // `onError` runs in addition to `onFinish` (which ai@6 also calls on error). // Log the raw failure here for devtools; the UI shows a friendly classified - // banner via `error` below. We still call `onTurnFinished()` (idempotent with - // the onFinish call) so a brand-new chat that fails its first turn is adopted - // and the chat list refreshes immediately rather than after a manual refresh. + // banner via `error` below. We still call `onTurnFinished()` with NO server id + // (idempotent with the onFinish call): for a brand-new chat that ARMS the + // bounded list-refetch fallback (adopt the single newly-appeared chat once the + // refetch lands); for an existing chat it just refreshes the chat list + // immediately rather than after a manual refresh. onError: (streamError) => { // Surface the raw failure in the browser console (devtools) for debugging; // the UI separately shows a friendly classified banner (see errorView). 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..8104d1e6 --- /dev/null +++ b/apps/client/src/features/ai-chat/hooks/use-chat-session.test.tsx @@ -0,0 +1,206 @@ +import { describe, it, expect, vi, beforeEach } from "vitest"; +import { renderHook } from "@testing-library/react"; +import { useChatSession } from "./use-chat-session"; +import type { UseChatSessionOptions } from "./use-chat-session"; + +// The props the test drives: the parent-owned subset of UseChatSessionOptions +// (the spies are injected by setup, not per-render). messagesLoading is optional +// here (defaulted to false in setup) for terser test call sites. +type DriverProps = Pick & { + messagesLoading?: boolean; +}; + +// 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: DriverProps) { + const setActiveChatId = vi.fn(); + const onInvalidateChatList = vi.fn(); + const onInvalidateChatMessages = vi.fn(); + const { result, rerender } = renderHook( + (props: DriverProps) => + 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("startNewChat while already in a new chat: cancelPendingAdoption stops a late refetch adopting the failed chat", () => { + // The Warning path the render-phase reconciler can't catch: pressing "New + // chat" while already in a new chat keeps activeChatId === null (a no-op for + // the atom), so only the explicit cancelPendingAdoption() disarms. + const { result, rerender, setActiveChatId } = setup({ + activeChatId: null, + chats: { items: [{ id: "x" }] }, + }); + result.current.onTurnFinished(undefined); // first turn failed → arm (before=["x"]) + result.current.cancelPendingAdoption(); // window calls this from startNewChat + // The just-failed row lands in a late refetch; it must NOT be adopted. + rerender({ + activeChatId: null, + chats: { items: [{ id: "x" }, { id: "failed" }] }, + }); + expect(setActiveChatId).not.toHaveBeenCalledWith("failed"); + }); + + it("onTurnFinished for an existing chat: no adoption, invalidates that chat's messages", () => { + const { + result, + setActiveChatId, + onInvalidateChatList, + onInvalidateChatMessages, + } = setup({ activeChatId: "chat-1", chats: { items: [{ id: "chat-1" }] } }); + result.current.onTurnFinished("chat-1"); + expect(setActiveChatId).not.toHaveBeenCalled(); // existing chat is never re-adopted + expect(onInvalidateChatList).toHaveBeenCalled(); + expect(onInvalidateChatMessages).toHaveBeenCalledWith("chat-1"); + }); + + it("double onTurnFinished on a failed-after-start turn: primary adopt, 2nd no-id call does NOT re-arm the fallback", () => { + // ai@6 fires onFinish AND onError on a failed turn. If the failure happened + // AFTER the `start` chunk, onFinish carries the streamed id and onError does + // not — so onTurnFinished runs twice in one turn (id, then no-id) before any + // re-render. The 2nd call must NOT re-arm the fallback off the still-null + // closure; otherwise a late refetch (parent hasn't reflected the adoption yet) + // would wrongly adopt a sibling row. + const { result, rerender, setActiveChatId } = setup({ + activeChatId: null, + chats: { items: [{ id: "x" }] }, + }); + result.current.onTurnFinished("A"); // onFinish: primary adoption + expect(setActiveChatId).toHaveBeenCalledWith("A"); + result.current.onTurnFinished(undefined); // onError: same turn, no id + // Even in the worst case (the parent has NOT yet reflected activeChatId="A" + // and a late refetch lands a new row), the just-failed sibling must NOT be + // adopted. Two layers guarantee this: the ref guard keeps the 2nd call from + // re-arming at the source, and the render-phase reconciler disarms anything + // stale once thread.chatId ("A") diverges from the still-null activeChatId. + rerender({ + activeChatId: null, + chats: { items: [{ id: "x" }, { id: "late" }] }, + }); + expect(setActiveChatId).not.toHaveBeenCalledWith("late"); + }); + + 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..998f2631 --- /dev/null +++ b/apps/client/src/features/ai-chat/hooks/use-chat-session.ts @@ -0,0 +1,238 @@ +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"; + +/** Inputs to {@link useChatSession}. `activeChatId`/`setActiveChatId` are the + * public selection atom (also written from outside the window, e.g. page + * history); the rest is read-only context the hook needs. */ +export interface UseChatSessionOptions { + 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; +} + +/** What the window needs from a chat session: the ChatThread mount key, the + * history-loader gate, and the turn-finished callback. */ +export interface UseChatSessionResult { + /** 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; + /** Disarm any pending error-path new-chat fallback. The window calls this from + * startNewChat/selectChat so a late refetch can't yank the user back into a + * just-failed chat after they explicitly moved on. */ + cancelPendingAdoption: () => void; +} + +/** Project a chat list to its id array (the before/after snapshot for the + * error-path fallback). */ +function chatIdSnapshot( + chats: { items?: { id: string }[] } | undefined, +): string[] { + return chats?.items?.map((c) => c.id) ?? []; +} + +/** + * 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: UseChatSessionOptions, +): UseChatSessionResult { + const { + activeChatId, + setActiveChatId, + chats, + messagesLoading, + onInvalidateChatList, + onInvalidateChatMessages, + } = params; + + // Live mirror of `activeChatId`, read by onTurnFinished. ai@6 fires both + // onFinish AND onError on a failed turn, so onTurnFinished can run twice in one + // turn (once with the streamed id, once without) BEFORE a re-render. Reading + // the ref — which the primary-adoption branch updates imperatively — makes that + // second call see the just-adopted id, so it cannot re-arm the fallback. (A + // plain closure over `activeChatId` would still read null on the second call.) + const activeChatIdRef = useRef(activeChatId); + activeChatIdRef.current = activeChatId; + + // 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) => { + // Read the live id from the ref, not the closure: on a failed turn this can + // run twice in one turn (onFinish + onError) before any re-render, and the + // primary branch below updates the ref so the second call sees the adopted id. + const current = activeChatIdRef.current; + const adopted = resolveAdoptedChatId(current, 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. + activeChatIdRef.current = adopted; // a same-turn 2nd call now sees the id + setActiveChatId(adopted); + dispatch({ type: "adopt", chatId: adopted }); + // Primary adoption won — disarm any previously-armed fallback. + pendingNewChatRef.current = null; + } else if (current === 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 = chatIdSnapshot(chats); + } + 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 `current` is + // still null here; later turns hit this with the adopted id. + if (current) { + onInvalidateChatMessages(current); + } + }, + [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 = chatIdSnapshot(chats); + 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; + + // Explicit disarm for startNewChat/selectChat. The render-phase reconciler only + // disarms when activeChatId actually changes, but "New chat" pressed while the + // user is ALREADY in a new chat is a no-op for the atom (activeChatId stays + // null), so the reconciler never fires — without this an armed fallback could + // adopt the just-failed chat from a late refetch and yank the user out of their + // fresh chat. Stable identity (writes a ref). + const cancelPendingAdoption = useCallback(() => { + pendingNewChatRef.current = null; + }, []); + + return { + threadKey: thread.key, + waitingForHistory, + onTurnFinished, + cancelPendingAdoption, + }; +} 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 new file mode 100644 index 00000000..c9ff117a --- /dev/null +++ b/apps/client/src/features/ai-chat/utils/adopt-chat-id.test.ts @@ -0,0 +1,72 @@ +import { describe, it, expect } from "vitest"; +import { + resolveAdoptedChatId, + newlyAddedChatIds, + extractServerChatId, +} from "./adopt-chat-id"; + +describe("resolveAdoptedChatId", () => { + it("adopts the server id for a brand-new chat (activeChatId null + id)", () => { + expect(resolveAdoptedChatId(null, "chat-1")).toBe("chat-1"); + }); + + it("returns null for an existing chat even with a server id", () => { + expect(resolveAdoptedChatId("chat-existing", "chat-1")).toBeNull(); + }); + + it("returns null for a new chat with no server id", () => { + expect(resolveAdoptedChatId(null, undefined)).toBeNull(); + expect(resolveAdoptedChatId(null, null)).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 new file mode 100644 index 00000000..1993dccc --- /dev/null +++ b/apps/client/src/features/ai-chat/utils/adopt-chat-id.ts @@ -0,0 +1,70 @@ +/** + * 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'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 — + * `newlyAddedChatIds` (the fallback effect adopts only when exactly one id is + * new). This is unambiguous and does not race a second tab the way the old + * "newest chat in the list" guess did. + * ============================================================================ + */ + +/** + * Resolve the id to adopt from the server-streamed metadata. Returns + * `serverChatId` only for a brand-new chat (`activeChatId === null`) that + * received a truthy id; otherwise null (existing chat, or no id streamed). + */ +export function resolveAdoptedChatId( + activeChatId: string | null, + serverChatId: string | null | undefined, +): string | null { + 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))); +} 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 new file mode 100644 index 00000000..eab1eadb --- /dev/null +++ b/apps/client/src/features/ai-chat/utils/thread-identity.test.ts @@ -0,0 +1,79 @@ +import { describe, it, expect } from "vitest"; +import { + newThread, + switchThread, + adoptThread, + threadSessionReducer, +} from "./thread-identity"; + +describe("newThread", () => { + it("uses the supplied key and has no chat id yet", () => { + expect(newThread("new-abc")).toEqual({ key: "new-abc", chatId: null }); + }); +}); + +describe("switchThread", () => { + it("switches to an existing chat: key becomes the chat id", () => { + expect(switchThread("chat-1")).toEqual({ + key: "chat-1", + chatId: "chat-1", + }); + }); +}); + +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({ + key: "new-abc", + chatId: "chat-1", + }); + }); + + it("is a no-op for an already-persisted chat", () => { + const prev: { key: string; chatId: string | null } = { + key: "chat-1", + chatId: "chat-1", + }; + expect(adoptThread(prev, "chat-2")).toBe(prev); + }); +}); + +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("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 new file mode 100644 index 00000000..f5408fce --- /dev/null +++ b/apps/client/src/features/ai-chat/utils/thread-identity.ts @@ -0,0 +1,73 @@ +/** + * Pure transitions for the AI-chat thread's identity: the single source of + * truth tying ChatThread's mount key to the chat id that mounted thread holds. + * + * The window keeps exactly ONE of these in state. Consolidating the mount key + * and the live thread's chat id into one atomic value makes the "stale chat id + * vs key" state unrepresentable: every change goes through one of the explicit + * transitions below, so the key and chatId can never silently diverge. + * + * - `newThread`/`switchThread` produce a key that forces a remount (+ reseed): + * `newThread` for a brand-new (id-less) chat, `switchThread` for an existing + * one. The caller picks which based on whether there is a chat id. + * - `adoptThread` keeps the SAME key so a brand-new chat learns its real id + * WITHOUT remounting (the live useChat store, holding the just-finished turn, + * is preserved and the next turn sends the real chatId). + * + * `newThread` takes the session key from the impure `generateId()` at the call + * site so these stay pure and unit-testable. + */ +export type ThreadIdentity = { key: string; chatId: string | null }; + +/** + * A brand-new chat: a fresh session key and no chat id yet. `newKey` is + * supplied by the caller (generateId() is impure) so this stays pure/testable. + */ +export function newThread(newKey: string): ThreadIdentity { + return { key: newKey, chatId: null }; +} + +/** + * Switch to an EXISTING chat: the mount key becomes the chat id, forcing a + * remount + reseed from the persisted history. (A switch to a brand-new chat + * goes through `newThread` instead — there is no id to key on.) + */ +export function switchThread(chatId: string): ThreadIdentity { + return { key: chatId, chatId }; +} + +/** + * In-place adoption: a brand-new chat (`prev.chatId === null`) learns its real + * id WITHOUT remounting — keep the SAME key, set the chat id. If `prev` already + * has a chatId (not a new chat), this is a no-op (returns `prev`): adoption only + * applies to an as-yet-unadopted new thread. + */ +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.spec.ts b/apps/server/src/core/ai-chat/ai-chat.service.spec.ts index d6f333d0..8f6e48d5 100644 --- a/apps/server/src/core/ai-chat/ai-chat.service.spec.ts +++ b/apps/server/src/core/ai-chat/ai-chat.service.spec.ts @@ -5,6 +5,7 @@ import { rowToUiMessage, prepareAgentStep, buildPartialAssistantRecord, + chatStreamStartMetadata, MAX_AGENT_STEPS, FINAL_STEP_INSTRUCTION, } from './ai-chat.service'; @@ -295,3 +296,20 @@ describe('buildPartialAssistantRecord', () => { expect(rec.text).toBe('half'); }); }); + +/** + * chatStreamStartMetadata: attach the authoritative chatId to the streamed + * assistant UI message ONLY on the `start` part (so the client adopts the real + * created chat id at the first chunk — see #137). Any non-start part adds none. + */ +describe('chatStreamStartMetadata', () => { + it('returns { chatId } for the start part', () => { + expect(chatStreamStartMetadata({ type: 'start' }, 'chat-1')).toEqual({ + chatId: 'chat-1', + }); + }); + + it('returns undefined for a finish part (any non-start part)', () => { + expect(chatStreamStartMetadata({ type: 'finish' }, 'chat-1')).toBeUndefined(); + }); +}); 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 e4932cc2..ed8b77cf 100644 --- a/apps/server/src/core/ai-chat/ai-chat.service.ts +++ b/apps/server/src/core/ai-chat/ai-chat.service.ts @@ -532,6 +532,15 @@ export class AiChatService { stripStreamingHopByHopHeaders(res.raw); result.pipeUIMessageStreamToResponse(res.raw, { headers: { 'X-Accel-Buffering': 'no' }, + // Surface the authoritative chatId on the streamed assistant UI message so + // the client adopts the REAL id of the row we created, instead of guessing + // the newest chat in its list. `messageMetadata` is invoked by the AI SDK + // on the `start` and `finish` stream parts (ai@6); we attach `chatId` on the + // `start` part so it reaches the client (as message.metadata.chatId) at the + // very first chunk — before any second tab can race a newer chat into the + // list. This fixes the two-tab "adoption race" (#137) where a new chat in + // tab A could adopt tab B's id and leak its turns into the wrong row. + messageMetadata: ({ part }) => chatStreamStartMetadata(part, chatId), onError: (error: unknown) => { // Reuse the shared formatter so provider error formatting stays // unified between the log line and the streamed error message. @@ -582,6 +591,18 @@ export class AiChatService { } } +/** + * 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 }, + chatId: string, +): { chatId: string } | undefined { + return part.type === 'start' ? { chatId } : undefined; +} + /** The last message with role 'user' from a useChat payload, if any. */ function lastUserMessage( messages: UIMessage[] | undefined,