Keep completed chat turns completed across cold restarts

vendor/agents/.changeset/persist-chat-before-completion.md

@@ -0,0 +1,5 @@
+---
+"@cloudflare/think": patch
+---
+
+Persist assistant messages before announcing completion and retain terminal stream evidence until the next turn so a cold restart cannot continue an already-completed answer.

vendor/agents/docs/fork-diff.md

@@ -14,6 +14,32 @@ The [current audit](audit/vendor-fork-audit.md) owns measured costs, upstream
 issue findings and unresolved coverage questions. Test names below describe
 inspected coverage; fresh follow-up results are identified explicitly in the audit.
 
+## 2026-09-14 — Preserve completion across a cold restart
+
+- `packages/think/src/think.ts`: both chat streaming paths persist the assistant
+  before broadcasting completion. The cutover retains the existing terminal
+  stream record until the next stream start reclaims it. A recovery task can
+  outlive the message transaction; deleting its stream in that transaction made
+  a cold wake continue an already-completed answer. This supersedes the 0.23
+  refresh's immediate-discard behavior and its resume-specific discard policy.
+- Wire order remains completion, then canonical transcript. Sending the snapshot
+  while an observer accumulator is still active makes its terminal merge replace
+  persisted duration/status metadata. The native broadcast regression asserts
+  both durability at completion and this frame order; Rook's real Chrome stopped
+  activity-label test covers the rendered result.
+- No new storage shape, migration, timer, or recovery protocol. The existing
+  `discard: false` cutover and start-time reclaim bound retention to the previous
+  turn, as already used for agent-tool tailing and reconnect replay.
+- Reproduced while validating [Rook #166](https://github.com/WebMCP-org/rook/pull/166):
+  upgrading the published 1.1.0.509 profile, completing a second turn, and cold
+  restarting could issue a third model request and duplicate the answer.
+- `packages/think/src/tests/think-session.test.ts` observes durable assistant rows
+  at terminal broadcast on both paths and recovers a real completed turn whose
+  enclosing run survived cutover. Each regression failed before the fix.
+  `stream-cleanup.test.ts` checks that the next turn reclaims terminal evidence.
+- Why no host shim: completion and stream-to-message transactions are private
+  Think internals, shared by browser and Worker consumers. Upstreamable: yes.
+
 ## Ownership map
 
 Rook's `ThreadApp` keeps the conversation socket and Agent-tool subscription

vendor/agents/packages/think/src/tests/agents/think-session.ts

@@ -604,6 +604,46 @@ class TestCollectingCallback implements StreamCallback {
 // _transformInferenceResult (error injection).
 
 export class ThinkTestAgent extends Think {
+  private _assistantRowsAtDone: number[] = [];
+  private _completionFrames: string[] = [];
+
+  override broadcast(
+    msg: string | ArrayBuffer | ArrayBufferView,
+    without?: string[]
+  ): void {
+    if (typeof msg === "string") {
+      const frame = JSON.parse(msg) as {
+        type?: string;
+        done?: boolean;
+        messages?: UIMessage[];
+      };
+      if (
+        frame.type === "cf_agent_chat_messages" &&
+        frame.messages?.some((message) => message.role === "assistant")
+      ) {
+        this._completionFrames.push("messages");
+      }
+      if (frame.type === "cf_agent_use_chat_response" && frame.done) {
+        this._completionFrames.push("done");
+        this._assistantRowsAtDone.push(
+          this.sql<{ count: number }>`
+            SELECT COUNT(*) AS count FROM cf_agents_session_messages
+            WHERE role = 'assistant'
+          `[0].count
+        );
+      }
+    }
+    super.broadcast(msg, without);
+  }
+
+  getAssistantRowsAtDoneForTest(): number[] {
+    return this._assistantRowsAtDone;
+  }
+
+  getCompletionFramesForTest(): string[] {
+    return this._completionFrames;
+  }
+
   private _response = "Hello from the assistant!";
   private _nextSubAgentConnectionSendDelayMs = 0;
   private _chatErrorLog: string[] = [];
@@ -7737,6 +7777,7 @@ export class ThinkRecoveryTestAgent extends Think {
     await this.chat(message, cb);
     return {
       events: cb.events,
+      requestId: cb.requestId,
       done: cb.doneCalled,
       error: cb.errorMessage,
       interruptedCalls: cb.interruptedCalls

vendor/agents/packages/think/src/tests/stream-cleanup.test.ts

@@ -4,8 +4,8 @@ import { getAgentByName } from "agents";
 import type { ThinkRecoveryTestAgent } from "./agents/think-session";
 
 // Resumable-stream buffers are reclaimed without an alarm: the cutover
-// deletes a finished stream's rows in the transaction that persists its
-// message, and the next stream start reclaims anything a crash left behind —
+// settles a finished stream in the transaction that persists its message,
+// and the next stream start reclaims the terminal recovery evidence —
 // finished streams of any age and in-flight rows abandoned past the stale
 // window. Uses ThinkRecoveryTestAgent, which carries the stream test helpers.
 
@@ -107,11 +107,17 @@ describe("Think — stream reclaim (no cleanup alarm)", () => {
     expect(snapshot?.chunkCount).toBeGreaterThan(0);
   });
 
-  it("a real turn leaves no stream rows behind", async () => {
+  it("a real turn retains completion evidence until the next stream starts", async () => {
     const agent = await freshAgent();
     const result = await agent.testChat("Cut over");
     expect(result.done).toBe(true);
-    expect(await agent.getLatestStreamSnapshot()).toBeNull();
+    const completed = await agent.getLatestStreamSnapshot();
+    expect(completed?.status).toBe("completed");
+    await agent.startStreamForTest("next-turn");
+    expect(await agent.runStreamCleanupForTest()).toBe(0);
+    expect((await agent.getLatestStreamSnapshot())?.requestId).toBe(
+      "next-turn"
+    );
     expect(
       await agent.getScheduledChatRecoveryCountForTest(CLEANUP_CALLBACK)
     ).toBe(0);

vendor/agents/packages/think/src/tests/think-session.test.ts

@@ -1394,6 +1394,23 @@ describe("Think — getConfig inside configureSession", () => {
 // ── onChatResponse hook ──────────────────────────────────────────
 
 describe("Think — onChatResponse", () => {
+  it.each(["rpc", "stream"] as const)(
+    "persists the assistant before announcing %s completion",
+    async (transport) => {
+      const agent = await freshAgent(`persist-before-done-${transport}`);
+      if (transport === "rpc") {
+        await agent.runChatTurnForTest({ input: "Hello!" });
+      } else {
+        await agent.runChannelTurnForTest({ input: "Hello!" });
+      }
+      expect(await agent.getAssistantRowsAtDoneForTest()).toEqual([1]);
+      expect(await agent.getCompletionFramesForTest()).toEqual([
+        "done",
+        "messages"
+      ]);
+    }
+  );
+
   it("should fire onChatResponse after successful chat turn", async () => {
     const agent = await freshAgent("hook-success");
 
@@ -2653,18 +2670,16 @@ describe("Think — chatRecovery", () => {
     expect(fibers).toHaveLength(0);
   });
 
-  it("chat() discards the stream once its message is persisted", async () => {
+  it("chat() retains terminal stream evidence after its message is persisted", async () => {
     const agent = await freshRecoveryAgent("chat-stream-metadata");
 
     const result = await agent.testChat("Record the stream");
     expect(result.done).toBe(true);
 
-    // The stream's rows were the recovery evidence while the turn was in
-    // flight; once the assistant message is durable they are redundant and
-    // are dropped in place, leaving nothing for the retention sweep.
-    // Interrupted turns keep their rows (covered by the recovery tests).
+    // Recovery can outlive the message commit; its terminal stream must still
+    // distinguish a completed turn from an interrupted one until the next turn.
     const snapshot = await agent.getLatestStreamSnapshot();
-    expect(snapshot).toBeNull();
+    expect(snapshot?.status).toBe("completed");
     const messages = (await agent.getStoredMessages()) as UIMessage[];
     expect(messages.at(-1)?.role).toBe("assistant");
   });
@@ -4594,6 +4609,23 @@ describe("Think — onChatRecovery", () => {
     expect(assistants[0].id).toBe("a-dup");
   });
 
+  it("does not recover a real completed turn whose recovery task outlived cutover", async () => {
+    const agent = await freshRecoveryAgent("completed-cutover-recovery");
+    const result = await agent.testChat("Finish before restart");
+    expect(result.done).toBe(true);
+    expect(result.requestId).toBeTruthy();
+    // A crash after the message commit can leave the enclosing recovery run.
+    await agent.insertInterruptedFiber(
+      `__cf_internal_chat_turn:${result.requestId}`
+    );
+    expect(await agent.triggerFiberRecovery()).toEqual({
+      scheduledContinueCount: 0,
+      scheduledRetryCount: 0
+    });
+    expect(await agent.getTurnCallCount()).toBe(1);
+    expect(await agent.getStoredMessages()).toHaveLength(2);
+  });
+
   it("does not continue a recovered chat fiber whose stream already completed", async () => {
     const agent = await freshRecoveryAgent("completed-stream-recovery");