Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
55 changes: 52 additions & 3 deletions src/App/src/hooks/usePlanWebSocket.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -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) {
Expand Down Expand Up @@ -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 ────────────────────────────────
Expand Down Expand Up @@ -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();
}
},
);
Expand Down
3 changes: 2 additions & 1 deletion src/App/src/models/enums.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down
29 changes: 29 additions & 0 deletions src/App/src/store/WebSocketService.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down Expand Up @@ -59,6 +61,7 @@ class WebSocketService {
clearTimeout(this.reconnectTimer);
this.reconnectTimer = null;
}
this.startHeartbeat();
this.emit('connection_status', { connected: true });
resolve();
};
Expand All @@ -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 &&
Expand All @@ -99,6 +103,7 @@ class WebSocketService {

disconnect(): void {
this.intentionalDisconnect = true;
this.stopHeartbeat();
if (this.reconnectTimer) {
clearTimeout(this.reconnectTimer);
this.reconnectTimer = null;
Expand Down Expand Up @@ -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());
Expand Down Expand Up @@ -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;
Expand Down
34 changes: 32 additions & 2 deletions src/backend/api/router.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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)


Expand Down
1 change: 1 addition & 0 deletions src/backend/models/messages.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
Loading