fix(agents): preserve optimistic send dropped on reconnect (#1983)

.changeset/reconnect-optimistic-send.md

@@ -0,0 +1,5 @@
+---
+"agents": patch
+---
+
+Fix `useAgentChat` dropping an in-flight optimistic send on reconnect (#1983). A message buffered by PartySocket while the socket was down was missing from the transcript the server sends on connect, so the whole-array `setMessages` erased it from the UI until the turn completed. The hook now keeps such buffered sends when a snapshot is marked `connect: true` (the new optional field on `cf_agent_chat_messages`). Every other snapshot stays authoritative, so a server rollback of a rejected send (for example under `messageConcurrency: "drop"`) still removes it.

.changeset/think-connect-transcript-marker.md

@@ -0,0 +1,5 @@
+---
+"@cloudflare/think": patch
+---
+
+Mark the transcript Think sends to a newly connected client with `connect: true`, so `useAgentChat` keeps a message the user sent while disconnected instead of erasing it on reconnect (#1983).

packages/agents/src/chat/react.tsx

@@ -651,6 +651,25 @@ function prependMissingHydratedMessages<ChatMessage extends UIMessage>(
   return [...missingHydratedMessages, ...currentMessages];
 }
 
+// Re-append the specific buffered sends a connect transcript omits, restoring
+// only the tracked ids (in local order).
+function restoreBufferedSends<ChatMessage extends UIMessage>(
+  snapshot: ChatMessage[],
+  local: readonly ChatMessage[],
+  bufferedIds: ReadonlySet<string>
+): ChatMessage[] {
+  if (bufferedIds.size === 0) {
+    return snapshot;
+  }
+
+  const snapshotIds = new Set(snapshot.map((message) => message.id));
+  const restored = local.filter(
+    (message) => bufferedIds.has(message.id) && !snapshotIds.has(message.id)
+  );
+
+  return restored.length === 0 ? snapshot : [...snapshot, ...restored];
+}
+
 /**
  * React hook for building AI chat interfaces using an Agent
  * @param options Chat options including the agent connection
@@ -1047,6 +1066,10 @@ export function useAgentChat<
    * Used by onAgentMessage to skip messages already handled by the transport.
    */
   const localRequestIdsRef = useRef<Set<string>>(new Set());
+  // Ids of sends the transport buffered (socket not OPEN when `send()` ran, so
+  // the server hasn't seen them). Consumed by the next transcript snapshot to
+  // rescue the optimistic messages from the reconnect replay.
+  const pendingBufferedSendIdsRef = useRef<Set<string>>(new Set());
   const pendingReplayResumeRequestIdsRef = useRef<Set<string>>(new Set());
   const replayHydratedAssistantMessageIdsRef = useRef<Set<string>>(new Set());
   /**
@@ -1079,6 +1102,11 @@ export function useAgentChat<
       agent: agentRef.current,
       activeRequestIds: localRequestIdsRef.current,
       cancelOnClientAbort,
+      onRequestBuffered: (messageId) => {
+        if (messageId) {
+          pendingBufferedSendIdsRef.current.add(messageId);
+        }
+      },
       prepareBody: async ({ messages: msgs, trigger, messageId }) => {
         // Start with the top-level body option (static or dynamic)
         let extraBody: Record<string, unknown> = {};
@@ -1624,6 +1652,7 @@ export function useAgentChat<
     fallbackAckedResumeRequestIdsRef.current.clear();
     replayHydratedAssistantMessageIdsRef.current.clear();
     protectedStreamingAssistantRef.current = null;
+    pendingBufferedSendIdsRef.current.clear();
     customTransport.appliedChunks.clear();
     turnErrorsRef.current.clear();
     pendingTurnEndsRef.current = [];
@@ -1636,6 +1665,9 @@ export function useAgentChat<
 
   const sendMessageWithStreamingProtection: typeof sendMessage = useCallback(
     async (message, options) => {
+      // Whether this send was buffered is detected at the real send site in the
+      // transport (via `onRequestBuffered`), not here — the socket can drop
+      // during the async request preparation between now and the actual send().
       const request = sendMessage(message, options);
 
       if (
@@ -2126,6 +2158,14 @@ export function useAgentChat<
           break;
 
         case MessageType.CF_AGENT_CHAT_MESSAGES: {
+          // One-shot, consumed outside the updater so a re-invoked updater
+          // sees the same ids. Only a connect transcript predates the buffered
+          // sends; any other snapshot (e.g. a `drop` rollback after a
+          // mid-stream reconnect) has already seen them and wins.
+          const bufferedSendIds = data.connect
+            ? new Set(pendingBufferedSendIdsRef.current)
+            : new Set<string>();
+          pendingBufferedSendIdsRef.current.clear();
           setMessages((currentMessages: ChatMessage[]) => {
             let next = preserveProtectedStreamingAssistant(
               data.messages,
@@ -2165,7 +2205,9 @@ export function useAgentChat<
                 next = observed.accumulator.mergeInto(next) as ChatMessage[];
               }
             }
-            return next;
+            // Rescue sends the transport buffered while the socket was down
+            // that the reconnect replay would otherwise drop.
+            return restoreBufferedSends(next, currentMessages, bufferedSendIds);
           });
           break;
         }

packages/agents/src/chat/wire-types.ts

@@ -77,6 +77,13 @@ export type OutgoingMessage<ChatMessage extends UIMessage = UIMessage> =
       type: MessageType.CF_AGENT_CHAT_MESSAGES;
       /** Array of chat messages */
       messages: readonly ChatMessage[];
+      /**
+       * Set on the transcript a server sends to a newly connected client. It
+       * predates any request the client buffered while disconnected, so the
+       * client keeps those optimistic sends; any other snapshot is
+       * authoritative (#1983).
+       */
+      connect?: boolean;
     }
   | {
       /** Indicates this message is a response to a chat request */

packages/agents/src/chat/ws-chat-transport.ts

@@ -55,7 +55,9 @@ function interruptChatStream(
  * `useAgent` from `agents/react`.
  */
 export interface AgentConnection {
-  send: (data: string) => void;
+  // PartySocket returns `false` when the frame was buffered (socket not OPEN)
+  // rather than sent immediately; `void` keeps non-PartySocket providers valid.
+  send: (data: string) => boolean | void;
   addEventListener: (
     type: string,
     listener: (event: MessageEvent) => void,
@@ -93,6 +95,13 @@ export type WebSocketChatTransportOptions<
    * @default false
    */
   cancelOnClientAbort?: boolean;
+  /**
+   * Called when a `submit-message` request frame was buffered rather than sent
+   * immediately (the socket was not OPEN, so PartySocket queued it for the next
+   * reconnect). The hook uses this to preserve the optimistic message across the
+   * reconnect transcript replay. `messageId` is the submitted user message's id.
+   */
+  onRequestBuffered?: (messageId: string | undefined) => void;
 };
 
 /**
@@ -111,6 +120,7 @@ export class WebSocketChatTransport<
   private prepareBody?: WebSocketChatTransportOptions<ChatMessage>["prepareBody"];
   private activeRequestIds?: Set<string>;
   private cancelOnClientAbort: boolean;
+  private onRequestBuffered?: WebSocketChatTransportOptions<ChatMessage>["onRequestBuffered"];
 
   // Pending resume resolver — set by reconnectToStream, called by
   // handleStreamResuming when onAgentMessage sees CF_AGENT_STREAM_RESUMING.
@@ -150,6 +160,7 @@ export class WebSocketChatTransport<
     this.prepareBody = options.prepareBody;
     this.activeRequestIds = options.activeRequestIds;
     this.cancelOnClientAbort = options.cancelOnClientAbort ?? false;
+    this.onRequestBuffered = options.onRequestBuffered;
   }
 
   /**
@@ -496,9 +507,13 @@ export class WebSocketChatTransport<
       return stream;
     }
 
-    // Send the request over WebSocket
+    // Send the request over WebSocket. `send()` returns false when the frame
+    // was buffered (socket not OPEN) rather than sent immediately; a buffered
+    // submit means the server hasn't seen the message yet, so tell the hook to
+    // preserve it across the reconnect transcript replay. Only an explicit
+    // false counts — non-PartySocket doubles return undefined (treated as sent).
     requestSent = true;
-    agent.send(
+    const sent = agent.send(
       JSON.stringify({
         id: requestId,
         init: {
@@ -508,6 +523,12 @@ export class WebSocketChatTransport<
         type: MessageType.CF_AGENT_USE_CHAT_REQUEST
       })
     );
+    if (sent === false && options.trigger === "submit-message") {
+      const lastUserMessage = [...options.messages]
+        .reverse()
+        .find((message) => message.role === "user");
+      this.onRequestBuffered?.(lastUserMessage?.id);
+    }
 
     return stream;
   }