From 5b59a70e3ff1d99f456615c9ee1e41fc1b16c277 Mon Sep 17 00:00:00 2001 From: claude code agent 227 Date: Fri, 26 Jun 2026 05:22:20 +0300 Subject: [PATCH] fix(ai-chat): New chat during first-turn stream resets the chat, not just the role badge (#161) Co-Authored-By: Claude Opus 4.8 (1M context) --- .../ai-chat/components/ai-chat-window.tsx | 17 +++++- .../ai-chat/components/chat-thread.tsx | 20 +++++-- .../ai-chat/hooks/use-chat-session.test.tsx | 46 +++++++++++++- .../ai-chat/hooks/use-chat-session.ts | 60 ++++++++++++++++++- 4 files changed, 133 insertions(+), 10 deletions(-) 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 de0b9923..b3d003db 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 @@ -193,6 +193,7 @@ export default function AiChatWindow() { const { threadKey, waitingForHistory, + startFreshThread, onTurnFinished, onServerChatId, cancelPendingAdoption, @@ -215,12 +216,25 @@ export default function AiChatWindow() { // just-failed chat after they chose a fresh one. const startNewChat = useCallback((): void => { cancelPendingAdoption(); + // Force a fresh, empty thread UNCONDITIONALLY (#161). Pressing "New chat" + // while a brand-new chat's first turn is still streaming leaves activeChatId + // null (the real id is adopted only at turn end), so setActiveChatId(null) + // alone is a no-op and the reconciler never remounts — the chat/stream/history + // would persist and only the role badge would drop. This always remounts the + // thread into a clean new chat. + startFreshThread(); setActiveChatId(null); setHistoryOpen(false); setDraft(""); // Default the picker back to "Universal assistant" for the fresh chat. setSelectedRoleId(null); - }, [cancelPendingAdoption, setActiveChatId, setDraft, setSelectedRoleId]); + }, [ + cancelPendingAdoption, + startFreshThread, + setActiveChatId, + setDraft, + setSelectedRoleId, + ]); const selectChat = useCallback( (chatId: string): void => { @@ -622,6 +636,7 @@ export default function AiChatWindow() { ) : ( void; + * undefined on a failed turn — see adopt-chat-id.ts for the full #137 design. + * `finishingThreadKey` (this thread's mount key) lets the session ignore a turn + * finishing on a thread already abandoned by New chat mid-stream (#161). */ + onTurnFinished: (serverChatId?: string, finishingThreadKey?: string) => void; /** Called EARLY (at the stream's `start` chunk) with the authoritative server * chat id streamed on the assistant message metadata, so a brand-new chat * adopts its real id WHILE the first turn is still streaming (#174 — makes the @@ -109,6 +116,7 @@ function rowToUiMessage(row: IAiChatMessageRow): UIMessage { */ export default function ChatThread({ chatId, + threadKey, initialRows, openPage, roleId, @@ -257,8 +265,10 @@ export default function ChatThread({ 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)); + // chat — see adopt-chat-id.ts for the full #137 design. `threadKey` lets the + // session ignore this finish if it belongs to a thread abandoned by New chat + // mid-stream (#161). + onTurnFinished(extractServerChatId(message), threadKey); // 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); @@ -279,7 +289,7 @@ export default function ChatThread({ // Surface the raw failure in the browser console (devtools) for debugging; // the UI separately shows a friendly classified banner (see errorView). console.error("AI chat stream error:", streamError); - onTurnFinished(); + onTurnFinished(undefined, threadKey); }, }); 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 index 0080cc80..39a72628 100644 --- 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 @@ -1,5 +1,5 @@ import { describe, it, expect, vi, beforeEach } from "vitest"; -import { renderHook } from "@testing-library/react"; +import { renderHook, act } from "@testing-library/react"; import { useChatSession } from "./use-chat-session"; import type { UseChatSessionOptions } from "./use-chat-session"; @@ -227,6 +227,50 @@ describe("useChatSession", () => { expect(result.current.threadKey).toBe("C"); }); + it("#161: New chat during a streaming first turn forces a fresh thread (remount), not just a no-op", () => { + // Brand-new chat whose first turn is still streaming: the id is adopted only + // at turn end, so activeChatId AND thread.chatId are both null. Pressing "New + // chat" must still remount to a clean thread even though the atom is unchanged + // — the render-phase reconciler (null === null) would otherwise do nothing, + // leaving the old chat/stream/history in place (the bug: only the role badge + // dropped). + const { result } = setup({ activeChatId: null, chats: { items: [] } }); + const keyBefore = result.current.threadKey; + act(() => result.current.startFreshThread()); + expect(result.current.threadKey).not.toBe(keyBefore); + }); + + it("#161: an abandoned thread's late onTurnFinished does NOT adopt its chat (thread-aware guard)", () => { + // New chat mid-stream remounts to a fresh thread, but @ai-sdk/react does not + // abort the abandoned stream on unmount: its onFinish still fires later with + // the real server id, tagged with the OLD (abandoned) mount key. That must not + // adopt — it would yank the user back into the chat they just left. + const { result, setActiveChatId, onInvalidateChatList } = setup({ + activeChatId: null, + chats: { items: [] }, + }); + const abandonedKey = result.current.threadKey; + act(() => result.current.startFreshThread()); + expect(result.current.threadKey).not.toBe(abandonedKey); + // The abandoned turn finishes in the background, streaming its real id "A". + result.current.onTurnFinished("A", abandonedKey); + expect(setActiveChatId).not.toHaveBeenCalledWith("A"); + // It still refreshes the chat list so the left-behind chat shows in history. + expect(onInvalidateChatList).toHaveBeenCalled(); + }); + + it("#161: a turn finishing on the CURRENT thread still adopts (guard is key-scoped, not blanket)", () => { + // The happy path must keep working: onTurnFinished tagged with the mounted + // thread's own key adopts in place as before. + const { result, setActiveChatId } = setup({ + activeChatId: null, + chats: { items: [] }, + }); + const currentKey = result.current.threadKey; + result.current.onTurnFinished("A", currentKey); + expect(setActiveChatId).toHaveBeenCalledWith("A"); + }); + 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({ 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 index d21ebd11..14420ad0 100644 --- a/apps/client/src/features/ai-chat/hooks/use-chat-session.ts +++ b/apps/client/src/features/ai-chat/hooks/use-chat-session.ts @@ -31,9 +31,19 @@ export interface UseChatSessionResult { threadKey: string; /** Show the history loader instead of the live thread. */ waitingForHistory: boolean; + /** Force a brand-new, empty thread (new mount key, no chat id) UNCONDITIONALLY, + * even when `activeChatId` is unchanged. The window calls this from + * startNewChat so "New chat" pressed WHILE a brand-new chat's first turn is + * still streaming (activeChatId still null, nothing to diverge) actually + * resets the chat instead of only dropping the role badge (#161). */ + startFreshThread: () => void; /** 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; + * (undefined on a failed turn). `finishingThreadKey` is the mount key of the + * thread that produced the turn (omit => "current thread", back-compatible): + * a turn ABANDONED by New chat mid-stream still fires this after its thread + * unmounted, so adoption is gated to the still-mounted thread (#161). Handles + * new-chat id adoption + invalidations. */ + onTurnFinished: (serverChatId?: string, finishingThreadKey?: string) => void; /** Call EARLY (at the stream's `start` chunk) with the authoritative streamed * chat id so a brand-new chat adopts its real id WHILE its first turn is still * streaming — making `activeChatId`-gated affordances (e.g. the Copy/export @@ -98,6 +108,15 @@ export function useChatSession( : switchThread(activeChatId), ); + // Live mirror of the mounted thread's mount key, read by onTurnFinished to tell + // the CURRENT thread from one ABANDONED by New chat mid-stream. @ai-sdk/react + // does not abort a stream on unmount and proxies callbacks through a ref, so an + // abandoned turn's onFinish/onError still fires AFTER its ChatThread unmounted; + // matching its key against this ref keeps that late finish from adopting the + // abandoned chat and yanking the user out of the fresh chat they opened (#161). + const threadKeyRef = useRef(thread.key); + threadKeyRef.current = thread.key; + // 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 @@ -115,7 +134,23 @@ export function useChatSession( // 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) => { + (serverChatId?: string, finishingThreadKey?: string) => { + // Thread-aware guard (#161). A turn ABANDONED by "New chat" mid-stream still + // fires onFinish/onError after its ChatThread unmounted (@ai-sdk/react does + // not abort on unmount and proxies callbacks through a ref). If that late + // finish ran the adoption path it would set activeChatId to the abandoned + // chat's real id and yank the user out of the fresh chat they just opened. + // So adopt / arm the fallback ONLY for the still-mounted thread; an + // abandoned one merely refreshes the chat list (so the left-behind chat + // surfaces in history) and does nothing else. A missing key (undefined) + // means "current thread" — keeps old call sites/tests working. + if ( + finishingThreadKey !== undefined && + finishingThreadKey !== threadKeyRef.current + ) { + onInvalidateChatList(); + return; + } // 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. @@ -258,9 +293,28 @@ export function useChatSession( pendingNewChatRef.current = null; }, []); + // Force a fresh, empty thread regardless of `activeChatId` (#161). The render- + // phase reconciler only remounts when activeChatId diverges from thread.chatId, + // so "New chat" pressed while a brand-new chat's first turn is still streaming + // (activeChatId AND thread.chatId both null — the real id is adopted only at the + // end of the turn) is a no-op for it and the abandoned thread/stream/history + // would persist. Dispatching reconcile with a fresh key and chatId:null here + // always produces a new mount key, so React remounts ChatThread (a clean useChat + // store) and the post-dispatch state (activeChatId null === thread.chatId null) + // keeps the reconciler from interfering. Also disarms any pending fallback. + const startFreshThread = useCallback(() => { + pendingNewChatRef.current = null; + dispatch({ + type: "reconcile", + chatId: null, + newKey: `new-${generateId()}`, + }); + }, []); + return { threadKey: thread.key, waitingForHistory, + startFreshThread, onTurnFinished, onServerChatId, cancelPendingAdoption,