refactor(ai-chat): extract useChatSession hook + lock the id lifecycle with tests

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 <noreply@anthropic.com>

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();
+  });
+});

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<string> {
+  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;
 }

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,
+    );
   });
 });

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);
+  }
+}