fix(coach): keep the reply's bubble when its persisted copy arrives

The persisted copy of a streamed reply used to mount as a new bubble, so
what the reader had done on the streamed one (the chart or table view, the
open steps list, the evidence panel) reset when the conversation refetch
landed. A turn now keeps one key from its first streamed frame to its
persisted copy, both render in the same wrapper, and the copy keeps the
live tables instead of refetching them, so React keeps the one bubble.
The wrapper stops being a live region once the turn is settled. The
150 ms twin-hiding window this replaces is gone.

e2e/coach-dialog.spec.ts

@@ -357,7 +357,13 @@ function persisted(completed: number) {
  * Serves the conversation as the stubbed server would hold it: one more
  * persisted turn after every POST. Returns the parsed body of every POST.
  */
-async function mockCoach(page: Page): Promise<Array<Record<string, unknown>>> {
+async function mockCoach(
+  page: Page,
+  opts: {
+    /** Holds the first detail refetch after a turn until this resolves. */
+    holdRefetch?: Promise<void>;
+  } = {},
+): Promise<Array<Record<string, unknown>>> {
   const posts: Array<Record<string, unknown>> = [];
   let completed = 0;
 
@@ -401,6 +407,7 @@ async function mockCoach(page: Page): Promise<Array<Record<string, unknown>>> {
         return fulfilJson(route, { results: [] });
       }
       if (url.pathname === `/api/insights/chat/${CONVERSATION_ID}`) {
+        if (completed > 0 && opts.holdRefetch) await opts.holdRefetch;
         return fulfilJson(route, {
           ...summary(),
           attachmentCount: 0,
@@ -448,8 +455,9 @@ async function expectNoAxeViolations(page: Page, label: string) {
 
 /**
  * The streamed reply is replaced by its persisted copy once the detail
- * refetch after `done` lands. The two are separate mounts, so anything
- * toggled on the streamed one before the swap is lost; wait for it.
+ * refetch after `done` lands. The copy takes over the same bubble, so what
+ * the reader toggled survives the swap (pinned below); the steps wait for it
+ * only to read the settled state.
  */
 async function waitForPersistedTwin(page: Page) {
   await expect(
@@ -614,4 +622,55 @@ test.describe("Coach dialog", () => {
       ).toBeVisible();
     });
   }
+
+  test("keeps the reader's view when the persisted copy replaces the streamed reply", async ({
+    page,
+  }, testInfo) => {
+    test.skip(
+      testInfo.project.name !== "chromium-desktop",
+      "viewport-independent; desktop run suffices",
+    );
+    await serveAiBlock(page, aiBlockAvailable());
+    let release: () => void = () => {};
+    const holdRefetch = new Promise<void>((resolve) => {
+      release = resolve;
+    });
+    await mockCoach(page, { holdRefetch });
+    await useTheme(page, "light");
+
+    await page.goto("/coach", { waitUntil: "domcontentloaded" });
+    await page
+      .locator('[data-slot="coach-input-textarea"]')
+      .fill(FIRST_QUESTION);
+    await page.locator('[data-slot="coach-input-send"]').click();
+
+    // The streamed reply has settled; its persisted copy is held back.
+    const streamed = page.locator(
+      '[role="log"] [data-slot="coach-bubble-assistant"]',
+    );
+    await expect(streamed).toContainText(FIRST_ANSWER, { timeout: 15_000 });
+    await streamed
+      .locator('[data-slot="coach-result-view-table"]')
+      .first()
+      .click();
+    await expect(
+      streamed.locator('[data-slot="coach-result-table"][data-ref="r1"]'),
+    ).toBeVisible();
+
+    release();
+    await waitForPersistedTwin(page);
+
+    // The same bubble, still on the table the reader chose.
+    const bubble = page
+      .locator('[data-slot="coach-bubble-assistant"]')
+      .filter({ hasText: FIRST_ANSWER });
+    await expect(bubble).toHaveCount(1);
+    const table = bubble.locator(
+      '[data-slot="coach-result-table"][data-ref="r1"]',
+    );
+    await expect(table).toBeVisible();
+    await expect(
+      table.locator('[data-slot="coach-result-view-table"]'),
+    ).toHaveAttribute("aria-pressed", "true");
+  });
 });

src/components/insights/coach-panel/__tests__/live-turn-keys.test.ts

@@ -0,0 +1,66 @@
+/**
+ * v1.39.4 — a turn keeps one React key from its first streamed frame to
+ * its persisted copy, so the bubble (and what the reader did in it) survives
+ * the swap.
+ */
+import { describe, expect, it } from "vitest";
+
+import {
+  EMPTY_LIVE_TURN_KEYS,
+  messageRenderKey,
+  nextLiveTurnKeys,
+} from "../live-turn-keys";
+
+type Frame = { inProgress: boolean; messageId: string | null } | undefined;
+
+function run(frames: Frame[]) {
+  let keys = EMPTY_LIVE_TURN_KEYS;
+  const seen: Array<string | null> = [];
+  for (const frame of frames) {
+    keys = nextLiveTurnKeys(keys, frame);
+    seen.push(keys.current);
+  }
+  return { keys, seen };
+}
+
+describe("live turn keys", () => {
+  it("gives the persisted message the key its turn streamed under", () => {
+    const { keys, seen } = run([
+      undefined,
+      { inProgress: true, messageId: null },
+      { inProgress: true, messageId: null },
+      { inProgress: false, messageId: "m-7" },
+    ]);
+    expect(seen.slice(1)).toEqual([
+      "coach-turn-1",
+      "coach-turn-1",
+      "coach-turn-1",
+    ]);
+    expect(messageRenderKey(keys, "m-7")).toBe("coach-turn-1");
+    // A message that never streamed here keeps its own id.
+    expect(messageRenderKey(keys, "m-3")).toBe("m-3");
+  });
+
+  it("keeps an earlier turn's key when the next one streams", () => {
+    const { keys } = run([
+      { inProgress: true, messageId: null },
+      { inProgress: false, messageId: "m-7" },
+      { inProgress: true, messageId: null },
+      { inProgress: false, messageId: "m-9" },
+    ]);
+    expect(messageRenderKey(keys, "m-7")).toBe("coach-turn-1");
+    expect(messageRenderKey(keys, "m-9")).toBe("coach-turn-2");
+  });
+
+  it("answers the same keys for the same frame twice, so render can call it", () => {
+    const first = nextLiveTurnKeys(EMPTY_LIVE_TURN_KEYS, {
+      inProgress: true,
+      messageId: null,
+    });
+    const again = nextLiveTurnKeys(first, {
+      inProgress: true,
+      messageId: null,
+    });
+    expect(again).toBe(first);
+  });
+});

src/components/insights/coach-panel/__tests__/message-thread.test.tsx

@@ -388,6 +388,35 @@ describe("<MessageThread>", () => {
     expect(html).toContain("fresh persisted");
   });
 
+  it("renders a streaming turn and a persisted one in the same wrapper", () => {
+    // The persisted copy replaces the streamed one under the same key; for
+    // React to keep the bubble, both must sit in the same element chain.
+    const live = render(
+      <MessageThread
+        conversation={baseConversation}
+        streaming={streaming({
+          content: "Looking at your data ",
+          metricSource: null,
+          suggestion: null,
+          inProgress: true,
+          messageId: null,
+          errorCode: null,
+        })}
+      />,
+    );
+    const liveTurn = live.match(
+      /<div data-slot="coach-assistant-turn" data-turn-key="coach-turn-1"[^>]*>/,
+    )?.[0];
+    expect(liveTurn).toContain('role="log"');
+    expect(liveTurn).toContain('aria-live="polite"');
+    const persisted = render(<MessageThread conversation={baseConversation} />);
+    const turns =
+      persisted.match(/<div data-slot="coach-assistant-turn"[^>]*>/g) ?? [];
+    expect(turns.length).toBeGreaterThan(0);
+    // Settled, the wrapper is no longer a live region.
+    for (const turn of turns) expect(turn).not.toContain("aria-live");
+  });
+
   it("still renders the streaming bubble while inProgress and id is null", () => {
     // Mid-stream — `messageId` is null, persisted twin can't exist
     // yet. Confirm the streaming branch survives the v1.4.20.1

src/components/insights/coach-panel/live-turn-keys.ts

@@ -0,0 +1,62 @@
+/**
+ * The React key a Coach turn keeps from its first streamed frame to its
+ * persisted copy.
+ *
+ * A reply streams under a key minted when the turn starts; once the `done`
+ * frame names its message id, the persisted message is rendered under that
+ * same key. React then keeps one bubble through the swap, and with it what
+ * the reader did while it streamed: the chart or table view, the open steps
+ * list, the open evidence panel. Keyed by message id alone, the persisted
+ * copy mounted as a new bubble and reset all of it.
+ *
+ * Pure and idempotent for the same input, so it can run during render.
+ */
+import type { CoachStreamingMessage } from "./use-coach";
+
+export interface LiveTurnKeys {
+  /** Turns started so far in this thread view. */
+  seq: number;
+  /** The key of the turn streaming now, or the last one that streamed. */
+  current: string | null;
+  wasInProgress: boolean;
+  /** Message id → the key its turn streamed under. */
+  byMessage: ReadonlyMap<string, string>;
+}
+
+export const EMPTY_LIVE_TURN_KEYS: LiveTurnKeys = {
+  seq: 0,
+  current: null,
+  wasInProgress: false,
+  byMessage: new Map(),
+};
+
+export function nextLiveTurnKeys(
+  prev: LiveTurnKeys,
+  streaming:
+    Pick<CoachStreamingMessage, "inProgress" | "messageId"> | undefined,
+): LiveTurnKeys {
+  let next = prev;
+  const inProgress = !!streaming?.inProgress;
+  if (inProgress && !prev.wasInProgress) {
+    const seq = prev.seq + 1;
+    next = { ...next, seq, current: `coach-turn-${seq}` };
+  }
+  if (inProgress !== next.wasInProgress) {
+    next = { ...next, wasInProgress: inProgress };
+  }
+  const id = streaming?.messageId;
+  if (id && next.current && !next.byMessage.has(id)) {
+    const byMessage = new Map(next.byMessage);
+    byMessage.set(id, next.current);
+    next = { ...next, byMessage };
+  }
+  return next;
+}
+
+/** The key a message renders under: its turn's, else its own id. */
+export function messageRenderKey(
+  keys: LiveTurnKeys,
+  messageId: string,
+): string {
+  return keys.byMessage.get(messageId) ?? messageId;
+}

src/components/insights/coach-panel/message-thread.tsx

@@ -11,6 +11,12 @@ import { PlanProposalCards } from "./plan-proposal-card";
 import { ChatBubble } from "./chat-bubble";
 import { CoachFollowUpChips } from "./follow-up-chips";
 import { CoachMessageDatesProvider } from "./coach-results";
+import {
+  EMPTY_LIVE_TURN_KEYS,
+  messageRenderKey,
+  nextLiveTurnKeys,
+  type LiveTurnKeys,
+} from "./live-turn-keys";
 import type {
   CoachConversationDetailDTO,
   CoachOptimisticUserMessage,
@@ -210,6 +216,42 @@ export const COACH_SCROLLBAR = cn(
   "hover:[&::-webkit-scrollbar-thumb]:bg-[color-mix(in_srgb,var(--primary)_45%,transparent)]",
 );
 
+/**
+ * The wrapper every assistant turn renders in, streaming or persisted, so a
+ * turn keeps one element through the swap. While it streams it is the live
+ * region: role=log + aria-live=polite so screen-reader users hear the prose
+ * as tokens land (aria-relevant=text limits announcements to the streamed
+ * content; additions covers the bubble-mount edge case), with a soft
+ * fade/slide-in so the hand-off from the thinking beat reads as one motion.
+ */
+function AssistantTurn({
+  turnKey,
+  live,
+  children,
+}: {
+  turnKey: string;
+  live: boolean;
+  children: React.ReactNode;
+}) {
+  return (
+    <div
+      data-slot="coach-assistant-turn"
+      data-turn-key={turnKey}
+      {...(live
+        ? {
+            role: "log",
+            "aria-live": "polite" as const,
+            "aria-relevant": "additions text" as const,
+            className:
+              "motion-safe:animate-in motion-safe:fade-in-0 motion-safe:slide-in-from-bottom-1 motion-safe:duration-300",
+          }
+        : {})}
+    >
+      {children}
+    </div>
+  );
+}
+
 function isPinnedToBottom(el: HTMLElement, slack = 64): boolean {
   return el.scrollHeight - el.scrollTop - el.clientHeight <= slack;
 }
@@ -247,43 +289,21 @@ export function MessageThread({
   );
   // v1.16.5 — locally-rendered guided bubbles (see `placeInterleaved`).
   const interleavedItems = interleaved ?? [];
-  // v1.4.20.1 — once the SSE stream emits `done`, the route's
-  // invalidate-then-refetch pulls the persisted assistant message into
-  // `conversation.messages`. The streaming bubble was still rendering
-  // because the hook keeps `streaming.content` populated to support
-  // the in-flight render path; the result was two assistant bubbles
-  // side by side until the next `send` reset the streaming state.
-  // Suppress the streaming bubble at render time as soon as the
-  // persisted twin lands — comparing on `messageId` keeps the
-  // transition seamless while never accidentally hiding an in-flight
-  // bubble (during streaming `messageId` is null).
-  // v1.4.22 W5 reconcile (Code-MED-3) — streaming/persisted-twin
-  // race. On slow connections SSE `done` fires before the
-  // invalidate-refetch resolves; the persisted twin lands while the
-  // streaming bubble is still painted, producing a 200-500ms
-  // duplicate render. Hide the persisted twin (matched on
-  // streaming.messageId) for a 150ms grace window after it first
-  // appears so the streaming bubble stays alone until the streaming
-  // state cleans up naturally on the next `send`.
-  const [graceWindow, setGraceWindow] = useState(false);
-  const lastPersistedIdRef = useRef<string | null>(null);
-  useEffect(() => {
-    const persistedId = streaming?.messageId ?? null;
-    if (
-      persistedId &&
-      persistedId !== lastPersistedIdRef.current &&
-      messages.some((m) => m.id === persistedId)
-    ) {
-      lastPersistedIdRef.current = persistedId;
-      setGraceWindow(true);
-      const handle = setTimeout(() => setGraceWindow(false), 150);
-      return () => clearTimeout(handle);
-    }
-  }, [streaming?.messageId, messages]);
-  const suppressedTwinId = graceWindow ? streaming?.messageId : null;
+  // v1.39.4 — one key per turn from its first streamed frame to its
+  // persisted copy (`live-turn-keys.ts`). The persisted twin renders under
+  // the key the turn streamed under, in the same place and wrapper, so React
+  // keeps the one bubble and the reader's state in it. The streaming copy
+  // steps aside the moment the twin is in the history; with a shared key the
+  // two never render side by side.
+  // Derived from the previous render's keys (React's "store information
+  // from previous renders" pattern): the update is idempotent, so it settles
+  // after one extra render and never loops.
+  const [storedLiveKeys, setLiveKeys] =
+    useState<LiveTurnKeys>(EMPTY_LIVE_TURN_KEYS);
+  const liveKeys = nextLiveTurnKeys(storedLiveKeys, streaming);
+  if (liveKeys !== storedLiveKeys) setLiveKeys(liveKeys);
 
   const streamingPersisted =
-    !graceWindow &&
     streaming?.messageId != null &&
     messages.some((m) => m.id === streaming.messageId);
   const streamingActive =
@@ -422,29 +442,24 @@ export function MessageThread({
           COACH_SCROLLBAR,
         )}
       >
-        {messages.map((m, idx) => {
-          // v1.4.22 W5 reconcile (Code-MED-3) — suppress the persisted
-          // twin during the 150ms grace window so the streaming bubble
-          // stays alone on slow connections.
-          if (m.id === suppressedTwinId) return null;
-          // v1.16.5 — a guided question renders directly above the user
-          // message that answered it.
-          const guidedBefore = placement.before.get(m.id);
-          // v1.22 — "Try again": resolve the user message that produced this
-          // assistant reply (the nearest preceding user turn) so the surface
-          // can resubmit it. Null when there is none → no regenerate action.
-          let precedingUserContent: string | null = null;
-          if (m.role === "assistant" && onRegenerate) {
-            for (let j = idx - 1; j >= 0; j--) {
-              if (messages[j].role === "user") {
-                precedingUserContent = messages[j].content;
-                break;
+        {[
+          ...messages.map((m, idx) => {
+            // v1.16.5 — a guided question renders directly above the user
+            // message that answered it.
+            const guidedBefore = placement.before.get(m.id);
+            // v1.22 — "Try again": resolve the user message that produced this
+            // assistant reply (the nearest preceding user turn) so the surface
+            // can resubmit it. Null when there is none → no regenerate action.
+            let precedingUserContent: string | null = null;
+            if (m.role === "assistant" && onRegenerate) {
+              for (let j = idx - 1; j >= 0; j--) {
+                if (messages[j].role === "user") {
+                  precedingUserContent = messages[j].content;
+                  break;
+                }
               }
             }
-          }
-          return (
-            <Fragment key={m.id}>
-              {guidedBefore?.node}
+            const bubble = (
               <ChatBubble
                 role={m.role}
                 content={m.content}
@@ -455,64 +470,75 @@ export function MessageThread({
                 tokensUsed={m.tokensUsed}
                 model={m.model}
                 createdAt={m.createdAt}
+                // The turn that just streamed keeps its live tables, so the
+                // swap to the persisted copy does not refetch and remount them.
+                {...(m.id === streaming?.messageId &&
+                streaming.results.length > 0
+                  ? { results: streaming.results }
+                  : {})}
                 onRegenerate={
                   precedingUserContent !== null && onRegenerate
                     ? () => onRegenerate(precedingUserContent as string)
                     : undefined
                 }
               />
-            </Fragment>
-          );
-        })}
-        {/* v1.4.25 W5 — optimistic user bubble surfaces between the
-          persisted history and the streaming assistant placeholder so
-          the visible order matches the user's mental model. The
-          send-hook drops it as soon as the SSE `done` frame fires +
-          the invalidate-refetch lands the persisted twin. */}
-        {/* v1.16.5 — guided question whose answer is still optimistic-only. */}
-        {placement.beforeOptimistic.map((i) => (
-          <Fragment key={i.key}>{i.node}</Fragment>
-        ))}
-        {optimisticActive && optimisticUser && (
-          <ChatBubble
-            key={optimisticUser.localId}
-            role="user"
-            content={optimisticUser.content}
-          />
-        )}
-        {streamingActive && streaming && (
-          // role=log + aria-live=polite so screen-reader users hear the
-          // assistant prose announce as tokens land. aria-relevant=text
-          // limits announcements to the streamed content; additions
-          // covers the bubble-mount edge case.
-          <div
-            role="log"
-            aria-live="polite"
-            aria-relevant="additions text"
-            // v1.18.7 — a soft fade/slide-in on the assistant turn so the
-            // hand-off from the thinking beat to the first streamed tokens
-            // reads as one continuous motion, not a hard swap. v1.18.9 — the
-            // per-word fade now lives in <StreamedProse>; this stays a light
-            // container fade for the disclosure → prose transition.
-            className="motion-safe:animate-in motion-safe:fade-in-0 motion-safe:slide-in-from-bottom-1 motion-safe:duration-300"
-          >
+            );
+            const key = messageRenderKey(liveKeys, m.id);
+            return (
+              <Fragment key={key}>
+                {guidedBefore?.node}
+                {m.role === "assistant" ? (
+                  <AssistantTurn turnKey={key} live={false}>
+                    {bubble}
+                  </AssistantTurn>
+                ) : (
+                  bubble
+                )}
+              </Fragment>
+            );
+          }),
+          // v1.16.5 — guided question whose answer is still optimistic-only.
+          ...placement.beforeOptimistic.map((i) => (
+            <Fragment key={i.key}>{i.node}</Fragment>
+          )),
+          // v1.4.25 W5 — the optimistic user bubble sits between the
+          // persisted history and the streaming assistant placeholder, so
+          // the visible order matches the person's mental model. The send
+          // hook drops it once the persisted twin lands.
+          optimisticActive && optimisticUser ? (
             <ChatBubble
-              role="assistant"
-              content={streaming.content}
-              metricSource={streaming.metricSource}
-              suggestion={streaming.suggestion}
-              suggestedAction={streaming.suggestedAction}
-              steps={streaming.steps}
-              results={streaming.results}
-              providerType={streaming.inProgress ? "streaming" : null}
-              inProgress={streaming.inProgress}
-              errorCode={streaming.errorCode}
-              // v1.18.9 — live word-fade + the just-landed token footer.
-              streaming
-              usage={streaming.usage}
+              key={optimisticUser.localId}
+              role="user"
+              content={optimisticUser.content}
             />
-          </div>
-        )}
+          ) : null,
+          streamingActive && streaming ? (
+            // Same key, fragment and wrapper as its persisted copy above.
+            <Fragment key={liveKeys.current ?? "coach-turn-live"}>
+              {null}
+              <AssistantTurn
+                turnKey={liveKeys.current ?? "coach-turn-live"}
+                live
+              >
+                <ChatBubble
+                  role="assistant"
+                  content={streaming.content}
+                  metricSource={streaming.metricSource}
+                  suggestion={streaming.suggestion}
+                  suggestedAction={streaming.suggestedAction}
+                  steps={streaming.steps}
+                  results={streaming.results}
+                  providerType={streaming.inProgress ? "streaming" : null}
+                  inProgress={streaming.inProgress}
+                  errorCode={streaming.errorCode}
+                  // v1.18.9 — live word-fade + the just-landed token footer.
+                  streaming
+                  usage={streaming.usage}
+                />
+              </AssistantTurn>
+            </Fragment>
+          ) : null,
+        ]}
         {/* v1.39.4 — follow-up chips under the latest assistant turn only. */}
         {onFollowUp && chips ? (
           <CoachFollowUpChips