fix(ai-chat): reset chat on New chat during the first turn's stream (#161)

Pressing "New chat" while a brand-new chat's first turn was still streaming
only removed the role badge — the thread, session and stream were kept. The
thread remount is driven by a render-phase reconciler in useChatSession that
fires only when activeChatId !== thread.chatId; on a not-yet-adopted new chat
both are null, so startNewChat's setActiveChatId(null) was a no-op and nothing
remounted.

- useChatSession: add startFreshThread(), which unconditionally dispatches a
  reconcile to a fresh mount key so ChatThread remounts even when activeChatId
  stays null. startNewChat now calls it.
- Guard against the abandoned thread's late finish: ai@6's useChat does not
  abort on unmount and proxies its callbacks, so onFinish/onError of the chat
  the user just left still fire and would adopt it back. onTurnFinished now
  takes the finishing thread's key and, when it no longer matches the mounted
  thread, only refreshes the chat list — no adoption, no fallback arming, no
  per-chat message invalidation. ChatThread forwards its threadKey in both
  onFinish and onError. An omitted key (legacy callers/tests) counts as the
  current thread, preserving the #137 adoption behavior.

Tests: 3 new useChatSession cases (fresh-thread remount with activeChatId null,
abandoned-thread finish rejected, current-thread finish still adopts). Full
ai-chat suite green (183).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

apps/client/src/features/ai-chat/components/ai-chat-window.tsx

