From 84c41693a0528ef8fff649c995f05c27b1cdba1b Mon Sep 17 00:00:00 2001 From: Alex Nahas Date: Mon, 14 Sep 2026 08:27:51 -0700 Subject: [PATCH 1/2] Preserve completed chat evidence until recovery can observe it A cold restart can interrupt the enclosing recovery task after its answer commits. Retain the existing terminal stream through that window and announce completion only after persistence, so waking the task cannot infer another answer. The next turn already reclaims terminal streams. --- .../persist-chat-before-completion.md | 5 + vendor/agents/docs/fork-diff.md | 21 + .../think/src/tests/agents/think-session.ts | 25 + .../think/src/tests/stream-cleanup.test.ts | 14 +- .../think/src/tests/think-session.test.ts | 40 +- vendor/agents/packages/think/src/think.ts | 708 +++++++++--------- 6 files changed, 435 insertions(+), 378 deletions(-) create mode 100644 vendor/agents/.changeset/persist-chat-before-completion.md diff --git a/vendor/agents/.changeset/persist-chat-before-completion.md b/vendor/agents/.changeset/persist-chat-before-completion.md new file mode 100644 index 0000000..aab01eb --- /dev/null +++ b/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. diff --git a/vendor/agents/docs/fork-diff.md b/vendor/agents/docs/fork-diff.md index 957fc7b..eea549a 100644 --- a/vendor/agents/docs/fork-diff.md +++ b/vendor/agents/docs/fork-diff.md @@ -14,6 +14,27 @@ 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. +- 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 diff --git a/vendor/agents/packages/think/src/tests/agents/think-session.ts b/vendor/agents/packages/think/src/tests/agents/think-session.ts index 53a24fd..72bfc3a 100644 --- a/vendor/agents/packages/think/src/tests/agents/think-session.ts +++ b/vendor/agents/packages/think/src/tests/agents/think-session.ts @@ -604,6 +604,30 @@ class TestCollectingCallback implements StreamCallback { // _transformInferenceResult (error injection). export class ThinkTestAgent extends Think { + private _assistantRowsAtDone: number[] = []; + + override broadcast( + msg: string | ArrayBuffer | ArrayBufferView, + without?: string[] + ): void { + if (typeof msg === "string") { + const frame = JSON.parse(msg) as { type?: string; done?: boolean }; + if (frame.type === "cf_agent_use_chat_response" && frame.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; + } + private _response = "Hello from the assistant!"; private _nextSubAgentConnectionSendDelayMs = 0; private _chatErrorLog: string[] = []; @@ -7737,6 +7761,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 diff --git a/vendor/agents/packages/think/src/tests/stream-cleanup.test.ts b/vendor/agents/packages/think/src/tests/stream-cleanup.test.ts index b03e292..5079343 100644 --- a/vendor/agents/packages/think/src/tests/stream-cleanup.test.ts +++ b/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); diff --git a/vendor/agents/packages/think/src/tests/think-session.test.ts b/vendor/agents/packages/think/src/tests/think-session.test.ts index e9e82a8..9ae6f7d 100644 --- a/vendor/agents/packages/think/src/tests/think-session.test.ts +++ b/vendor/agents/packages/think/src/tests/think-session.test.ts @@ -1394,6 +1394,19 @@ 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]); + } + ); + it("should fire onChatResponse after successful chat turn", async () => { const agent = await freshAgent("hook-success"); @@ -2653,18 +2666,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 +4605,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"); diff --git a/vendor/agents/packages/think/src/think.ts b/vendor/agents/packages/think/src/think.ts index d2ba685..3f41682 100644 --- a/vendor/agents/packages/think/src/think.ts +++ b/vendor/agents/packages/think/src/think.ts @@ -13448,13 +13448,6 @@ export class Think< this._finishResumableStream(streamId); } streamFinalized = true; - this._broadcastChat({ - type: MSG_CHAT_RESPONSE, - id: requestId, - body: "", - done: true - }); - doneSent = true; terminalStatus = streamError ? "error" @@ -13467,12 +13460,7 @@ export class Think< startedAt ); if (accumulator.parts.length > 0) { - await this._persistAssistantMessageWithCutover( - streamId, - assistantMsg, - undefined, - { discard: this._discardStreamAtCutover(requestId) } - ); + await this._persistAssistantMessageWithCutover(streamId, assistantMsg); // Vendor divergence: the terminal message is on the row, so the // catch/finally fallback below must not write it a second time. terminalMessagePersisted = true; @@ -13481,6 +13469,15 @@ export class Think< // Nothing to persist (or the persist threw): settle the finished stream. this._resumableStream.finalizePending(); + // Completion is observable only after the message and stream commit. + this._broadcastChat({ + type: MSG_CHAT_RESPONSE, + id: requestId, + body: "", + done: true + }); + doneSent = true; + if (terminalStatus === "error") { await this._fireResponseHook({ message: assistantMsg, @@ -13787,15 +13784,21 @@ export class Think< let streamAborted = false; let streamError: string | undefined; let output: unknown; - // Vendor divergence: the cutover's discard decision has to be taken while - // connections are still parked on the resume handshake — by the time the - // cutover runs, every terminal path has already released them. Memoized - // at the first release; see `_discardStreamAtCutover`. - let discardStream: boolean | undefined; const releaseResumeConnections = (): void => { - discardStream ??= this._discardStreamAtCutover(requestId); this._pendingResumeConnections.clear(); }; + const sendDone = (): void => { + // Reconnecting clients replay the committed buffer before terminal. + this._broadcastChat({ + type: MSG_CHAT_RESPONSE, + id: requestId, + body: "", + done: true, + ...(continuation && { continuation: true }) + }); + releaseResumeConnections(); + doneSent = true; + }; // Set when an in-stream overflow error is recoverable (opt-in): suppresses // terminal delivery so the driver can compact and re-run the turn. let overflowRetry = false; @@ -13812,368 +13815,357 @@ export class Think< // the abandoned tee branch. let streamDrainedNaturally = false; try { - this._insideInferenceLoop = true; try { - const guardedStream = iterateWithStallWatchdog( - result.toUIMessageStream({ - onError: streamErrorToString, - messageMetadata: turnMessageMetadata(startedAt) - }), - stallTimeoutMs, - () => { - this._emit("chat:stream:stalled", { - requestId, - timeoutMs: stallTimeoutMs - }); - // Tear down the upstream model stream so a hung provider/transport - // is released; the watchdog's throw drives the terminal error below. - this.abortRequest( - requestId, - new Error("chat stream stalled: inactivity watchdog fired") - ); - } - ); - for await (const chunk of guardedStream) { - if (abortSignal?.aborted) { - streamAborted = true; - break; - } - - const rawChunk = chunk as unknown as StreamChunkData; - if ( - continuationSeedParts && - isReplayChunk(continuationSeedParts, rawChunk) - ) { - continue; - } - const streamChunk = this._annotateActionApprovalChunk( - requestId, - rawChunk, - pendingActionCalls, - accumulator.parts - ); - const { action } = accumulator.applyChunk(streamChunk); - this._applyActionApprovalDescriptorToParts( - streamChunk, - accumulator.parts + this._insideInferenceLoop = true; + try { + const guardedStream = iterateWithStallWatchdog( + result.toUIMessageStream({ + onError: streamErrorToString, + messageMetadata: turnMessageMetadata(startedAt) + }), + stallTimeoutMs, + () => { + this._emit("chat:stream:stalled", { + requestId, + timeoutMs: stallTimeoutMs + }); + // Tear down the upstream model stream so a hung provider/transport + // is released; the watchdog's throw drives the terminal error below. + this.abortRequest( + requestId, + new Error("chat stream stalled: inactivity watchdog fired") + ); + } ); + for await (const chunk of guardedStream) { + if (abortSignal?.aborted) { + streamAborted = true; + break; + } - // Approved server tools execute during a continuation stream, but - // their original tool part lives in an earlier assistant message. - // The accumulator can only own this turn's new content, so it - // surfaces a terminal result for a prior message as a - // `cross-message-tool-update`. Persist + broadcast it directly so - // the approved result reaches clients and durable storage. The - // update builder is first-write-wins (replay-safe) and preserves a - // streamed `preliminary` flag; `_applyToolUpdateToMessages` skips - // the write/broadcast when the matched part is already settled. - if (action?.type === "cross-message-tool-update") { - await this._applyToolUpdateToMessages( - crossMessageToolResultUpdate( - action.toolCallId, - action.updateType, - action.output, - action.errorText, - action.preliminary - ) - ); - } - - if (action?.type === "error") { - streamError = action.error; - // Recoverable context overflow (opt-in): don't terminalize. Persist - // the partial after the loop, then signal the driver to compact and - // re-run. No `message:error`/`chat:request:failed`/error frame here. + const rawChunk = chunk as unknown as StreamChunkData; if ( - options?.overflowRecovery && - this._isRecoverableContextOverflow(streamError, requestId) + continuationSeedParts && + isReplayChunk(continuationSeedParts, rawChunk) ) { - overflowRetry = true; + continue; + } + const streamChunk = this._annotateActionApprovalChunk( + requestId, + rawChunk, + pendingActionCalls, + accumulator.parts + ); + const { action } = accumulator.applyChunk(streamChunk); + this._applyActionApprovalDescriptorToParts( + streamChunk, + accumulator.parts + ); + + // Approved server tools execute during a continuation stream, but + // their original tool part lives in an earlier assistant message. + // The accumulator can only own this turn's new content, so it + // surfaces a terminal result for a prior message as a + // `cross-message-tool-update`. Persist + broadcast it directly so + // the approved result reaches clients and durable storage. The + // update builder is first-write-wins (replay-safe) and preserves a + // streamed `preliminary` flag; `_applyToolUpdateToMessages` skips + // the write/broadcast when the matched part is already settled. + if (action?.type === "cross-message-tool-update") { + await this._applyToolUpdateToMessages( + crossMessageToolResultUpdate( + action.toolCallId, + action.updateType, + action.output, + action.errorText, + action.preliminary + ) + ); + } + + if (action?.type === "error") { + streamError = action.error; + // Recoverable context overflow (opt-in): don't terminalize. Persist + // the partial after the loop, then signal the driver to compact and + // re-run. No `message:error`/`chat:request:failed`/error frame here. + if ( + options?.overflowRecovery && + this._isRecoverableContextOverflow(streamError, requestId) + ) { + overflowRetry = true; + break; + } + if (options?.captureProgrammaticStreamError) { + this._programmaticStreamErrors.set(requestId, streamError); + } + this._emit("message:error", { error: streamError }); + // An AI-SDK error surfaces as a stream error part (not a thrown + // exception), so it lands here rather than in the `catch` below. + // Bridge it to `chat:request:failed` too — observers shouldn't have + // to know whether the failure threw or arrived as a chunk (the + // post-`beforeTurn`, in-stream provider 400 class), and turn-count + // telemetry needs the failed signal to balance `turn.started`. + this._emit("chat:request:failed", { + requestId, + stage: "stream", + messagesPersisted: true, + error: streamError + }); + this._broadcastChat({ + type: MSG_CHAT_RESPONSE, + id: requestId, + body: action.error, + done: false, + error: true, + ...(continuation && { continuation: true }) + }); break; } - if (options?.captureProgrammaticStreamError) { - this._programmaticStreamErrors.set(requestId, streamError); + + this._alignStreamStartId( + streamChunk, + action, + accumulator, + continuation + ); + if (streamChunk.type === "start" && continuationStart) { + streamChunk.continuationStart = continuationStart; } - this._emit("message:error", { error: streamError }); - // An AI-SDK error surfaces as a stream error part (not a thrown - // exception), so it lands here rather than in the `catch` below. - // Bridge it to `chat:request:failed` too — observers shouldn't have - // to know whether the failure threw or arrived as a chunk (the - // post-`beforeTurn`, in-stream provider 400 class), and turn-count - // telemetry needs the failed signal to balance `turn.started`. - this._emit("chat:request:failed", { - requestId, - stage: "stream", - messagesPersisted: true, - error: streamError - }); + + const chunkBody = JSON.stringify(streamChunk); + // Vendor divergence: keep store + broadcast synchronous — a resume + // must not replay this chunk before its live broadcast, or the + // client receives it twice (2026-09-11 "Atomic replay storage and + // live chunk delivery"). + this._storeChunkDurably( + streamId, + streamChunk, + chunkBody, + flushState + ); this._broadcastChat({ type: MSG_CHAT_RESPONSE, id: requestId, - body: action.error, + body: chunkBody, done: false, - error: true, ...(continuation && { continuation: true }) }); - break; } - - this._alignStreamStartId( - streamChunk, - action, - accumulator, - continuation + streamDrainedNaturally = !( + streamAborted || + overflowRetry || + streamError !== undefined ); - if (streamChunk.type === "start" && continuationStart) { - streamChunk.continuationStart = continuationStart; + // A cancel that lands after the final chunk (the persist/settle + // window) never hits the in-loop check above, so without this the + // turn reports "completed" to the response hook — and hosts that + // gate auto/goal continuations on that status would start the very + // turn the user just stopped. + if (!streamAborted && abortSignal?.aborted) { + streamAborted = true; + } + } finally { + this._insideInferenceLoop = false; + // Only early exits leave an abandoned tee branch; a naturally + // exhausted stream needs no drain (consumeStream is not free — it + // tees the base stream and traverses the buffered branch). A thrown + // exit (stall watchdog) never reaches the assignment above, so it + // drains too. + if (!streamDrainedNaturally) { + this._drainInferenceStream(result); } + } - const chunkBody = JSON.stringify(streamChunk); - // Vendor divergence: keep store + broadcast synchronous — a resume - // must not replay this chunk before its live broadcast, or the - // client receives it twice (2026-09-11 "Atomic replay storage and - // live chunk delivery"). - this._storeChunkDurably(streamId, streamChunk, chunkBody, flushState); + // Recoverable context overflow: discard the partial, close this stream + // segment WITHOUT a terminal frame, and hand control back to the driver + // via `onRetry`. The inline retry runs in this same invocation and owns + // the terminal outcome, so we must NOT emit a `done` frame here — and + // `doneSent = true` keeps the outer `finally` from emitting one (it would + // otherwise prematurely terminate the client's stream mid-recovery and + // mark the segment errored). + // + // The partial is intentionally NOT persisted: the driver re-runs the turn + // from scratch (`continuation: false`) against the compacted history, so + // the retry produces a fresh assistant message. Persisting the truncated + // partial would leave an orphan beside the recovered answer — and any tool + // work it captured would be re-issued by the retry, duplicating records. + // The live-streamed chunks already reached clients; the retry's + // `_broadcastMessages()` reconciles them to the real answer. + if (overflowRetry && options?.overflowRecovery) { + this._completeResumableStream(streamId); + releaseResumeConnections(); + doneSent = true; + options.overflowRecovery.onRetry(streamError); + this._streamingAssistant = null; + return { status: "aborted" }; + } + + if (streamError) { + this._errorResumableStream(streamId); + } else { + this._finishResumableStream(streamId); + } + } catch (error) { + // #1626: a stream-stall watchdog abort is a recoverable interruption, not + // a terminal error. Persist the settled partial (so the continuation + // re-anchors without re-running completed tool calls), then route into + // bounded recovery; only fall through to the terminal path below once the + // budget is exhausted. + if (error instanceof ChatStreamStalledError) { + let targetAssistantId: string | undefined; + const partialMsg = accumulator.toMessage(); + if ( + this._historyGeneration === clearGen && + accumulator.parts.length > 0 + ) { + await this._persistAssistantMessage(partialMsg, parentId); + this._broadcastMessages(); + targetAssistantId = partialMsg.id; + } + const outcome = await this._routeStallToBoundedRecovery({ + requestId, + streamId, + partialParts: partialMsg.parts, + targetAssistantId + }); + if (outcome === "scheduled") { + // Recovering: close the stream cleanly (no terminal error frame); the + // scheduled continuation drives the turn to completion. Report + // `aborted` so the caller does not terminalize the turn. + this._completeResumableStream(streamId); + if (!doneSent) { + this._broadcastChat({ + type: MSG_CHAT_RESPONSE, + id: requestId, + body: "", + done: true, + ...(continuation && { continuation: true }) + }); + doneSent = true; + } + releaseResumeConnections(); + // `aborted` (not `error`): this attempt was aborted by the watchdog; + // the scheduled continuation owns the real terminal outcome. No + // response hook fires here (the continuation fires it), mirroring how + // a deploy-interrupted attempt is superseded by its continuation. + // Plain clear (no auto-continuation re-check): recovery re-runs the + // turn and its own stream finalize re-triggers the held barrier. + this._streamingAssistant = null; + return { status: "aborted" }; + } + if (outcome === "exhausted") { + // `_routeStallToBoundedRecovery` already delivered the terminal UX + // (configured `terminalMessage` + done/error frame + `onExhausted` + + // submission interrupted), identical to deploy-recovery exhaustion. + // Finalize the stream and report `aborted` (not `error`) so the caller + // does not re-run the generic terminal path on top of it. + this._errorResumableStream(streamId); + releaseResumeConnections(); + doneSent = true; + this._streamingAssistant = null; + return { status: "aborted" }; + } + } + streamError = error instanceof Error ? error.message : "Stream error"; + if (options?.captureProgrammaticStreamError) { + this._programmaticStreamErrors.set(requestId, streamError); + } + this._errorResumableStream(streamId); + if (!doneSent) { this._broadcastChat({ type: MSG_CHAT_RESPONSE, id: requestId, - body: chunkBody, - done: false, + body: streamError, + done: true, + error: true, ...(continuation && { continuation: true }) }); + doneSent = true; } - streamDrainedNaturally = !( - streamAborted || - overflowRetry || - streamError !== undefined - ); - // A cancel that lands after the final chunk (the persist/settle - // window) never hits the in-loop check above, so without this the - // turn reports "completed" to the response hook — and hosts that - // gate auto/goal continuations on that status would start the very - // turn the user just stopped. - if (!streamAborted && abortSignal?.aborted) { - streamAborted = true; - } - } finally { - this._insideInferenceLoop = false; - // Only early exits leave an abandoned tee branch; a naturally - // exhausted stream needs no drain (consumeStream is not free — it - // tees the base stream and traverses the buffered branch). A thrown - // exit (stall watchdog) never reaches the assignment above, so it - // drains too. - if (!streamDrainedNaturally) { - this._drainInferenceStream(result); - } - } - - // Recoverable context overflow: discard the partial, close this stream - // segment WITHOUT a terminal frame, and hand control back to the driver - // via `onRetry`. The inline retry runs in this same invocation and owns - // the terminal outcome, so we must NOT emit a `done` frame here — and - // `doneSent = true` keeps the outer `finally` from emitting one (it would - // otherwise prematurely terminate the client's stream mid-recovery and - // mark the segment errored). - // - // The partial is intentionally NOT persisted: the driver re-runs the turn - // from scratch (`continuation: false`) against the compacted history, so - // the retry produces a fresh assistant message. Persisting the truncated - // partial would leave an orphan beside the recovered answer — and any tool - // work it captured would be re-issued by the retry, duplicating records. - // The live-streamed chunks already reached clients; the retry's - // `_broadcastMessages()` reconciles them to the real answer. - if (overflowRetry && options?.overflowRecovery) { - this._completeResumableStream(streamId); releaseResumeConnections(); - doneSent = true; - options.overflowRecovery.onRetry(streamError); - this._streamingAssistant = null; - return { status: "aborted" }; } - if (streamError) { - this._errorResumableStream(streamId); - } else { - this._finishResumableStream(streamId); - } - // A reconnecting client must receive replay before terminal. Keep it - // excluded until this broadcast; its ACK replays the completed buffer. - this._broadcastChat({ - type: MSG_CHAT_RESPONSE, - id: requestId, - body: "", - done: true, - ...(continuation && { continuation: true }) - }); - releaseResumeConnections(); - doneSent = true; - } catch (error) { - // #1626: a stream-stall watchdog abort is a recoverable interruption, not - // a terminal error. Persist the settled partial (so the continuation - // re-anchors without re-running completed tool calls), then route into - // bounded recovery; only fall through to the terminal path below once the - // budget is exhausted. - if (error instanceof ChatStreamStalledError) { - let targetAssistantId: string | undefined; - const partialMsg = accumulator.toMessage(); - if ( - this._historyGeneration === clearGen && - accumulator.parts.length > 0 - ) { - await this._persistAssistantMessage(partialMsg, parentId); - this._broadcastMessages(); - targetAssistantId = partialMsg.id; - } - const outcome = await this._routeStallToBoundedRecovery({ - requestId, - streamId, - partialParts: partialMsg.parts, - targetAssistantId - }); - if (outcome === "scheduled") { - // Recovering: close the stream cleanly (no terminal error frame); the - // scheduled continuation drives the turn to completion. Report - // `aborted` so the caller does not terminalize the turn. - this._completeResumableStream(streamId); - if (!doneSent) { - this._broadcastChat({ - type: MSG_CHAT_RESPONSE, - id: requestId, - body: "", - done: true, - ...(continuation && { continuation: true }) - }); - doneSent = true; + if ( + options?.captureOutput && + result.output && + !streamError && + !streamAborted + ) { + try { + output = await result.output; + } catch (error) { + streamError = + error instanceof Error ? error.message : "Structured output error"; + if (options.captureProgrammaticStreamError) { + this._programmaticStreamErrors.set(requestId, streamError); } - releaseResumeConnections(); - // `aborted` (not `error`): this attempt was aborted by the watchdog; - // the scheduled continuation owns the real terminal outcome. No - // response hook fires here (the continuation fires it), mirroring how - // a deploy-interrupted attempt is superseded by its continuation. - // Plain clear (no auto-continuation re-check): recovery re-runs the - // turn and its own stream finalize re-triggers the held barrier. - this._streamingAssistant = null; - return { status: "aborted" }; - } - if (outcome === "exhausted") { - // `_routeStallToBoundedRecovery` already delivered the terminal UX - // (configured `terminalMessage` + done/error frame + `onExhausted` + - // submission interrupted), identical to deploy-recovery exhaustion. - // Finalize the stream and report `aborted` (not `error`) so the caller - // does not re-run the generic terminal path on top of it. - this._errorResumableStream(streamId); - releaseResumeConnections(); - doneSent = true; - this._streamingAssistant = null; - return { status: "aborted" }; } } - streamError = error instanceof Error ? error.message : "Stream error"; - if (options?.captureProgrammaticStreamError) { - this._programmaticStreamErrors.set(requestId, streamError); + if (!streamAborted && abortSignal?.aborted) { + streamAborted = true; } - this._errorResumableStream(streamId); - if (!doneSent) { - this._broadcastChat({ - type: MSG_CHAT_RESPONSE, - id: requestId, - body: streamError, - done: true, - error: true, - ...(continuation && { continuation: true }) - }); - doneSent = true; + if (output !== undefined) { + this._persistAgentToolOutputForRequest(requestId, output); } - releaseResumeConnections(); - } finally { - if (!doneSent) { - this._errorResumableStream(streamId); - this._broadcastChat({ - type: MSG_CHAT_RESPONSE, - id: requestId, - body: "", - done: true, - ...(continuation && { continuation: true }) - }); - releaseResumeConnections(); - } - } - if ( - options?.captureOutput && - result.output && - !streamError && - !streamAborted - ) { - try { - output = await result.output; - } catch (error) { - streamError = - error instanceof Error ? error.message : "Structured output error"; - if (options.captureProgrammaticStreamError) { - this._programmaticStreamErrors.set(requestId, streamError); + if (this._historyGeneration === clearGen) { + try { + const status: ThinkTerminalMessageStatus = streamError + ? "error" + : streamAborted + ? "aborted" + : "completed"; + const assistantMsg = withTerminalMessageTiming( + accumulator.toMessage(), + status, + startedAt + ); + + if (accumulator.parts.length > 0) { + await this._persistAssistantMessageWithCutover( + streamId, + assistantMsg, + parentId + ); + this._broadcastMessages(); + } + // Nothing to persist (or the persist threw): settle the finished + // stream so it is not mistaken for an interrupted turn. + this._resumableStream.finalizePending(); + if (!doneSent) sendDone(); + + await this._fireResponseHook({ + message: assistantMsg, + requestId, + continuation, + status, + error: streamError + }); + } catch (e) { + console.error("Failed to persist assistant message:", e); } } - } - if (!streamAborted && abortSignal?.aborted) { - streamAborted = true; - } - if (output !== undefined) { - this._persistAgentToolOutputForRequest(requestId, output); - } + this._resumableStream.finalizePending(); - if (this._historyGeneration === clearGen) { - try { - const status: ThinkTerminalMessageStatus = streamError - ? "error" - : streamAborted - ? "aborted" - : "completed"; - const assistantMsg = withTerminalMessageTiming( - accumulator.toMessage(), - status, - startedAt - ); + if (!doneSent) sendDone(); - if (accumulator.parts.length > 0) { - await this._persistAssistantMessageWithCutover( - streamId, - assistantMsg, - parentId, - { - discard: discardStream ?? this._discardStreamAtCutover(requestId) - } - ); - this._broadcastMessages(); - } - // Nothing to persist (or the persist threw): settle the finished - // stream so it is not mistaken for an interrupted turn. - this._resumableStream.finalizePending(); + // The message is now persisted (or the turn was cleared), so subsequent + // tool results resolve against storage; stop exposing the accumulator and + // re-check any continuation the stream-active barrier held (#1650). + this._onStreamingTurnFinalized(); - await this._fireResponseHook({ - message: assistantMsg, - requestId, - continuation, - status, - error: streamError - }); - } catch (e) { - console.error("Failed to persist assistant message:", e); + return streamError + ? { status: "error", error: streamError } + : { + status: streamAborted ? "aborted" : "completed", + ...(output !== undefined && { output }) + }; + } finally { + if (!doneSent) { + this._errorResumableStream(streamId); + sendDone(); } } - this._resumableStream.finalizePending(); - - // The message is now persisted (or the turn was cleared), so subsequent - // tool results resolve against storage; stop exposing the accumulator and - // re-check any continuation the stream-active barrier held (#1650). - this._onStreamingTurnFinalized(); - - return streamError - ? { status: "error", error: streamError } - : { - status: streamAborted ? "aborted" : "completed", - ...(output !== undefined && { output }) - }; } // ── Session-backed persistence ────────────────────────────────── @@ -14208,16 +14200,14 @@ export class Think< /** * The cutover: persist the finished turn's assistant message, settle its - * resumable stream and delete the stream's rows in ONE SQLite transaction, - * so a crash leaves either the live stream (recovery rebuilds the message - * from it) or the message — never neither, never both. The session + * resumable stream in ONE SQLite transaction, so a crash leaves either the + * live stream or the message with terminal recovery evidence. The session * change feed and auto-compaction run once the transaction has committed. */ private async _persistAssistantMessageWithCutover( streamId: string, msg: UIMessage, - parentId?: string, - options: { discard?: boolean } = {} + parentId?: string ): Promise { const toPersist = this._strippedForPersist(msg); if (toPersist === null) return; @@ -14254,7 +14244,10 @@ export class Think< source: "server" }).after; }, - { discard: options.discard ?? true } + // The enclosing recovery task may still be running when this message + // commits. Keep its terminal evidence until the next stream starts; + // deleting it here makes a cold wake continue an already-finished turn. + { discard: false } ); } catch (error) { // The settle transaction rolled back: the row never landed, but the @@ -14265,27 +14258,6 @@ export class Think< await after?.(); } - /** - * Whether this turn's stream rows can go with its cutover. An agent-tool - * child turn keeps them: the parent tails the stored chunks after the - * child completes (`getAgentToolChunks`), so the rows are reclaimed by - * the child's next `start()` instead, as `AIChatAgent` does. - * - * Vendor divergence: a turn with connections parked on the resume handshake - * keeps them too. Those clients were deliberately excluded from the live - * terminal broadcast so their queued ACK replays the completed buffer and - * its terminal frame (2026-08-11 "Reconnecting streams receive replay - * before terminal"); the ACK is processed AFTER this cutover, so discarding - * the rows here leaves it nothing to replay — the client would receive a - * bare `done` with none of the turn's content and, for a continuation, no - * `continuation: true` (2026-08-30 "Think retains continuation metadata on - * resume"). Call this BEFORE `_pendingResumeConnections.clear()`. - */ - private _discardStreamAtCutover(requestId: string): boolean { - if (this._agentToolRunsByRequestId.get(requestId)) return false; - return this._pendingResumeConnections.size === 0; - } - /** * Remove parts belonging to Think's internal structured-output final-answer * tool (`think_final_answer`, or a collision-suffixed variant) from a UI From 03bed92b56fc922ccd43ecc008d20162336ff742 Mon Sep 17 00:00:00 2001 From: Alex Nahas Date: Mon, 14 Sep 2026 08:47:25 -0700 Subject: [PATCH 2/2] Send canonical terminal metadata after closing the observer stream Persistence must precede completion, but the observer accumulator still needs done before the canonical snapshot. Keep that wire order so stopping a turn preserves its elapsed label instead of replacing metadata with the final accumulator merge. --- vendor/agents/docs/fork-diff.md | 5 +++++ .../think/src/tests/agents/think-session.ts | 18 +++++++++++++++++- .../think/src/tests/think-session.test.ts | 4 ++++ vendor/agents/packages/think/src/think.ts | 6 ++++-- 4 files changed, 30 insertions(+), 3 deletions(-) diff --git a/vendor/agents/docs/fork-diff.md b/vendor/agents/docs/fork-diff.md index eea549a..83ec358 100644 --- a/vendor/agents/docs/fork-diff.md +++ b/vendor/agents/docs/fork-diff.md @@ -22,6 +22,11 @@ inspected coverage; fresh follow-up results are identified explicitly in the aud 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. diff --git a/vendor/agents/packages/think/src/tests/agents/think-session.ts b/vendor/agents/packages/think/src/tests/agents/think-session.ts index 72bfc3a..361ffe3 100644 --- a/vendor/agents/packages/think/src/tests/agents/think-session.ts +++ b/vendor/agents/packages/think/src/tests/agents/think-session.ts @@ -605,14 +605,26 @@ class TestCollectingCallback implements StreamCallback { 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 }; + 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 @@ -628,6 +640,10 @@ export class ThinkTestAgent extends Think { return this._assistantRowsAtDone; } + getCompletionFramesForTest(): string[] { + return this._completionFrames; + } + private _response = "Hello from the assistant!"; private _nextSubAgentConnectionSendDelayMs = 0; private _chatErrorLog: string[] = []; diff --git a/vendor/agents/packages/think/src/tests/think-session.test.ts b/vendor/agents/packages/think/src/tests/think-session.test.ts index 9ae6f7d..45c3853 100644 --- a/vendor/agents/packages/think/src/tests/think-session.test.ts +++ b/vendor/agents/packages/think/src/tests/think-session.test.ts @@ -1404,6 +1404,10 @@ describe("Think — onChatResponse", () => { await agent.runChannelTurnForTest({ input: "Hello!" }); } expect(await agent.getAssistantRowsAtDoneForTest()).toEqual([1]); + expect(await agent.getCompletionFramesForTest()).toEqual([ + "done", + "messages" + ]); } ); diff --git a/vendor/agents/packages/think/src/think.ts b/vendor/agents/packages/think/src/think.ts index 3f41682..9f24373 100644 --- a/vendor/agents/packages/think/src/think.ts +++ b/vendor/agents/packages/think/src/think.ts @@ -13464,7 +13464,6 @@ export class Think< // Vendor divergence: the terminal message is on the row, so the // catch/finally fallback below must not write it a second time. terminalMessagePersisted = true; - this._broadcastMessages(); } // Nothing to persist (or the persist threw): settle the finished stream. this._resumableStream.finalizePending(); @@ -13477,6 +13476,9 @@ export class Think< done: true }); doneSent = true; + // The observer accumulator clears at done; publish terminal metadata + // afterwards so its final merge cannot replace the persisted snapshot. + if (terminalMessagePersisted) this._broadcastMessages(); if (terminalStatus === "error") { await this._fireResponseHook({ @@ -14127,12 +14129,12 @@ export class Think< assistantMsg, parentId ); - this._broadcastMessages(); } // Nothing to persist (or the persist threw): settle the finished // stream so it is not mistaken for an interrupted turn. this._resumableStream.finalizePending(); if (!doneSent) sendDone(); + if (accumulator.parts.length > 0) this._broadcastMessages(); await this._fireResponseHook({ message: assistantMsg,