diff --git a/src/App/src/hooks/usePlanWebSocket.tsx b/src/App/src/hooks/usePlanWebSocket.tsx index 7b7895362..7bac625c1 100644 --- a/src/App/src/hooks/usePlanWebSocket.tsx +++ b/src/App/src/hooks/usePlanWebSocket.tsx @@ -105,6 +105,11 @@ export function usePlanWebSocket({ const streamingMessageBuffer = useAppSelector(selectStreamingMessageBuffer); const processingStartedAtRef = React.useRef(null); + // Coalesce high-frequency streaming tokens into one flush per animation frame + // to avoid a synchronous re-render per token freezing the UI on fast streams. + const streamingChunkQueueRef = React.useRef([]); + const streamingFlushHandleRef = React.useRef(null); + useEffect(() => { if (showProcessingPlanSpinner) { if (processingStartedAtRef.current === null) { @@ -144,15 +149,38 @@ export function usePlanWebSocket({ // ── AGENT_MESSAGE_STREAMING ─────────────────────────────────── useEffect(() => { + const flushStreamingChunks = () => { + streamingFlushHandleRef.current = null; + const chunks = streamingChunkQueueRef.current; + if (chunks.length === 0) return; + streamingChunkQueueRef.current = []; + dispatch(setShowBufferingText(true)); + dispatch(appendToStreamingBuffer(chunks.join(''))); + }; + const unsub = webSocketService.on( WebsocketMessageType.AGENT_MESSAGE_STREAMING, (msg: any) => { const line = PlanDataService.simplifyHumanClarification(msg.data?.content || msg.content || ''); - dispatch(setShowBufferingText(true)); - dispatch(appendToStreamingBuffer(line)); + streamingChunkQueueRef.current.push(line); + if (streamingFlushHandleRef.current === null) { + streamingFlushHandleRef.current = requestAnimationFrame(flushStreamingChunks); + } }, ); - return unsub; + return () => { + unsub(); + // Cancel pending frame and flush leftovers so no streamed text is lost + if (streamingFlushHandleRef.current !== null) { + cancelAnimationFrame(streamingFlushHandleRef.current); + streamingFlushHandleRef.current = null; + } + if (streamingChunkQueueRef.current.length > 0) { + const remaining = streamingChunkQueueRef.current.join(''); + streamingChunkQueueRef.current = []; + dispatch(appendToStreamingBuffer(remaining)); + } + }; }, [dispatch]); // ── USER_CLARIFICATION_REQUEST ──────────────────────────────── @@ -241,6 +269,27 @@ export function usePlanWebSocket({ scrollToBottom(); showToast(errorContent, 'error'); webSocketService.disconnect(); + } else { + // Any other terminal status (e.g. "terminated"): clear the spinner + // so the UI doesn't hang after the answer has already arrived. + const content = finalMessage.data?.content; + if (content) { + const terminalMessage: AgentMessageData = { + agent: AgentType.GROUP_CHAT_MANAGER, + agent_type: AgentMessageType.AI_AGENT, + timestamp: Date.now(), + steps: [], + next_steps: [], + content, + raw_data: finalMessage, + }; + dispatch(addAgentMessage(terminalMessage)); + } + dispatch(setShowBufferingText(false)); + dispatch(setShowProcessingPlanSpinner(false)); + processingStartedAtRef.current = null; + scrollToBottom(); + webSocketService.disconnect(); } }, ); diff --git a/src/App/src/models/enums.tsx b/src/App/src/models/enums.tsx index 6f0deac9a..33fc95d14 100644 --- a/src/App/src/models/enums.tsx +++ b/src/App/src/models/enums.tsx @@ -253,7 +253,8 @@ export enum WebsocketMessageType { USER_CLARIFICATION_REQUEST = "user_clarification_request", USER_CLARIFICATION_RESPONSE = "user_clarification_response", FINAL_RESULT_MESSAGE = "final_result_message", - ERROR_MESSAGE = 'error_message' + ERROR_MESSAGE = 'error_message', + PING = "ping" } export enum AgentMessageType { diff --git a/src/App/src/store/WebSocketService.tsx b/src/App/src/store/WebSocketService.tsx index 7fff80a56..c87a796b7 100644 --- a/src/App/src/store/WebSocketService.tsx +++ b/src/App/src/store/WebSocketService.tsx @@ -15,6 +15,8 @@ class WebSocketService { private intentionalDisconnect = false; private lastPlanId: string | undefined; private lastProcessId: string | undefined; + private heartbeatTimer: ReturnType | null = null; + private heartbeatIntervalMs = 20000; // 20s client keepalive ping private buildSocketUrl(processId?: string, planId?: string): string { @@ -59,6 +61,7 @@ class WebSocketService { clearTimeout(this.reconnectTimer); this.reconnectTimer = null; } + this.startHeartbeat(); this.emit('connection_status', { connected: true }); resolve(); }; @@ -75,6 +78,7 @@ class WebSocketService { this.ws.onclose = (event) => { this.isConnecting = false; this.ws = null; + this.stopHeartbeat(); this.emit('connection_status', { connected: false }); /* P1: Only auto-reconnect if not intentional and not a clean close */ if (!this.intentionalDisconnect && event.code !== 1000 && @@ -99,6 +103,7 @@ class WebSocketService { disconnect(): void { this.intentionalDisconnect = true; + this.stopHeartbeat(); if (this.reconnectTimer) { clearTimeout(this.reconnectTimer); this.reconnectTimer = null; @@ -130,6 +135,26 @@ class WebSocketService { } + private startHeartbeat(): void { + this.stopHeartbeat(); + this.heartbeatTimer = setInterval(() => { + if (this.ws && this.ws.readyState === WebSocket.OPEN) { + try { + this.ws.send(JSON.stringify({ type: WebsocketMessageType.PING })); + } catch { + /* onclose handles real drops */ + } + } + }, this.heartbeatIntervalMs); + } + + private stopHeartbeat(): void { + if (this.heartbeatTimer) { + clearInterval(this.heartbeatTimer); + this.heartbeatTimer = null; + } + } + on(eventType: string, callback: (message: StreamMessage) => void): () => void { if (!this.listeners.has(eventType)) { this.listeners.set(eventType, new Set()); @@ -255,6 +280,10 @@ class WebSocketService { } break; } + case WebsocketMessageType.PING: { + // Server keepalive heartbeat — ignore. + break; + } case WebsocketMessageType.ERROR_MESSAGE: { this.emit(WebsocketMessageType.ERROR_MESSAGE, message.data); // Emit the data break; diff --git a/src/backend/api/router.py b/src/backend/api/router.py index dbf11afe6..f683e2dab 100644 --- a/src/backend/api/router.py +++ b/src/backend/api/router.py @@ -72,6 +72,33 @@ async def start_comms( ws_props["session_id"] = session_id track_event_if_configured("WebSocket_Connected", ws_props) + # Keepalive: reasoning models (gpt-5.4/-mini) stream nothing during long + # thinking gaps; a periodic frame stops the ingress proxy idle-timing-out + # and dropping the socket (which would lose the final-result message). + HEARTBEAT_INTERVAL_SECONDS = 20 + + async def _heartbeat() -> None: + while True: + await asyncio.sleep(HEARTBEAT_INTERVAL_SECONDS) + try: + await websocket.send_text( + json.dumps( + { + "type": WebsocketMessageType.PING, + "data": {"ts": asyncio.get_event_loop().time()}, + }, + default=str, + ) + ) + except Exception as hb_exc: + logging.debug( + "Heartbeat stopped for user %s, process %s: %s", + user_id, process_id, hb_exc, + ) + break + + heartbeat_task = asyncio.create_task(_heartbeat()) + # Keep the connection open - FastAPI will close the connection if this returns try: # Keep the connection open - FastAPI will close the connection if this returns @@ -95,10 +122,13 @@ async def start_comms( logging.info(f"Client disconnected from batch {process_id}") break except Exception as e: - # Fixed logging syntax - removed the error= parameter logging.error(f"Error in WebSocket connection: {str(e)}") finally: - # Always clean up the connection + heartbeat_task.cancel() + try: + await heartbeat_task + except (asyncio.CancelledError, Exception): + pass await connection_config.close_connection(process_id=process_id) diff --git a/src/backend/models/messages.py b/src/backend/models/messages.py index e774b648a..501366c7c 100644 --- a/src/backend/models/messages.py +++ b/src/backend/models/messages.py @@ -152,6 +152,7 @@ class WebsocketMessageType(str, Enum): FINAL_RESULT_MESSAGE = "final_result_message" TIMEOUT_NOTIFICATION = "timeout_notification" ERROR_MESSAGE = "error_message" + PING = "ping" @dataclass(slots=True)