fix: flush replayed stream state and handle orphaned streams after hibernation

.changeset/fix-hibernated-stream-state.md

@@ -0,0 +1,9 @@
+---
+"@cloudflare/ai-chat": patch
+---
+
+Fix active streams losing UI state after reconnect and dead streams after DO hibernation.
+
+- Send `replayComplete` signal after replaying stored chunks for live streams, so the client flushes accumulated parts to React state immediately instead of waiting for the next live chunk.
+- Detect orphaned streams (restored from SQLite after hibernation with no live LLM reader) via `_isLive` flag on `ResumableStream`. On reconnect, send `done: true`, complete the stream, and reconstruct/persist the partial assistant message from stored chunks.
+- Client-side: flush `activeStreamRef` on `replayComplete` (keeps stream alive for subsequent live chunks) and on `done` during replay (finalizes orphaned streams).

packages/ai-chat/src/index.ts

@@ -509,10 +509,18 @@ export class AIChatAgent<
             this._resumableStream.hasActiveStream() &&
             this._resumableStream.activeRequestId === data.id
           ) {
-            this._resumableStream.replayChunks(
+            const orphanedStreamId = this._resumableStream.replayChunks(
               connection,
               this._resumableStream.activeRequestId
             );
+
+            // If the stream was orphaned (restored from SQLite after
+            // hibernation with no live reader), reconstruct the partial
+            // assistant message from stored chunks and persist it so it
+            // survives further page refreshes.
+            if (orphanedStreamId) {
+              this._persistOrphanedStream(orphanedStreamId);
+            }
           }
           return;
         }
@@ -766,6 +774,62 @@ export class AIChatAgent<
     this._resumableStream.markError(streamId);
   }
 
