fix(agents): let the server snapshot heal a diverged observed assistant (#2166)

.changeset/observer-snapshot-heal.md

@@ -0,0 +1,5 @@
+---
+"agents": patch
+---
+
+`useAgentChat` now lets the server's message snapshot replace an observed (cross-tab or resumed) assistant message when the live copy's text no longer extends the server's copy. An interleaved or duplicated observed stream previously overrode every clean snapshot, including the final one, so the scrambled text stayed until a page reload.

packages/agents/src/chat/__tests__/broadcast-state.test.ts

@@ -1,5 +1,5 @@
 import { describe, it, expect } from "vitest";
-import { transition } from "../broadcast-state";
+import { observedDivergesFrom, transition } from "../broadcast-state";
 import type { BroadcastStreamState } from "../broadcast-state";
 import type { UIMessage } from "ai";
 
@@ -587,6 +587,88 @@ describe("broadcast stream state machine", () => {
     expect(fullText).toContain(" second half");
   });
 
+  // ── diverged accumulator vs stored copy (#2166) ──────────────────
+
+  function observe(deltas: string[]): BroadcastStreamState {
+    let state: BroadcastStreamState = idle;
+    for (const chunkData of [
+      { type: "start", messageId: "a1" },
+      { type: "text-start", id: "t1" },
+      ...deltas.map((delta) => ({ type: "text-delta", id: "t1", delta }))
+    ]) {
+      state = transition(state, {
+        type: "response",
+        streamId: "s-heal",
+        messageId: "tmp",
+        chunkData
+      }).state;
+    }
+    return state;
+  }
+
+  function stored(text: string): UIMessage[] {
+    return [
+      { id: "u1", role: "user", parts: [{ type: "text", text: "hi" }] },
+      { id: "a1", role: "assistant", parts: [{ type: "text", text }] }
+    ] as UIMessage[];
+  }
+
+  function assistantText(messages: UIMessage[]): string {
+    return messages
+      .filter((m) => m.role === "assistant")
+      .flatMap((m) => m.parts)
+      .filter((p) => p.type === "text")
+      .map((p) => (p as { text: string }).text)
+      .join("");
+  }
+
+  it("observedDivergesFrom flags an accumulator that no longer extends the stored copy", () => {
+    const scrambled = observe(["You", " can", "You", " can"]);
+    const healthy = observe(["You", " can", " see"]);
+    if (scrambled.status !== "observing" || healthy.status !== "observing") {
+      throw new Error("expected observing");
+    }
+    expect(observedDivergesFrom(scrambled.accumulator, stored("You can"))).toBe(
+      false
+    );
+    expect(
+      observedDivergesFrom(scrambled.accumulator, stored("You can see"))
+    ).toBe(true);
+    expect(observedDivergesFrom(healthy.accumulator, stored("You can"))).toBe(
+      false
+    );
+    expect(observedDivergesFrom(healthy.accumulator, stored(""))).toBe(false);
+    expect(observedDivergesFrom(healthy.accumulator, makeMessages("hi"))).toBe(
+      false
+    );
+  });
+
+  it("done keeps a stored copy the accumulator diverged from", () => {
+    const state = observe(["You", " canYou", " see can", " see"]);
+    const done = transition(state, {
+      type: "response",
+      streamId: "s-heal",
+      messageId: "tmp",
+      done: true
+    });
+    expect(assistantText(done.messagesUpdate!(stored("You can see")))).toBe(
+      "You can see"
+    );
+  });
+
+  it("done still merges an accumulator that extends the stored copy", () => {
+    const state = observe(["You", " can", " see"]);
+    const done = transition(state, {
+      type: "response",
+      streamId: "s-heal",
+      messageId: "tmp",
+      done: true
+    });
+    expect(assistantText(done.messagesUpdate!(stored("You can")))).toBe(
+      "You can see"
+    );
+  });
+
   // ── resume-fallback then response chunks ─────────────────────────
 
   it("response chunks reuse accumulator created by resume-fallback", () => {

packages/agents/src/chat/broadcast-state.ts

@@ -58,6 +58,31 @@ export interface TransitionResult {
   isStreaming: boolean;
 }
 
+// ── Snapshot reconciliation ────────────────────────────────────────
+
+function textOf(parts: UIMessage["parts"]): string {
+  let text = "";
+  for (const part of parts) {
+    if (part.type === "text") text += part.text;
+  }
+  return text;
+}
+
+/**
+ * Whether `messages` already holds a copy of the observed message whose text
+ * the accumulator does not extend. A healthy live accumulator is always the
+ * stored copy plus more; one that diverges holds interleaved or duplicated
+ * chunks (#2166), so the stored copy must win or the corruption never heals.
+ */
+export function observedDivergesFrom(
+  accumulator: StreamAccumulator,
+  messages: UIMessage[]
+): boolean {
+  const existing = messages.find((m) => m.id === accumulator.messageId);
+  if (!existing) return false;
+  return !textOf(accumulator.parts).startsWith(textOf(existing.parts));
+}
+
 // ── Transition ─────────────────────────────────────────────────────
 
 export function transition(
@@ -118,7 +143,10 @@ export function transition(
       let messagesUpdate: ((prev: UIMessage[]) => UIMessage[]) | undefined;
 
       if (event.done) {
-        messagesUpdate = (prev) => accumulator.mergeInto(prev);
+        messagesUpdate = (prev) =>
+          observedDivergesFrom(accumulator, prev)
+            ? prev
+            : accumulator.mergeInto(prev);
         return {
           state: { status: "idle" },
           messagesUpdate,

packages/agents/src/chat/react.tsx

@@ -14,6 +14,7 @@ import type { OutgoingMessage } from "./wire-types";
 import { STREAM_RESUME_NONE_REASONS } from "./protocol";
 import { MessageType } from "./wire-types";
 import {
+  observedDivergesFrom,
   transition as broadcastTransition,
   type BroadcastStreamState
 } from "./broadcast-state";
@@ -1925,13 +1926,17 @@ export function useAgentChat<
               // briefly trail a fully-persisted snapshot; merging then would drop
               // parts until replay catches up. In steady-state live observing the
               // accumulator is always at or ahead of the snapshot, so this still
-              // fixes the disappear/reappear flicker.
+              // fixes the disappear/reappear flicker. An accumulator whose text
+              // no longer extends the snapshot's is corrupt, not ahead (#2166).
               const snapshotIdx = next.findIndex(
                 (m) => m.id === observed.accumulator.messageId
               );
               const snapshotParts =
                 snapshotIdx >= 0 ? next[snapshotIdx].parts.length : 0;
-              if (observed.accumulator.parts.length >= snapshotParts) {
+              if (
+                observed.accumulator.parts.length >= snapshotParts &&
+                !observedDivergesFrom(observed.accumulator, next)
+              ) {
                 next = observed.accumulator.mergeInto(next) as ChatMessage[];
               }
             }

packages/ai-chat/src/react-tests/use-agent-chat.test.tsx

@@ -6117,6 +6117,93 @@ describe("useAgentChat overlapping submits (issue #1231)", () => {
       .toHaveTextContent("One two");
   });
 
+  it("lets the server snapshot replace an OBSERVED assistant whose text diverged from it (#2166)", async () => {
+    const { agent, target } = createAgentWithTarget({
+      name: "observer-heal",
+      url: "ws://localhost:3000/agents/chat/observer-heal?_pk=abc"
+    });
+
+    const TestComponent = () => {
+      const chat = useAgentChat({
+        agent,
+        getInitialMessages: null,
+        messages: [] as UIMessage[],
+        resume: false
+      });
+      const assistant = chat.messages.find((m) => m.role === "assistant");
+      const text = (assistant?.parts ?? [])
+        .filter((p) => p.type === "text")
+        .map((p) => (p as { text: string }).text)
+        .join("");
+      return <div data-testid="assistant-text">{text}</div>;
+    };
+
+    const screen = await act(async () => {
+      const screen = render(<TestComponent />, {
+        wrapper: ({ children }) => (
+          <StrictMode>
+            <Suspense fallback="Loading...">{children}</Suspense>
+          </StrictMode>
+        )
+      });
+      await sleep(10);
+      return screen;
+    });
+
+    // Two copies of the reply interleave into the observed text part.
+    const chunks = [
+      '{"type":"start","messageId":"obs-assistant"}',
+      '{"type":"text-start","id":"t1"}',
+      ...["You", " can", "You", " see", " can", " see"].map(
+        (delta) => `{"type":"text-delta","id":"t1","delta":"${delta}"}`
+      )
+    ];
+    await act(async () => {
+      for (const body of chunks) {
+        dispatch(target, {
+          type: "cf_agent_use_chat_response",
+          id: "req-scrambled",
+          body,
+          done: false
+        });
+      }
+      await sleep(10);
+    });
+
+    await expect
+      .element(screen.getByTestId("assistant-text"))
+      .toHaveTextContent("You canYou see can see");
+
+    // The server persists the clean reply and broadcasts it before `done`.
+    await act(async () => {
+      dispatch(target, {
+        type: "cf_agent_chat_messages",
+        messages: [
+          { id: "u1", role: "user", parts: [{ type: "text", text: "Hi" }] },
+          {
+            id: "obs-assistant",
+            role: "assistant",
+            parts: [{ type: "text", text: "You can see" }]
+          }
+        ]
+      });
+      dispatch(target, {
+        type: "cf_agent_use_chat_response",
+        id: "req-scrambled",
+        body: "",
+        done: true
+      });
+      await sleep(10);
+    });
+
+    await expect
+      .element(screen.getByTestId("assistant-text"))
+      .toHaveTextContent("You can see");
+    expect(screen.getByTestId("assistant-text").element().textContent).toBe(
+      "You can see"
+    );
+  });
+
   it("clears protection when CF_AGENT_CHAT_CLEAR arrives mid-stream", async () => {
     const { agent, target, sentMessages } = createAgentWithTarget({
       name: "clear-mid-stream",