useAgentChat keeps a cut-off reply after the socket closes mid-stream, and the next send stores it over the real one #2464

Describe the bug

If the WebSocket closes while a reply streams, and the client reconnects after the reply has finished, the client never shows the finished reply. It keeps the cut-off copy for good. The next message then sends that copy to the server, and AIChatAgent stores it. When the reply carries its stored id, the cut-off copy overwrites the stored reply. When it does not, the server stores the copy as a second reply.

We hit this in production with an iOS home-screen PWA. The user sent a message and switched apps at once, and iOS closed the socket 0.7 s later. The server finished the turn and stored the full reply. The app showed no reply when the user came back. The user's next message erased the stored reply.

To Reproduce

Repro app: https://github.com/chhoumann/agents-chat-repro. It is a plain AIChatAgent that streams a scripted reply (no model), a plain useAgentChat page, and a Playwright script that closes the page's socket mid-reply. The README has the steps.

1. bun install && bunx playwright install chromium
2. bunx vite dev --port 3520
3. START=1 bun repro.ts down-until-done

START=1 makes onChatMessage write { type: "start" }, so the reply carries its stored id. The script sends a message, closes the socket after the first words, keeps reconnects failing until the reply has finished, and lets PartySocket reconnect. Then it sends a follow-up:

PASS  down-until-done: the server stored the full reply
PASS  down-until-done: the client reconnects
FAIL  down-until-done: after reconnect the full reply is on screen  -- "assistant: REPLY to first-a-murhyqpu "
PASS  down-until-done: the follow-up gets its full reply
FAIL  down-until-done: after the follow-up the first reply is complete on screen  -- "assistant: REPLY to first-a-murhyqpu "
FAIL  down-until-done: after the follow-up the first reply is stored complete  -- "REPLY to first-a-murhyqpu "

Without START=1, the follow-up stores the cut-off copy as a second assistant message next to the full one. A second scenario, back-mid-stream, reconnects while the reply still streams. Without a stored id, the reply's text then shows twice, and the follow-up stores a second message with the doubled text.

Expected behavior

After the reconnect, the client shows the stored reply. A later send never replaces or duplicates a stored reply with the client's copy of it.

Version:

agents@0.26.0 with @cloudflare/ai-chat@0.12.1. Also agents@0.17.4 with @cloudflare/ai-chat@0.9.3.

Additional context

Three things combine. Line numbers are against main at 2f3176b.

1. The client never releases the protection when the socket closes. useAgentChat protects the streaming reply when the local request's start chunk arrives (protectedStreamingAssistantRef). Only that request's done or error frame releases it (restoreProtectedStreamingAssistant, line 2501). The other release sites are a snapshot with a later assistant message, a chat reset, and the effect cleanup. onAgentClose (packages/agents/src/chat/react.tsx:2735) does not release it, and after a close the done frame never arrives. From then on, preserveProtectedStreamingAssistant swaps the stale local copy back into every cf_agent_chat_messages snapshot. A snapshot that contains the full reply cannot repair the client.

2. AIChatAgent sends no transcript to a client that connects between turns. On connect it sends cf_agent_stream_resuming, cf_agent_stream_pending, or a recovering frame. A client that connects after the turn ended gets cf_agent_stream_resume_none with reason: "idle" and nothing else, so the finished reply never reaches it. chat-agents.md says that "stream resumption and message sync are handled for you automatically". Think sends an idle-connect snapshot (see Think: reply vanishes for clients that reconnect around stream end (turn-end broadcast, connect snapshot and /get-messages all serve a stale message cache) #2132). AIChatAgent does not.

3. The server stores the client's copy of messages it already has. The chat-request handler calls persistMessages(transformedMessages, [connection.id], { _deleteStaleRows: true }). reconcileMessages upserts a client assistant message with the stored id over the stored row, even when its text is a cut-off prefix. It stores a client copy with an unknown id as a new row. fix(ai-chat): preserve server-generated messages when client appends new messages #1038 stopped the deletion of server replies when a client appends, but this path still writes the client's stale copy.

Proposed fix

• For (1), release the protection in onAgentClose. client-fix.patch in the repro is a 2-line change:

          serverReportsRecovery = false;
   +      protectedStreamingAssistantRef.current = null;
   +      localResponseIds.clear();
          fallbackAckedResumeRequestIds.clear();

• For (2), send the stored transcript on connect when no turn is streaming, as Think does. Send it only between turns. During a turn the transcript lacks the streaming reply. With (1) applied, it then takes the partial reply off screen while the replay rebuilds it, and the text doubles. We saw this in our app. On 0.25.0 and later the frame probably needs connect: true (useAgentChat: cf_agent_chat_messages handler replaces the message array wholesale, dropping client-buffered sends on reconnect #1983, fix(agents): preserve optimistic send dropped on reconnect (#1983) #2039).

• With (1) and (2), every check in the repro passes when the reply carries its stored id (FIX=1 START=1). Without a stored id, back-mid-stream still doubles the text. resumable-streaming.md says that replays of non-continuation streams "are not filtered", and fix(chat): skip replayed continuation chunks the client already applied (#1951) #2348 says the doubled part shows "until the final snapshot replaces the message". In this case no final snapshot reaches the reconnected client.

• For (3), we suggest the server keep its stored version of every message it already has. A request would then only add new messages, or drop the messages after a resent one (regenerate). We did this in our app by overriding persistMessages for writes that carry _deleteStaleRows. Our app has no client-side tools. Client tool results can arrive inside the client's assistant message, so an SDK fix would need to keep merging those.

Related: #2132 (Think, same symptom, now fixed), #2013 and #2019 (a close mid-stream now sets the error state, but the reply still never arrives), #1038, #1733, #1951 and #2348 (replay duplication), #2166, #1983 and #2039 (the connect flag).

I can open a PR with the client fix (1) and a regression test in your suite, if that helps.