fix(chat): keep AI SDK status correct when reconnect races the pre-stream window (#1784) (#1803)

* fix(chat): keep AI SDK status correct when reconnect races the pre-stream window (#1784)

A chat turn spends a window between "request accepted" and "first chunk
produced" in neither the active-stream nor terminal-replay state: it is queued,
debouncing, waiting on waitForMcpConnections, or running async setup inside
onChatMessage before a stream object exists. A client that reconnected or
re-mounted in that window was answered with cf_agent_stream_resume_none and gave
up, so the turn the server went on to stream never drove the client's AI SDK
`status` — the UI stayed stuck at "ready" until a full remount.

Server (agents/chat):
- New shared PreStreamTurns tracker (pre-stream-turns.ts) for accepted-but-not-
  yet-streamed turns and the connections parked waiting on them, mirroring
  ContinuationState (pure data + send-through-callback).
- New server->client cf_agent_stream_pending frame (protocol + wire-types +
  golden builder).
- ResumeHandshake now parks resume requests that arrive during the pre-stream
  window and emits STREAM_PENDING ("keep waiting") instead of RESUME_NONE, then
  flushes parked connections into the normal STREAM_RESUMING handshake on stream
  start. Continuation affinity is relaxed via an optional isConnectionPresent
  host hook so a transparent reconnect (whose connection id changed) can resume
  a continuation whose original owner connection is gone.

Client:
- ws-chat-transport: handleStreamPending() extends the resume probe from the 5s
  fast path to a 60s backstop so the probe stays open across the gap.
- useAgentChat re-probes the stream on a transparent socket reopen (e.g. a 1006
  reconnect that does not remount the component) so status recovers.

Hosts:
- Wired the begin/park/flush/settle lifecycle into both AIChatAgent and
  @cloudflare/think.
- Skipped turns (supersede / queue generation change) settle WITHOUT releasing
  parked connections (releaseParked: false), so a client parked during the
  window survives onto the successor turn instead of being cut loose by a
  premature RESUME_NONE in the supersede/settle microtask race.

Hibernation: PreStreamTurns is in-memory only and is safe because the pre-stream
window cannot overlap hibernation — a turn between begin() and stream start is an
unresolved message-handler promise that pins the DO in memory, so eviction only
happens once a durable stream exists (resumed via ResumableStream) or the turn
finished. Documented as an invariant on each _preStream field.

Tests:
- Unit: PreStreamTurns (incl. skip-path contract), handshake park/pending/flush/
  affinity, golden STREAM_PENDING frame.
- Transport: STREAM_PENDING keep-waiting timeout extension (incl. tool
  continuation path).
- Integration (workers runtime, real AIChatAgent): onConnect/resume-request park
  -> pending -> resuming, and parked-client-survives-overlapping-submits.
- React hook: transparent reconnect re-probe + STREAM_PENDING keepalive.

Co-authored-by: Cursor <cursoragent@cursor.com>

* fix(ai-chat): guarantee pre-stream turn settles if a pre-turn-body step throws (#1784)

AIChatAgent.begin()'d the pre-stream turn before persistMessages,
_mergeQueuedUserMessages, and the queued mcp.waitForConnections /
_setRequestContext steps. A throw in any of those never reached the
settle in chatTurnBody's finally, so the request id stayed stuck in
_preStream (hasInFlight() true forever) and every later client was
parked on STREAM_PENDING (60s) instead of getting an immediate
STREAM_RESUME_NONE — until chat clear or DO eviction.

Wrap the whole post-begin() scope in a top-level try/finally that always
calls _settlePreStreamTurn, mirroring @cloudflare/think. Idempotent with
the existing happy-path settle.

Co-authored-by: Cursor <cursoragent@cursor.com>

---------

Co-authored-by: Sunil Pai <18808+threepointone@users.noreply.github.com>
Co-authored-by: Cursor <cursoragent@cursor.com>

Author: threepointone

.changeset/prestream-resume-keep-waiting.md

@@ -0,0 +1,41 @@
+---
+"agents": patch
+"@cloudflare/ai-chat": patch
+"@cloudflare/think": patch
+---
+
+Fix AI SDK `status` getting stuck after a reconnect that races a turn's
+pre-stream window (#1784).
+
+A turn is "accepted but pre-stream" while it is queued, debouncing, or awaiting
+async setup before its resumable stream starts. A client that connected or sent
+a `STREAM_RESUME_REQUEST` in that window was answered with `STREAM_RESUME_NONE`
+("nothing to resume"), so its short resume probe resolved `null` and AI SDK
+`status` settled on `ready` even though the server went on to stream — leaving
+the UI unable to render the in-flight turn until a full remount.
+
+This adds a shared `PreStreamTurns` tracker (`agents/chat`) and a new
+server→client `cf_agent_stream_pending` frame:
+
+- The resume handshake now parks resume requests that arrive during the
+  pre-stream window and emits `STREAM_PENDING` ("keep waiting") instead of
+  `STREAM_RESUME_NONE`, then flushes parked connections into the normal
+  `STREAM_RESUMING` handshake once the stream actually starts (and releases them
+  with `STREAM_RESUME_NONE` if the turn is superseded/cleared before streaming).
+- On `STREAM_PENDING` the client transport extends its resume probe from the
+  5s fast-path to a 60s backstop so the probe stays open across the gap.
+- `useAgentChat` re-probes the stream on a transparent socket reopen (e.g. a
+  1006 reconnect that does not remount the component) so `status` recovers.
+- Continuation affinity is relaxed via an optional `isConnectionPresent` host
+  hook so a transparent reconnect (whose connection id changed) can resume a
+  continuation whose original owner connection is gone.
+
+Wired into both `AIChatAgent` and `@cloudflare/think`.
+
+The pre-stream tracker is in-memory only; it is hibernation-safe because a turn
+in its pre-stream window is an unresolved message-handler promise that pins the
+Durable Object in memory, so eviction only happens once a stream is durably
+recorded (and resumes via `ResumableStream`) or the turn has finished. Skipped
+turns (supersede/generation change) settle without releasing parked
+connections, so a client parked during the window survives onto the successor
+turn instead of being cut loose by a premature `STREAM_RESUME_NONE`.

packages/agents/src/chat/__tests__/pre-stream-turns.test.ts

@@ -0,0 +1,176 @@
+import { describe, it, expect, vi } from "vitest";
+import { PreStreamTurns } from "../pre-stream-turns";
+import type { ChatConnection } from "../connection";
+
+type SentFrame = Record<string, unknown>;
+
+function makeConnection(
+  id: string,
+  sink?: SentFrame[]
+): ChatConnection & { send: ReturnType<typeof vi.fn> } {
+  return {
+    id,
+    send: vi.fn((message: string) => {
+      sink?.push(JSON.parse(message) as SentFrame);
+    })
+  };
+}
+
+describe("PreStreamTurns (#1784)", () => {
+  it("starts idle: nothing in flight, park is a no-op that sends no frame", () => {
+    const pre = new PreStreamTurns();
+    expect(pre.hasInFlight()).toBe(false);
+
+    const conn = makeConnection("c1");
+    expect(pre.park(conn)).toBe(false);
+    expect(conn.send).not.toHaveBeenCalled();
+    expect(pre.awaitingConnections.size).toBe(0);
+  });
+
+  it("begin marks in flight and exposes the latest request id", () => {
+    const pre = new PreStreamTurns();
+    pre.begin("req-1");
+    expect(pre.hasInFlight()).toBe(true);
+    expect(pre.latestRequestId).toBe("req-1");
+  });
+
+  it("park enrolls a connection and sends STREAM_PENDING with the request id", () => {
+    const frames: SentFrame[] = [];
+    const pre = new PreStreamTurns();
+    pre.begin("req-1");
+
+    const conn = makeConnection("c1", frames);
+    expect(pre.park(conn)).toBe(true);
+    expect(pre.awaitingConnections.get("c1")).toBe(conn);
+    expect(frames).toEqual([{ type: "cf_agent_stream_pending", id: "req-1" }]);
+  });
+
+  it("flushOnStreamStart notifies each parked connection then clears the map (keeps accepted)", () => {
+    const pre = new PreStreamTurns();
+    pre.begin("req-1");
+    const c1 = makeConnection("c1");
+    const c2 = makeConnection("c2");
+    pre.park(c1);
+    pre.park(c2);
+
+    const notified: string[] = [];
+    pre.flushOnStreamStart((c) => notified.push(c.id));
+
+    expect(notified.sort()).toEqual(["c1", "c2"]);
+    expect(pre.awaitingConnections.size).toBe(0);
+    // The turn is still running until settle() — accepted set untouched.
+    expect(pre.hasInFlight()).toBe(true);
+  });
+
+  it("releaseAwaiting sends STREAM_RESUME_NONE to each parked connection then clears", () => {
+    const frames: SentFrame[] = [];
+    const pre = new PreStreamTurns();
+    pre.begin("req-1");
+    const c1 = makeConnection("c1", frames);
+    pre.park(c1);
+    frames.length = 0;
+
+    pre.releaseAwaiting();
+
+    expect(frames).toEqual([{ type: "cf_agent_stream_resume_none" }]);
+    expect(pre.awaitingConnections.size).toBe(0);
+  });
+
+  it("settle returns true only when the last accepted turn settles", () => {
+    const pre = new PreStreamTurns();
+    pre.begin("a");
+    pre.begin("b");
+
+    expect(pre.settle("a")).toBe(false);
+    expect(pre.hasInFlight()).toBe(true);
+    expect(pre.settle("b")).toBe(true);
+    expect(pre.hasInFlight()).toBe(false);
+    expect(pre.latestRequestId).toBeNull();
+  });
+
+  it("settling an unknown id is safe and reports idle when nothing else is in flight", () => {
+    const pre = new PreStreamTurns();
+    expect(pre.settle("never-began")).toBe(true);
+  });
+
+  it("a connection parked during the gap between queued turns survives the first turn's settle", () => {
+    // A begins, streams, settles while B is still queued (begin'd). A reconnect
+    // that parked must NOT be released until B also settles with no stream.
+    const pre = new PreStreamTurns();
+    pre.begin("a");
+    pre.begin("b");
+    const conn = makeConnection("c1");
+    pre.park(conn);
+
+    // A finishes first: not idle (B still in flight) → caller would NOT release.
+    expect(pre.settle("a")).toBe(false);
+    expect(pre.awaitingConnections.get("c1")).toBe(conn);
+
+    // B finishes: now idle.
+    expect(pre.settle("b")).toBe(true);
+  });
+
+  it("models the skip-path contract: settle WITHOUT releaseAwaiting keeps a parked client for the successor (#1784)", () => {
+    // Mirrors the host `_completeSkippedRequest` flow: a superseded turn settles
+    // out of the accepted set but must NOT release parked connections, so a
+    // client parked during the pre-stream window survives onto the turn that
+    // actually streams — even if the set momentarily empties before the
+    // successor's `begin()` runs (the supersede/settle microtask race).
+    const pre = new PreStreamTurns();
+    pre.begin("superseded");
+    const conn = makeConnection("c1");
+    pre.park(conn);
+
+    // The superseded turn settles. The host does NOT call releaseAwaiting here.
+    expect(pre.settle("superseded")).toBe(true); // set momentarily empty
+    expect(pre.awaitingConnections.get("c1")).toBe(conn); // still parked
+
+    // The successor turn begins and starts streaming → the parked client is
+    // flushed into the normal resume handshake rather than stranded.
+    pre.begin("successor");
+    const notified: string[] = [];
+    pre.flushOnStreamStart((c) => notified.push(c.id));
+    expect(notified).toEqual(["c1"]);
+    expect(pre.awaitingConnections.size).toBe(0);
+  });
+
+  it("release drops a single connection without touching the rest", () => {
+    const pre = new PreStreamTurns();
+    pre.begin("req-1");
+    const c1 = makeConnection("c1");
+    const c2 = makeConnection("c2");
+    pre.park(c1);
+    pre.park(c2);
+
+    pre.release("c1");
+
+    expect(pre.awaitingConnections.has("c1")).toBe(false);
+    expect(pre.awaitingConnections.get("c2")).toBe(c2);
+  });
+
+  it("park without a known request id omits the id from the keep-waiting frame", () => {
+    const frames: SentFrame[] = [];
+    const pre = new PreStreamTurns();
+    // settle the only id so latestRequestId is null but force in-flight via a
+    // second begin with an empty-ish id path: emulate by beginning then parking.
+    pre.begin("");
+    const conn = makeConnection("c1", frames);
+    pre.park(conn);
+    expect(frames).toEqual([{ type: "cf_agent_stream_pending" }]);
+  });
+
+  it("reset drops all state and sends no frames", () => {
+    const pre = new PreStreamTurns();
+    pre.begin("req-1");
+    const c1 = makeConnection("c1");
+    pre.park(c1);
+    c1.send.mockClear();
+
+    pre.reset();
+
+    expect(pre.hasInFlight()).toBe(false);
+    expect(pre.awaitingConnections.size).toBe(0);
+    expect(pre.latestRequestId).toBeNull();
+    expect(c1.send).not.toHaveBeenCalled();
+  });
+});

packages/agents/src/chat/__tests__/resume-handshake-frames.test.ts

@@ -4,6 +4,7 @@ import {
   HANDSHAKE_INVARIANTS,
   recoveringFrame,
   replayDoneFrame,
+  streamPendingFrame,
   streamResumeNoneFrame,
   streamResumingFrame,
   terminalErrorFrame
@@ -54,6 +55,16 @@ describe("resume-handshake golden frames", () => {
     expect("replay" in frame).toBe(false);
   });
 
+  it("STREAM_PENDING includes the id only when present (#1784)", () => {
+    expect(streamPendingFrame("req-1")).toEqual({
+      type: "cf_agent_stream_pending",
+      id: "req-1"
+    });
+    const noId = streamPendingFrame();
+    expect(noId).toEqual({ type: "cf_agent_stream_pending" });
+    expect("id" in noId).toBe(false);
+  });
+
   it("recovering frame includes the id only when present", () => {
     expect(recoveringFrame("req-1")).toEqual({
       type: "cf_agent_chat_recovering",
@@ -82,6 +93,9 @@ describe("resume-handshake golden frames", () => {
       CHAT_MESSAGE_TYPES.USE_CHAT_RESPONSE
     );
     expect(recoveringFrame("x").type).toBe(CHAT_MESSAGE_TYPES.CHAT_RECOVERING);
+    expect(streamPendingFrame("x").type).toBe(
+      CHAT_MESSAGE_TYPES.STREAM_PENDING
+    );
   });
 
   it("a host with a distinct use-chat-response constant is honored", () => {

packages/agents/src/chat/__tests__/resume-handshake-frames.ts

@@ -54,6 +54,20 @@ export function streamResumeNoneFrame() {
   };
 }
 
+/**
+ * Server -> client `STREAM_PENDING` keep-waiting frame (#1784): a turn is
+ * accepted but its resumable stream has not started yet. Carries `id` only when
+ * the accepted request id is known. The client cancels its short resume probe
+ * timeout and waits for a later `STREAM_RESUMING` or `STREAM_RESUME_NONE`.
+ * Emitted by `PreStreamTurns.park`.
+ */
+export function streamPendingFrame(requestId?: string) {
+  return {
+    type: CHAT_MESSAGE_TYPES.STREAM_PENDING,
+    ...(requestId ? { id: requestId } : {})
+  };
+}
+
 /**
  * Server -> client terminal replay-done frame: the ACK fallback when no live or
  * completed chunks remain to replay, closing the resumed stream cleanly.