fix: report continuation turns that fail before streaming (#2381)

.changeset/continuation-prestream-errors.md

@@ -0,0 +1,6 @@
+---
+"@cloudflare/think": patch
+"@cloudflare/ai-chat": patch
+---
+
+Report continuation turns that fail before they stream (#2381). When an auto-continuation (after a tool approval or a client tool result) or a connection-less continuation throws before producing a response, for example from `beforeTurn` or `onChatMessage`, clients now receive an error frame for the continuation and `onChatResponse` fires with `status: "error"` and `continuation: true`. The failure is also recorded as the terminal status, so a client that reconnects learns the turn failed. Think additionally calls `onChatError` with `stage: "turn"` and a new `continuation: true` field on `ChatErrorContext`, and emits `chat:request:failed`. Previously the failure was only logged, and the conversation looked like it was still waiting.

packages/ai-chat/src/index.ts

@@ -3462,6 +3462,7 @@ export class AIChatAgent<
               },
               async () => {
                 const autoContinuationBody = async () => {
+                  let replying = false;
                   try {
                     await this._repairInterruptedToolsBeforeTurn({
                       continuation: true
@@ -3484,7 +3485,10 @@ export class AIChatAgent<
                         [],
                         {
                           continuation: true,
-                          chatMessageId: requestId
+                          chatMessageId: requestId,
+                          onStreamStart: () => {
+                            replying = true;
+                          }
                         }
                       );
                       if (replyResult.status === "error") {
@@ -3496,6 +3500,15 @@ export class AIChatAgent<
                       this._clearPendingAutoContinuation(true);
                       this._activateDeferredAutoContinuation();
                     }
+                  } catch (error) {
+                    if (
+                      !replying &&
+                      !abortSignal?.aborted &&
+                      !isDurableObjectResetError(error)
+                    ) {
+                      await this._reportContinuationFailure(requestId, error);
+                    }
+                    throw error;
                   } finally {
                     this._abortRegistry.remove(requestId);
                   }
@@ -3526,6 +3539,42 @@ export class AIChatAgent<
     });
   }
 
+  /**
+   * Report an auto-continuation that failed before it produced a response,
+   * the way a stream error is reported: an error frame, a durable terminal
+   * record, and `onChatResponse` with `status: "error"` for the assistant
+   * message the continuation would have extended. Called inside the turn so
+   * the response hook drains when the turn settles.
+   */
+  private async _reportContinuationFailure(
+    requestId: string,
+    error: unknown
+  ): Promise<void> {
+    const errorMessage = error instanceof Error ? error.message : String(error);
+    this._broadcastChatMessage({
+      body: errorMessage,
+      done: true,
+      error: true,
+      id: requestId,
+      type: MessageType.CF_AGENT_USE_CHAT_RESPONSE,
+      continuation: true
+    });
+    const message = [...this.messages]
+      .reverse()
+      .find((candidate) => candidate.role === "assistant");
+    if (message) {
+      this._pendingChatResponseResults.push({
+        message,
+        requestId,
+        continuation: true,
+        status: "error",
+        error: errorMessage
+      });
+    } else {
+      await this._recordChatTerminal(requestId, errorMessage);
+    }
+  }
+
   /**
    * @returns Terminal status for the turn.
    */
@@ -7527,9 +7576,17 @@ export class AIChatAgent<
     id: string,
     response: Response,
     excludeBroadcastIds: string[] = [],
-    options: { continuation?: boolean; chatMessageId?: string } = {}
+    options: {
+      continuation?: boolean;
+      chatMessageId?: string;
+      /**
+       * Called once the stream is set up. From then on `_reply` reports its
+       * own failures to clients; a throw before it is the caller's to report.
+       */
+      onStreamStart?: () => void;
+    } = {}
   ): Promise<StreamResultStatus> {
-    const { continuation = false, chatMessageId } = options;
+    const { continuation = false, chatMessageId, onStreamStart } = options;
     // Look up the abort signal for this request so we can cancel the reader
     // loop if the client sends a cancel message. This is a safety net —
     // users should also pass abortSignal to streamText for proper cancellation.
@@ -7554,28 +7611,37 @@ export class AIChatAgent<
           return { status: "completed" };
         }
 
-        // Parsing state adapted from:
-        // https://github.com/vercel/ai/blob/main/packages/ai/src/ui-message-stream/ui-message-chunks.ts#L295
-        const message = this._createStreamingAssistantMessage(continuation);
-
-        // Start tracking this stream for resumability. The allocated message id
-        // is persisted in stream metadata so orphan recovery (#1691) can
-        // re-associate reconstructed chunks with the right assistant message —
-        // even when the provider stream carries no `start.messageId`. For a
-        // continuation this is the cloned last-assistant id, so recovery merges
-        // into it; for a new turn it is a fresh id, so recovery keeps it
-        // distinct.
-        // The continuation flag is persisted in stream metadata so replayed
-        // frames carry `continuation: true` exactly like the live broadcast
-        // frames below (#1733) — a reconnecting client needs it to append to
-        // the existing assistant message instead of rebuilding it from
-        // scratch and dropping the pre-continuation parts.
-        const streamId = this._startStream(id, {
-          messageId: message.id,
-          continuation
-        });
-
+        // Take the reader before the stream starts: an unreadable body must
+        // fail while nothing has been sent, not leave a started stream behind.
         const reader = response.body.getReader();
