fix: buffer websocket events during disconnects and replay on reconnect

agent/tools/papers_tool.py

@@ -249,7 +249,7 @@ def _format_read_paper_toc(parsed: dict[str, Any], arxiv_id: str) -> str:
         lines.append(f"{prefix}- **{s['title']}**: {preview}")
 
     lines.append(
-        '\n**Tip:** Call read_paper with section parameter (e.g. section="4" or section="Experiments") to read a specific section.'
+        '\nCall read_paper with section parameter (e.g. section="4" or section="Experiments") to read a specific section.'
     )
     return "\n".join(lines)
 

backend/routes/agent.py

@@ -386,20 +386,22 @@ async def websocket_endpoint(websocket: WebSocket, session_id: str) -> None:
         await websocket.close(code=4003, reason="Access denied")
         return
 
-    await ws_manager.connect(websocket, session_id)
-
-    # Send "ready" immediately on WebSocket connection so the frontend
-    # knows the session is alive.  The original ready event from _run_session
-    # fires before the WS is connected and is always lost.
-    try:
-        await websocket.send_json(
-            {
-                "event_type": "ready",
-                "data": {"message": "Agent initialized"},
-            }
-        )
-    except Exception as e:
-        logger.error(f"Failed to send ready event for session {session_id}: {e}")
+    had_buffered = await ws_manager.connect(websocket, session_id)
+
+    # Send "ready" on fresh connections so the frontend knows the session
+    # is alive.  Skip it when buffered events were flushed — those already
+    # contain the correct state and a ready would incorrectly reset
+    # isProcessing on the frontend.
+    if not had_buffered:
+        try:
+            await websocket.send_json(
+                {
+                    "event_type": "ready",
+                    "data": {"message": "Agent initialized"},
+                }
+            )
+        except Exception as e:
+            logger.error(f"Failed to send ready event for session {session_id}: {e}")
 
     try:
         while True:

backend/session_manager.py

@@ -320,6 +320,8 @@ class SessionManager:
         if not agent_session:
             return False
 
+        ws_manager.clear_buffer(session_id)
+
         # Clean up sandbox Space before cancelling the task
         await self._cleanup_sandbox(agent_session.session)
 

backend/websocket.py

@@ -1,12 +1,16 @@
 """WebSocket connection manager for real-time communication."""
 
 import logging
+from collections import deque
 from typing import Any
 
 from fastapi import WebSocket
 
 logger = logging.getLogger(__name__)
 
+# Max events buffered per session while WS is disconnected.
+MAX_EVENT_BUFFER = 500
+
 
 class ConnectionManager:
     """Manages WebSocket connections for multiple sessions."""
@@ -14,13 +18,36 @@ class ConnectionManager:
     def __init__(self) -> None:
         # session_id -> WebSocket
         self.active_connections: dict[str, WebSocket] = {}
+        # session_id -> events buffered while WS was disconnected
+        self._event_buffers: dict[str, deque[dict[str, Any]]] = {}
+
+    async def connect(self, websocket: WebSocket, session_id: str) -> bool:
+        """Accept a WebSocket connection, register it, and flush buffered events.
 
-    async def connect(self, websocket: WebSocket, session_id: str) -> None:
-        """Accept a WebSocket connection and register it."""
+        Returns True if buffered events were flushed (i.e. this is a reconnect
+        that had missed events).
+        """
         logger.info(f"Attempting to accept WebSocket for session {session_id}")
         await websocket.accept()
         self.active_connections[session_id] = websocket
+
+        # Flush events that were buffered while the WS was disconnected
+        buffered = self._event_buffers.pop(session_id, None)
+        if buffered:
+            logger.info(
+                f"Flushing {len(buffered)} buffered events for session {session_id}"
+            )
+            for message in buffered:
+                try:
+                    await websocket.send_json(message)
+                except Exception:
+                    logger.error(
+                        f"Error flushing buffered event for session {session_id}"
+                    )
+                    break
+
         logger.info(f"WebSocket connected and registered for session {session_id}")
+        return bool(buffered)
 
     def disconnect(self, session_id: str) -> None:
         """Remove a WebSocket connection."""
@@ -28,23 +55,39 @@ class ConnectionManager:
             del self.active_connections[session_id]
         logger.info(f"WebSocket disconnected for session {session_id}")
 
+    def clear_buffer(self, session_id: str) -> None:
+        """Clear the event buffer for a session (e.g. on session delete)."""
+        self._event_buffers.pop(session_id, None)
+
     async def send_event(
         self, session_id: str, event_type: str, data: dict[str, Any] | None = None
     ) -> None:
-        """Send an event to a specific session's WebSocket."""
-        if session_id not in self.active_connections:
-            logger.warning(f"No active connection for session {session_id}")
-            return
+        """Send an event to a specific session's WebSocket.
 
-        message = {"event_type": event_type}
+        If no WebSocket is connected, the event is buffered so it can be
+        replayed when the client reconnects.
+        """
+        message: dict[str, Any] = {"event_type": event_type}
         if data is not None:
             message["data"] = data
 
+        if session_id not in self.active_connections:
+            buf = self._event_buffers.setdefault(
+                session_id, deque(maxlen=MAX_EVENT_BUFFER)
+            )
+            buf.append(message)
+            return
+
         try:
             await self.active_connections[session_id].send_json(message)
         except Exception as e:
             logger.error(f"Error sending to session {session_id}: {e}")
             self.disconnect(session_id)
+            # Buffer the event that failed to send
+            buf = self._event_buffers.setdefault(
+                session_id, deque(maxlen=MAX_EVENT_BUFFER)
+            )
+            buf.append(message)
 
     async def broadcast(
         self, event_type: str, data: dict[str, Any] | None = None

frontend/src/lib/ws-chat-transport.ts

@@ -402,7 +402,13 @@ export class WebSocketChatTransport implements ChatTransport<UIMessage> {
     switch (event.event_type) {
       // -- Side-channel only events ----------------------------------------
       case 'ready':
-        this.sideChannel.onReady();
+        if (this.streamController) {
+          // Reconnect during an active turn — don't reset processing state,
+          // just signal that the connection is back.
+          this.sideChannel.onConnectionChange(true);
+        } else {
+          this.sideChannel.onReady();
+        }
         break;
 
       case 'shutdown':