@@ -185,8 +185,13 @@ export default function AiChatWindow() {
   // 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 } =
+  const {
-    useChatSession({
+    threadKey,
+    waitingForHistory,
+    onTurnFinished,
+    startFreshThread,
+    cancelPendingAdoption,
+  } = useChatSession({
       activeChatId,
       setActiveChatId,
       chats,
@@ -205,12 +210,25 @@ export default function AiChatWindow() {
   // just-failed chat after they chose a fresh one.
   const startNewChat = useCallback((): void => {
     cancelPendingAdoption();
+    // Force a fresh thread UNCONDITIONALLY. If the user is still on a brand-new,
+    // not-yet-adopted chat (activeChatId === null) while its first turn streams,
+    // setActiveChatId(null) is a no-op and the render-phase reconciler would not
+    // remount — leaving the streaming thread in place (#161). startFreshThread
+    // guarantees a clean remount; the abandoned thread's late finish is rejected
+    // by the threadKey guard in the session hook.
+    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 => {
@@ -621,6 +639,7 @@ export default function AiChatWindow() {
               onRolePicked={(role) => setSelectedRoleId(role.id)}
               assistantName={currentRole?.name}
               onTurnFinished={onTurnFinished}
+              threadKey={threadKey}
               liveStateRef={liveThreadRef}
               onLiveTurnTokens={setLiveTurnTokens}
             />

apps/client/src/features/ai-chat/components/chat-thread.tsx

@@ -67,7 +67,12 @@ interface ChatThreadProps {
    *  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;
+  onTurnFinished: (serverChatId?: string, finishingThreadKey?: string) => void;
+  /** This thread's mount key (the same value the parent uses as React `key`).
+   *  Forwarded back through onTurnFinished so the session hook can tell a finish
+   *  from THIS still-mounted thread from a late finish of an abandoned thread the
+   *  user already left via New chat / switch (#161). */
+  threadKey?: string;
   /** 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
@@ -123,6 +128,7 @@ export default function ChatThread({
   onRolePicked,
   assistantName,
   onTurnFinished,
+  threadKey,
   liveStateRef,
   onLiveTurnTokens,
 }: ChatThreadProps) {
@@ -265,7 +271,7 @@ export default function ChatThread({
       // 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));
+      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);
@@ -286,7 +292,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);
     },
   });
 

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";
 
@@ -187,6 +187,48 @@ describe("useChatSession", () => {
     expect(result.current.threadKey).toBe("C");
   });
 
+  it("#161: startFreshThread remounts even when activeChatId stays null (New chat mid-stream)", () => {
+    // The bug: pressing New chat while still on a brand-new, not-yet-adopted chat
+    // leaves activeChatId === null, so the render-phase reconciler never fires.
+    // startFreshThread must remount UNCONDITIONALLY (a new mount key).
+    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: a late finish from an ABANDONED thread does not adopt or invalidate messages", () => {
+    const {
+      result,
+      setActiveChatId,
+      onInvalidateChatList,
+      onInvalidateChatMessages,
+    } = setup({ activeChatId: null, chats: { items: [{ id: "x" }] } });
+    const abandonedKey = result.current.threadKey;
+    // User pressed New chat mid-stream → fresh thread (new mount key).
+    act(() => result.current.startFreshThread());
+    expect(result.current.threadKey).not.toBe(abandonedKey);
+    // The left-behind thread's onFinish fires late, carrying ITS (now stale) key
+    // and the server id of the chat the user just left. It must NOT be adopted.
+    act(() => result.current.onTurnFinished("A", abandonedKey));
+    expect(setActiveChatId).not.toHaveBeenCalled();
+    expect(onInvalidateChatMessages).not.toHaveBeenCalled();
+    // The abandoned chat should still surface in the history list.
+    expect(onInvalidateChatList).toHaveBeenCalled();
+  });
+
+  it("#161: a finish from the CURRENT thread (matching key) still adopts", () => {
+    const { result, setActiveChatId } = setup({
+      activeChatId: null,
+      chats: { items: [{ id: "x" }] },
+    });
+    // Same thread that is mounted reports its finish with the matching key.
+    act(() =>
+      result.current.onTurnFinished("A", result.current.threadKey),
+    );
+    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({

apps/client/src/features/ai-chat/hooks/use-chat-session.ts

@@ -32,8 +32,20 @@ export interface UseChatSessionResult {
   /** 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. */
+   *  (undefined on a failed turn). `finishingThreadKey` is the mount key of the
-  onTurnFinished: (serverChatId?: string) => void;
+   *  thread that produced this callback — when it no longer matches the mounted
+   *  thread (the user pressed New chat / switched mid-stream), the call is from an
+   *  abandoned thread and must NOT adopt or re-arm the fallback. Omitting it (old
+   *  callers / tests) treats the call as belonging to the current thread. Handles
+   *  new-chat id adoption + invalidations. */
+  onTurnFinished: (serverChatId?: string, finishingThreadKey?: string) => void;
+  /** Force a brand-new, empty thread (new mount key, no chat id) UNCONDITIONALLY.
+   *  The render-phase reconciler only remounts when `activeChatId` actually
+   *  changes; pressing "New chat" while already in a not-yet-adopted new chat
+   *  leaves `activeChatId === null` (a no-op for the atom), so the reconciler
+   *  never fires and the stale streaming thread stays mounted (#161). The window
+   *  calls this from startNewChat to guarantee a fresh thread regardless. */
+  startFreshThread: () => 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. */
@@ -78,6 +90,14 @@ export function useChatSession(
   const activeChatIdRef = useRef(activeChatId);
   activeChatIdRef.current = activeChatId;
 
+  // Live mirror of the mounted thread's key, read by onTurnFinished to tell a
+  // current-thread finish from an ABANDONED one. ai@6's useChat does not abort
+  // its request on unmount, and its callbacks are proxied so onFinish/onError of
+  // a thread the user already left (via New chat / switch) still fire AFTER that
+  // thread unmounts. By then this ref holds the NEW thread's key, so comparing it
+  // to the key the finishing thread reports rejects the abandoned turn (#161).
+  const threadKeyRef = useRef<string>("");
+
   // 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
@@ -94,6 +114,10 @@ export function useChatSession(
         : switchThread(activeChatId),
   );
 
+  // Keep the live mirror pointed at the currently-mounted thread's key so a late
+  // onTurnFinished can be matched against it (see threadKeyRef above).
+  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
@@ -111,7 +135,21 @@ 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) => {
+      // Reject a finish from an ABANDONED thread. After the user pressed New chat
+      // (or switched chats) mid-stream, the left-behind thread's onFinish/onError
+      // still fire (ai@6 does not abort on unmount). Adopting/arming off that late
+      // callback would yank the user back into the chat they just left (#161).
+      // `undefined` (legacy callers/tests) is treated as the current thread.
+      const isCurrentThread =
+        finishingThreadKey === undefined ||
+        finishingThreadKey === threadKeyRef.current;
+      if (!isCurrentThread) {
+        // Still surface the abandoned chat in the history list, but do NOT adopt,
+        // arm the fallback, or invalidate per-chat messages (no thread shows it).
+        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.
@@ -229,10 +267,28 @@ export function useChatSession(
     pendingNewChatRef.current = null;
   }, []);
 
+  // Force a fresh, empty thread regardless of the current `activeChatId`. The
+  // render-phase reconciler only remounts on an activeChatId CHANGE, so "New chat"
+  // pressed while already in a not-yet-adopted new chat (activeChatId stays null)
+  // would otherwise leave the in-flight streaming thread mounted (#161). Dispatch
+  // `reconcile` to chatId:null with a brand-new key so React remounts ChatThread
+  // (a fresh useChat store). Disarm any armed fallback too. After this dispatch
+  // thread.chatId is null; the window also sets activeChatId to null, so the
+  // render-phase reconciler then finds them equal and does not double-remount.
+  const startFreshThread = useCallback(() => {
+    pendingNewChatRef.current = null;
+    dispatch({
+      type: "reconcile",
+      chatId: null,
+      newKey: `new-${generateId()}`,
+    });
+  }, []);
+
   return {
     threadKey: thread.key,
     waitingForHistory,
     onTurnFinished,
+    startFreshThread,
     cancelPendingAdoption,
   };
 }