+        let message: UIMessage;
+        let streamId: string;
+        try {
+          // Parsing state adapted from:
+          // https://github.com/vercel/ai/blob/main/packages/ai/src/ui-message-stream/ui-message-chunks.ts#L295
+          message = this._createStreamingAssistantMessage(continuation);
+
+          // Start tracking this stream for resumability. The allocated message id
+          // is persisted in stream metadata so orphan recovery (#1691) can
+          // re-associate reconstructed chunks with the right assistant message —
+          // even when the provider stream carries no `start.messageId`. For a
+          // continuation this is the cloned last-assistant id, so recovery merges
+          // into it; for a new turn it is a fresh id, so recovery keeps it
+          // distinct.
+          // The continuation flag is persisted in stream metadata so replayed
+          // frames carry `continuation: true` exactly like the live broadcast
+          // frames below (#1733) — a reconnecting client needs it to append to
+          // the existing assistant message instead of rebuilding it from
+          // scratch and dropping the pre-continuation parts.
+          streamId = this._startStream(id, {
+            messageId: message.id,
+            continuation
+          });
+        } catch (error) {
+          // The streaming `finally` that releases the reader is not entered yet.
+          reader.cancel(error).catch(() => {});
+          throw error;
+        }
 
         // Track the streaming message so tool results can be applied before persistence
         this._streamingMessage = message;
@@ -7607,6 +7673,7 @@ export class AIChatAgent<
         this._streamingTurnActive = true;
         this._heldTerminalFrames.set(id, []);
         let persisted = false;
