Merge pull request microsoft#1086 from microsoft/feat/model-changes

fix: Intermittent UI freeze and infinite reasoning

Author: Roopan-Microsoft

src/App/src/hooks/usePlanWebSocket.tsx

@@ -105,6 +105,11 @@ export function usePlanWebSocket({
     const streamingMessageBuffer = useAppSelector(selectStreamingMessageBuffer);
     const processingStartedAtRef = React.useRef<number | null>(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<string[]>([]);
+    const streamingFlushHandleRef = React.useRef<number | null>(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();
                 }
             },
         );

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 {

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<typeof setInterval> | 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;

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)
 
 

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)