+  /**
+   * Reconstruct and persist a partial assistant message from an orphaned
+   * stream's stored chunks. Called when the DO wakes from hibernation and
+   * discovers an active stream with no live LLM reader.
+   *
+   * Replays each chunk body through `applyChunkToParts` to rebuild the
+   * message parts, then persists the result so it survives further refreshes.
+   * @internal
+   */
+  private _persistOrphanedStream(streamId: string) {
+    const chunks = this._resumableStream.getStreamChunks(streamId);
+    if (!chunks.length) return;
+
+    const message: ChatMessage = {
+      id: `assistant_${Date.now()}_${Math.random().toString(36).slice(2, 11)}`,
+      role: "assistant",
+      parts: []
+    };
+
+    for (const chunk of chunks) {
+      try {
+        const data = JSON.parse(chunk.body);
+
+        // Capture message ID from the "start" event if present
+        if (data.type === "start" && data.messageId != null) {
+          message.id = data.messageId;
+        }
+        if (
+          (data.type === "start" ||
+            data.type === "finish" ||
+            data.type === "message-metadata") &&
+          data.messageMetadata != null
+        ) {
+          message.metadata = message.metadata
+            ? { ...message.metadata, ...data.messageMetadata }
+            : data.messageMetadata;
+        }
+
+        applyChunkToParts(message.parts, data);
+      } catch {
+        // Skip malformed chunk bodies
+      }
+    }
+
+    if (message.parts.length > 0) {
+      // Check if a message with this ID already exists (e.g., from an
+      // early persist during tool approval). Update in place if so.
+      const existingIdx = this.messages.findIndex((m) => m.id === message.id);
+      const updatedMessages =
+        existingIdx >= 0
+          ? this.messages.map((m, i) => (i === existingIdx ? message : m))
+          : [...this.messages, message];
+      this.persistMessages(updatedMessages);
+    }
+  }
+
   /**
    * Restore _lastBody and _lastClientTools from SQLite.
    * Called in the constructor so these values survive DO hibernation.

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

@@ -1433,3 +1433,294 @@ describe("useAgentChat stale agent ref (issue #929)", () => {
     expect(oldSendCalls.length).toBe(0);
   });
 });
+
+describe("useAgentChat stream resumption (issue #896)", () => {
+  function createAgentWithTarget({ name, url }: { name: string; url: string }) {
+    const target = new EventTarget();
+    const sentMessages: string[] = [];
+    const agent = createAgent({
+      name,
+      url,
+      send: (data: string) => sentMessages.push(data)
+    });
+    // Wire up the target so we can dispatch messages to the hook
+    (agent as unknown as Record<string, unknown>).addEventListener =
+      target.addEventListener.bind(target);
+    (agent as unknown as Record<string, unknown>).removeEventListener =
+      target.removeEventListener.bind(target);
+    return { agent, target, sentMessages };
+  }
+
+  function dispatch(target: EventTarget, data: Record<string, unknown>) {
+    target.dispatchEvent(
+      new MessageEvent("message", { data: JSON.stringify(data) })
+    );
+  }
+
+  it("should flush messages to state after replayComplete for live streams", async () => {
+    const { agent, target } = createAgentWithTarget({
+      name: "replay-complete-test",
+      url: "ws://localhost:3000/agents/chat/replay-complete-test?_pk=abc"
+    });
+
+    const TestComponent = () => {
+      const chat = useAgentChat({
+        agent,
+        getInitialMessages: null,
+        messages: [] as UIMessage[]
+      });
+      const assistantMsg = chat.messages.find(
+        (m: UIMessage) => m.role === "assistant"
+      );
+      const textPart = assistantMsg?.parts.find(
+        (p: UIMessage["parts"][number]) => p.type === "text"
+      ) as { text?: string } | undefined;
+      return (
+        <div>
+          <div data-testid="count">{chat.messages.length}</div>
+          <div data-testid="text">{textPart?.text ?? ""}</div>
+        </div>
+      );
+    };
+
+    const screen = await act(async () => {
+      const screen = render(<TestComponent />, {
+        wrapper: ({ children }) => (
+          <StrictMode>
+            <Suspense fallback="Loading...">{children}</Suspense>
+          </StrictMode>
+        )
+      });
+      await sleep(10);
+      return screen;
+    });
+
+    // Initially no messages
+    await expect.element(screen.getByTestId("count")).toHaveTextContent("0");
+
+    // Simulate server sending CF_AGENT_STREAM_RESUMING
+    await act(async () => {
+      dispatch(target, {
+        type: "cf_agent_stream_resuming",
+        id: "req-1"
+      });
+      await sleep(10);
+    });
+
+    // Simulate replay chunks (these are batched, no per-chunk flush)
+    await act(async () => {
+      dispatch(target, {
+        type: "cf_agent_use_chat_response",
+        id: "req-1",
+        body: '{"type":"text-start","id":"t1"}',
+        done: false,
+        replay: true
+      });
+      dispatch(target, {
+        type: "cf_agent_use_chat_response",
+        id: "req-1",
+        body: '{"type":"text-delta","id":"t1","delta":"Hello world"}',
+        done: false,
+        replay: true
+      });
+      await sleep(10);
+    });
+
+    // Should still be 0 messages — replay chunks are not flushed yet
+    await expect.element(screen.getByTestId("count")).toHaveTextContent("0");
+
+    // Now send replayComplete — this should trigger a flush
+    await act(async () => {
+      dispatch(target, {
+        type: "cf_agent_use_chat_response",
+        id: "req-1",
+        body: "",
+        done: false,
+        replay: true,
+        replayComplete: true
+      });
+      await sleep(10);
+    });
+
+    // Now the assistant message should appear
+    await expect.element(screen.getByTestId("count")).toHaveTextContent("1");
+    await expect
+      .element(screen.getByTestId("text"))
+      .toHaveTextContent("Hello world");
+  });
+
+  it("should flush and finalize after done:true for orphaned streams", async () => {
+    const { agent, target } = createAgentWithTarget({
+      name: "orphaned-done-test",
+      url: "ws://localhost:3000/agents/chat/orphaned-done-test?_pk=abc"
+    });
+
+    const TestComponent = () => {
+      const chat = useAgentChat({
+        agent,
+        getInitialMessages: null,
+        messages: [] as UIMessage[]
+      });
+      const assistantMsg = chat.messages.find(
+        (m: UIMessage) => m.role === "assistant"
+      );
+      const textPart = assistantMsg?.parts.find(
+        (p: UIMessage["parts"][number]) => p.type === "text"
+      ) as { text?: string } | undefined;
+      return (
+        <div>
+          <div data-testid="count">{chat.messages.length}</div>
+          <div data-testid="text">{textPart?.text ?? ""}</div>
+          <div data-testid="status">{chat.status}</div>
+        </div>
+      );
+    };
+
+    const screen = await act(async () => {
+      const screen = render(<TestComponent />, {
+        wrapper: ({ children }) => (
+          <StrictMode>
+            <Suspense fallback="Loading...">{children}</Suspense>
+          </StrictMode>
+        )
+      });
+      await sleep(10);
+      return screen;
+    });
+
+    // Simulate resume + replay + done (orphaned stream path)
+    await act(async () => {
+      dispatch(target, {
+        type: "cf_agent_stream_resuming",
+        id: "req-orphaned"
+      });
+      await sleep(5);
+
+      dispatch(target, {
+        type: "cf_agent_use_chat_response",
+        id: "req-orphaned",
+        body: '{"type":"text-start","id":"t1"}',
+        done: false,
+        replay: true
+      });
+      dispatch(target, {
+        type: "cf_agent_use_chat_response",
+        id: "req-orphaned",
+        body: '{"type":"text-delta","id":"t1","delta":"partial from hibernation"}',
+        done: false,
+        replay: true
+      });
+
+      // done:true signals orphaned stream is finalized
+      dispatch(target, {
+        type: "cf_agent_use_chat_response",
+        id: "req-orphaned",
+        body: "",
+        done: true,
+        replay: true
+      });
+      await sleep(10);
+    });
+
+    // Message should be flushed with the accumulated text
+    await expect.element(screen.getByTestId("count")).toHaveTextContent("1");
+    await expect
+      .element(screen.getByTestId("text"))
+      .toHaveTextContent("partial from hibernation");
+  });
+
+  it("should continue receiving live chunks after replayComplete", async () => {
+    const { agent, target } = createAgentWithTarget({
+      name: "replay-then-live-test",
+      url: "ws://localhost:3000/agents/chat/replay-then-live-test?_pk=abc"
+    });
+
+    const TestComponent = () => {
+      const chat = useAgentChat({
+        agent,
+        getInitialMessages: null,
+        messages: [] as UIMessage[]
+      });
+      const assistantMsg = chat.messages.find(
+        (m: UIMessage) => m.role === "assistant"
+      );
+      const textPart = assistantMsg?.parts.find(
+        (p: UIMessage["parts"][number]) => p.type === "text"
+      ) as { text?: string } | undefined;
+      return (
+        <div>
+          <div data-testid="count">{chat.messages.length}</div>
+          <div data-testid="text">{textPart?.text ?? ""}</div>
+        </div>
+      );
+    };
+
+    const screen = await act(async () => {
+      const screen = render(<TestComponent />, {
+        wrapper: ({ children }) => (
+          <StrictMode>
+            <Suspense fallback="Loading...">{children}</Suspense>
+          </StrictMode>
+        )
+      });
+      await sleep(10);
+      return screen;
+    });
+
+    // Replay phase
+    await act(async () => {
+      dispatch(target, {
+        type: "cf_agent_stream_resuming",
+        id: "req-live"
+      });
+      await sleep(5);
+
+      dispatch(target, {
+        type: "cf_agent_use_chat_response",
+        id: "req-live",
+        body: '{"type":"text-start","id":"t1"}',
+        done: false,
+        replay: true
+      });
+      dispatch(target, {
+        type: "cf_agent_use_chat_response",
+        id: "req-live",
+        body: '{"type":"text-delta","id":"t1","delta":"replayed-"}',
+        done: false,
+        replay: true
+      });
+      dispatch(target, {
+        type: "cf_agent_use_chat_response",
+        id: "req-live",
+        body: "",
+        done: false,
+        replay: true,
+        replayComplete: true
+      });
+      await sleep(10);
+    });
+
+    // After replay, message should show replayed text
+    await expect.element(screen.getByTestId("count")).toHaveTextContent("1");
+    await expect
+      .element(screen.getByTestId("text"))
+      .toHaveTextContent("replayed-");
+
+    // Now simulate a live chunk arriving (no replay flag)
+    await act(async () => {
+      dispatch(target, {
+        type: "cf_agent_use_chat_response",
+        id: "req-live",
+        body: '{"type":"text-delta","id":"t1","delta":"and live!"}',
+        done: false
+      });
+      await sleep(10);
+    });
+
+    // The live chunk should append to the same message
+    await expect.element(screen.getByTestId("count")).toHaveTextContent("1");
+    await expect
+      .element(screen.getByTestId("text"))
+      .toHaveTextContent("replayed-and live!");
+  });
+});