+        onStreamStart?.();
         try {
           try {
             if (isSSE) {

packages/ai-chat/src/tests/on-stream-end.test.ts

@@ -382,6 +382,83 @@ describe("onChatResponse with continuation", () => {
 
     ws.close(1000);
   });
+
+  it.each([
+    ["onChatMessage throws", "throw" as const],
+    ["the response body cannot be read", "locked-body" as const]
+  ])(
+    "reports an auto-continuation that fails before streaming when %s (#2381)",
+    async (_case, failure) => {
+      const room = crypto.randomUUID();
+      const { ws } = await connectContinuationAgent(room);
+
+      const initialDone = waitForDone(ws, "req-init");
+      sendChatRequest(ws, "req-init", [userMessage]);
+      expect(await initialDone).toBe(true);
+
+      const agentStub = await getAgentByName(
+        env.ResponseContinuationAgent,
+        room
+      );
+      await agentStub.persistMessages([
+        ...((await agentStub.getPersistedMessages()) as ChatMessage[]),
+        {
+          id: "assistant-tool",
+          role: "assistant",
+          parts: [
+            {
+              type: "tool-testTool",
+              toolCallId: "call_cont_fail",
+              state: "input-available",
+              input: { query: "test" }
+            }
+          ] as ChatMessage["parts"]
+        }
+      ]);
+      await agentStub.setFailContinuation(failure);
+      const expectedError =
+        failure === "throw"
+          ? "continuation failed before streaming"
+          : expect.any(String);
+
+      const terminal = new Promise<Record<string, unknown>>((resolve) => {
+        ws.addEventListener("message", function handler(e: MessageEvent) {
+          const data = JSON.parse(e.data as string);
+          if (isUseChatResponseMessage(data) && data.done) {
+            ws.removeEventListener("message", handler);
+            resolve(data);
+          }
+        });
+      });
+      ws.send(
+        JSON.stringify({
+          type: "cf_agent_tool_result",
+          toolCallId: "call_cont_fail",
+          toolName: "testTool",
+          output: { result: "success" },
+          autoContinue: true
+        })
+      );
+      expect(await terminal).toMatchObject({
+        error: true,
+        continuation: true,
+        body: expectedError
+      });
+      await new Promise((r) => setTimeout(r, 300));
+
+      const results =
+        (await agentStub.getChatResponseResults()) as ChatResponseResult[];
+      expect(results).toHaveLength(2);
+      expect(results[1]).toMatchObject({
+        continuation: true,
+        status: "error",
+        error: expectedError,
+        message: { id: "assistant-tool" }
+      });
+
+      ws.close(1000);
+    }
+  );
 });
 
 describe("onChatResponse error resilience", () => {

packages/ai-chat/src/tests/worker.ts

@@ -1871,11 +1871,22 @@ export class ResponseAgent extends AIChatAgent<Env> {
  */
 export class ResponseContinuationAgent extends AIChatAgent<Env> {
   private _responseResults: ChatResponseResult[] = [];
+  private _failContinuation: false | "throw" | "locked-body" = false;
 
   async onChatMessage(
     _onFinish: GenerateTextOnFinishCallback<ToolSet>,
-    _options?: OnChatMessageOptions
+    options?: OnChatMessageOptions
   ) {
+    if (options?.continuation && this._failContinuation === "throw") {
+      throw new Error("continuation failed before streaming");
+    }
+    if (options?.continuation && this._failContinuation === "locked-body") {
+      const response = new Response("unreadable", {
+        headers: { "Content-Type": "text/plain" }
+      });
+      response.body?.getReader();
+      return response;
+    }
     return new Response("Continuation response", {
       headers: { "Content-Type": "text/plain" }
     });
@@ -1889,6 +1900,10 @@ export class ResponseContinuationAgent extends AIChatAgent<Env> {
     return [...this._responseResults];
   }
 
+  setFailContinuation(value: false | "throw" | "locked-body"): void {
+    this._failContinuation = value;
+  }
+
   getPersistedMessages(): Promise<ChatMessage[]> {
     return persistedMessages(this);
   }

packages/think/src/tests/actions-durable-pause.test.ts

@@ -344,6 +344,29 @@ describe("durable-pause actions (turn-driven, connection-less)", () => {
     expect(await agent.listActionPendingForTest()).toHaveLength(0);
   });
 
+  it("reports a connection-less continuation that fails before streaming (#2381)", async () => {
+    const agent = await freshPauseAgent(`dp-fail-${crypto.randomUUID()}`);
+    await agent.useDurablePauseActionForTest();
+
+    const first = await agent.testChat("call pauseAction");
+    expect(first.done).toBe(true);
+    const [pending] = await agent.listActionPendingForTest();
+
+    await agent.failContinuationBeforeStreamForTest();
+    await agent.approveExecutionForTest(pending.execution_id);
+
+    await vi.waitFor(
+      async () => {
+        expect((await agent.getResponseStatusesForTest()).at(-1)).toEqual({
+          status: "error",
+          continuation: true,
+          error: "continuation failed before streaming"
+        });
+      },
+      { timeout: 5000, interval: 50 }
+    );
+  });
+
   it("labels orphaned durable-pause outcomes without re-invoking the action", async () => {
     const agent = await freshPauseAgent(`dp-orphan-${crypto.randomUUID()}`);
     await agent.useDurablePauseActionForTest();

packages/think/src/tests/agents/client-tools.ts

@@ -10,6 +10,7 @@ import { tool } from "ai";
 import { z } from "zod";
 import { Think } from "../../think";
 import type {
+  ChatErrorContext,
   ChatResponseResult,
   MessageConcurrency,
   StreamCallback,
@@ -823,9 +824,27 @@ export class ThinkClientToolsAgent extends Think {
   private _responseLog: ChatResponseResult[] = [];
   private _lastTurnToolNames: string[] = [];
   private _stampMetadata = false;
+  private _failContinuationBeforeStream = false;
+  private _chatErrorLog: ChatErrorContext[] = [];
+
+  override onChatError(error: unknown, ctx: ChatErrorContext): unknown {
+    this._chatErrorLog.push(ctx);
+    return error;
+  }
+
+  async setFailContinuationBeforeStream(value: boolean): Promise<void> {
+    this._failContinuationBeforeStream = value;
+  }
+
+  async getChatErrorLog(): Promise<ChatErrorContext[]> {
+    return this._chatErrorLog;
+  }
 
   override beforeTurn(ctx: TurnContext): TurnConfig | void {
     this._lastTurnToolNames = Object.keys(ctx.tools);
+    if (ctx.continuation && this._failContinuationBeforeStream) {
+      throw new Error("continuation failed before streaming");
+    }
     if (this._stampMetadata) {
       // Per-turn write path (`TurnConfig.messageMetadata`) for issue #1873.
       // `createdAt` on `start` and `source` on `finish` prove the two are
@@ -866,8 +885,24 @@ export class ThinkClientToolsAgent extends Think {
     return this._lastTurnToolNames;
   }
 
-  override onChatResponse(result: ChatResponseResult): void {
+  private _responseHookGate: Promise<void> | undefined;
+  private _releaseResponseHook: (() => void) | undefined;
+
+  /** While set, `onChatResponse` does not return until it is cleared. */
+  async setStallResponseHook(value: boolean): Promise<void> {
+    if (value) {
+      this._responseHookGate = new Promise((resolve) => {
+        this._releaseResponseHook = resolve;
+      });
+      return;
+    }
+    this._releaseResponseHook?.();
+    this._responseHookGate = undefined;
+  }
+
+  override async onChatResponse(result: ChatResponseResult): Promise<void> {
     this._responseLog.push(result);
+    await this._responseHookGate;
   }
 
   async getResponseLog(): Promise<ChatResponseResult[]> {