import React, { createContext, useContext, useEffect, useLayoutEffect, useState, useCallback, useMemo, useRef, } from "react"; import { ConversationClient } from "@openhands/typescript-client/clients"; import { useQueryClient } from "@tanstack/react-query"; import { useWebSocket, WebSocketHookOptions } from "#/hooks/use-websocket"; import { SERVER_CONNECTION_ERROR_MESSAGE } from "#/constants/server-connection-error"; import { useEventStore } from "#/stores/use-event-store"; import { useErrorMessageStore } from "#/stores/error-message-store"; import { useOptimisticUserMessageStore } from "#/stores/optimistic-user-message-store"; import { useConversationStateStore } from "#/stores/conversation-state-store"; import { useCommandStore } from "#/stores/command-store"; import { useBrowserStore } from "#/stores/browser-store"; import { useGoalStore } from "#/stores/goal-store"; import { isAgentServerEvent, isAgentErrorEvent, isUserMessageEvent, isActionEvent, isConversationStateUpdateEvent, isFullStateConversationStateUpdateEvent, isAgentStatusConversationStateUpdateEvent, isStatsConversationStateUpdateEvent, isGoalConversationStateUpdateEvent, isExecuteBashActionEvent, isExecuteBashObservationEvent, isDisplayableErrorEvent, isPlanningFileEditorObservationEvent, isBrowserObservationEvent, isBrowserNavigateActionEvent, isSwitchLLMObservationEvent, isCanvasUIActionEvent, isStreamingDeltaEvent, isLaunchChildConversationActionEvent, } from "#/types/agent-server/type-guards"; import { createStreamingDeltaBatcher, StreamingDeltaBatcher, } from "#/utils/streaming-delta-batcher"; import { handleCanvasUIAction } from "#/services/canvas-ui"; import { handleLaunchChildConversationAction } from "#/services/child-conversation-launch"; import { ConversationStateUpdateEventStats } from "#/types/agent-server/core/events/conversation-state-event"; import type { ConversationErrorEvent, ServerErrorEvent, } from "#/types/agent-server/core/events/conversation-state-event"; import { handleActionEventCacheInvalidation } from "#/utils/cache-utils"; import { buildWebSocketUrl } from "#/utils/websocket-url"; import type { AppConversation, SendMessageRequest, } from "#/api/conversation-service/agent-server-conversation-service.types"; import EventService from "#/api/event-service/event-service.api"; import { getAgentServerClientOptions } from "#/api/agent-server-client-options"; import { useConversationStore } from "#/stores/conversation-store"; import { trackError } from "#/utils/error-handler"; import { useReadConversationFile } from "#/hooks/mutation/use-read-conversation-file"; import useMetricsStore, { type MetricsState } from "#/stores/metrics-store"; import { useConversationHistory } from "#/hooks/query/use-conversation-history"; import { setConversationState } from "#/utils/conversation-local-storage"; import { recordModelSwitchMessage, seedModelSwitchesFromHistory, stampActiveLlmProfile, } from "#/hooks/chat/record-model-switch-message"; import { invalidateConversationQueries, updateConversationLlmModelInCache, } from "#/hooks/mutation/conversation-mutation-utils"; import { isPlanFilePath } from "#/utils/plan-file"; export type WebSocketConnectionState = | "CONNECTING" | "OPEN" | "CLOSED" | "CLOSING"; interface SendMessageResult { queued: boolean; // true if message was queued for later delivery, false if sent immediately } interface ConversationWebSocketContextType { connectionState: WebSocketConnectionState; /** * The main connection's own state, unmerged with the planning connection — * see `useMainWebSocketStatus`. */ mainConnectionState: WebSocketConnectionState; sendMessage: (message: SendMessageRequest) => Promise; isLoadingHistory: boolean; reconnect: () => void; } const ConversationWebSocketContext = createContext< ConversationWebSocketContextType | undefined >(undefined); /** * Extract the text body of an echoed user `MessageEvent` for matching against * the optimistic pending-message queue. The server wraps the original * `args.content` string in one or more `TextContent` entries (alongside any * `ImageContent` entries for inline images), so concatenating the `text` * fields gives us back the exact prompt we sent. */ function extractMessageEventText( event: import("#/types/agent-server/core/events/message-event").MessageEvent, ): string { return event.llm_message.content .filter( (part): part is { type: "text"; text: string } => part.type === "text", ) .map((part) => part.text) .join(""); } export function ConversationWebSocketProvider({ children, conversationId, conversationUrl, sessionApiKey, subConversations, subConversationIds, }: { children: React.ReactNode; conversationId?: string; conversationUrl?: string | null; sessionApiKey?: string | null; subConversations?: AppConversation[]; subConversationIds?: string[]; }) { // Separate connection state tracking for each WebSocket const [mainConnectionState, setMainConnectionState] = useState("CONNECTING"); const [planningConnectionState, setPlanningConnectionState] = useState("CONNECTING"); // Track if we've ever successfully connected for each connection // Don't show errors until after first successful connection const hasConnectedRefMain = React.useRef(false); const hasConnectedRefPlanning = React.useRef(false); const queryClient = useQueryClient(); const addEvent = useEventStore((state) => state.addEvent); const addEvents = useEventStore((state) => state.addEvents); const clearEventsForConversation = useEventStore( (state) => state.clearEventsForConversation, ); const { setErrorMessage, removeErrorMessage, clearConnectionError } = useErrorMessageStore(); const consumeMatchingPendingMessage = useOptimisticUserMessageStore( (state) => state.consumeMatchingPendingMessage, ); const { setExecutionStatus } = useConversationStateStore(); const { appendInput, appendOutput } = useCommandStore(); const resetBrowserStore = useBrowserStore((state) => state.reset); // Coalesce streaming deltas to ≤1 store commit/render per frame. // Separate batchers keep the main and planning streams from ever merging. const mainDeltaBatcherRef = useRef(null); if (mainDeltaBatcherRef.current === null) { mainDeltaBatcherRef.current = createStreamingDeltaBatcher((delta) => { useEventStore.getState().addEvent(delta); // A delta means connectivity recovered — mirror handleNonErrorEvent. useErrorMessageStore.getState().clearConnectionError(); }); } const planningDeltaBatcherRef = useRef(null); if (planningDeltaBatcherRef.current === null) { planningDeltaBatcherRef.current = createStreamingDeltaBatcher((delta) => { useEventStore .getState() .addEvent({ ...delta, isFromPlanningAgent: true }); useErrorMessageStore.getState().clearConnectionError(); }); } // History loading state. // - Main conversation history is now loaded via REST (`useConversationHistory`), // so its loading state mirrors the REST query state (see below). // - Planning sub-conversation history still streams over the WebSocket using // `resend_mode='all'`, so we keep the count-based detection for it. const [isLoadingHistoryPlanning, setIsLoadingHistoryPlanning] = useState(true); const [expectedEventCountPlanning, setExpectedEventCountPlanning] = useState< number | null >(null); const { setPlanContent } = useConversationStore(); useEffect(() => { setPlanContent(null); }, [conversationId, setPlanContent]); const { mutate: readConversationFile } = useReadConversationFile(); // Track planning-agent received events (still WS-driven). const receivedEventCountRefPlanning = useRef(0); // Track the latest PlanningFileEditorObservation for Plan.md during history replay const latestPlanningFileEventRef = useRef<{ path: string; conversationId: string; } | null>(null); const handleNonErrorEvent = useCallback(() => { // A normal event means connectivity recovered: clear a transient connection // error, but keep sticky conversation errors (e.g. a wrong API key). clearConnectionError(); }, [clearConnectionError]); // Helper function to update metrics from stats event const updateMetricsFromStats = useCallback( (event: ConversationStateUpdateEventStats) => { const usageToMetrics = event.value.usage_to_metrics; if (!usageToMetrics) { return; } // usage_to_metrics is keyed by arbitrary LLM usage ids ("default", // "condenser", "profile::", …) — combine across all of // them, mirroring getCombinedMetrics on the REST path. const combined = Object.values(usageToMetrics).reduce<{ cost: number; maxBudgetPerTask: number | null; usage: MetricsState["usage"]; }>( (acc, metrics) => { acc.cost += metrics.accumulated_cost; if ( acc.maxBudgetPerTask === null && metrics.max_budget_per_task !== null ) { acc.maxBudgetPerTask = metrics.max_budget_per_task; } const tokenUsage = metrics.accumulated_token_usage; if (tokenUsage) { acc.usage = { prompt_tokens: (acc.usage?.prompt_tokens ?? 0) + tokenUsage.prompt_tokens, completion_tokens: (acc.usage?.completion_tokens ?? 0) + tokenUsage.completion_tokens, cache_read_tokens: (acc.usage?.cache_read_tokens ?? 0) + tokenUsage.cache_read_tokens, cache_write_tokens: (acc.usage?.cache_write_tokens ?? 0) + tokenUsage.cache_write_tokens, context_window: Math.max( acc.usage?.context_window ?? 0, tokenUsage.context_window, ), per_turn_token: Math.max( acc.usage?.per_turn_token ?? 0, tokenUsage.per_turn_token, ), }; } return acc; }, { cost: 0, maxBudgetPerTask: null, usage: null }, ); useMetricsStore.getState().setMetrics({ cost: combined.cost, max_budget_per_task: combined.maxBudgetPerTask, usage: combined.usage, }); }, [], ); // Initial REST history load: fetch the most recent events and seed the // store. Older events are paginated in via `useLoadOlderEvents` when the // user scrolls to the top of the chat. The WebSocket connection waits for // this query so it can subscribe with `resend_mode='since'` and avoid // re-streaming everything REST already returned. const { data: preloadedHistory, isPending: isPreloadingHistory } = useConversationHistory(conversationId); // Skeleton only on the genuine first load (no cached data yet). On return the // cached page is present, so `isPending` is false and we render the // last-known discussion immediately while the tail refetch runs in the // background — the socket gate below also keys on `isPending`, so that // refetch never drops a live socket. const isLoadingHistoryMain = !!conversationId && isPreloadingHistory; // Clear the (global, not conversation-scoped) event store when the active // conversation changes, BEFORE the preloaded-history effect below re-seeds // it. This MUST live here rather than in the route component: a parent's // passive effect runs *after* this child's layout effects, so clearing from // the route would wipe the freshly seeded history. On a conversation switch // the history page is already cached, so `preloadedHistory` is available // synchronously — without ordering the clear first, the user's already-echoed // message gets seeded then immediately wiped, leaving only the `since` // WebSocket resend (the agent's reply). Re-entering the same conversation is // a no-op, so the store survives navigating away to Settings and back. useLayoutEffect(() => { const nextId = conversationId ?? null; if (useEventStore.getState().loadedConversationId === nextId) { return; } // Single atomic action: clears the previous conversation's events and // records the new loaded id in one `set`, so no subscriber can observe a // half-applied state (events gone but the old id still reported). clearEventsForConversation(nextId); resetBrowserStore(); // The metrics store is conversation-scoped state too: without a reset the // previous conversation's usage/cost keeps rendering in the new // conversation's meter until fresh WS stats arrive — and a brand-new // conversation sends none, so the stale figure stuck indefinitely. useMetricsStore.getState().resetMetrics(); }, [conversationId, clearEventsForConversation, resetBrowserStore]); useLayoutEffect(() => { if (!preloadedHistory || preloadedHistory.events.length === 0) { return; } addEvents(preloadedHistory.events); // The first user message of a cloud start-task conversation is persisted // server-side and reaches us via this REST preload, not over the WebSocket // (which subscribes with resend_mode='since' after the latest preloaded // timestamp). Consume any matching optimistic "Sending…" bubble here too — // mirroring the WS handler — so it doesn't linger as a duplicate of the echo. if (conversationId) { // Rebuild inline "Switched to" messages from the REST-preloaded history. // The live store writers (WS handler / user action) never see preloaded // events, so without this past model switches wouldn't render on reload. // Read the post-`addEvents` `uiEvents` (actions replaced by observations, // Think/Finish observations dropped) — not the raw history — so anchors // match the ids the renderer actually mounts. seedModelSwitchesFromHistory( conversationId, useEventStore.getState().uiEvents, ); for (const event of preloadedHistory.events) { if (isUserMessageEvent(event)) { consumeMatchingPendingMessage( conversationId, extractMessageEventText(event), ); } } } }, [ preloadedHistory, addEvents, conversationId, consumeMatchingPendingMessage, ]); /** * Timestamp of the latest event we already have from REST. Used as * `after_timestamp` when opening the WebSocket so the server only resends * events strictly after this point. `null` until the first REST page lands * (the WS connection is gated on that — see `wsUrl` below). During * background refetches `preloadedHistory` keeps the last-known page, so the * anchor holds steady instead of flipping to null; reconnects read the * freshest value from the options ref at connect time. */ const initialAfterTimestamp = useMemo(() => { const events = preloadedHistory?.events ?? []; const latest = events[events.length - 1]; if (!latest || !("timestamp" in latest) || !latest.timestamp) return null; return latest.timestamp; }, [preloadedHistory]); // Build WebSocket URL from props. // // We deliberately wait for the FIRST history load (`isPending`: no data for // this query key yet) before opening the socket, so the WS subscription can // use `resend_mode='since'` with a meaningful `after_timestamp` instead of // falling back to `resend_mode='all'`. The gate is intentionally NOT on // `isFetching`: background refetches (e.g. the `refetchOnMount` fired when // returning to a conversation) must never tear a live socket down — on a // flaky link that caused a refetch → teardown → reconnect → refetch loop // that kept the conversation stuck at "Connecting" for minutes. Connecting // during a background refetch anchors `since` to the cached tail; the // overlap with the refetched page is deduped by the event store and the // `isDuplicateEvent` guards in the message handlers. A query-key reset // (backend swap / new session key) makes `isPending` true again, so a // genuine reset still re-gates. If the initial load errors, `isPending` // flips false and we fall through to connect with `resend_mode='all'` so // the user still sees live events. const wsUrl = useMemo(() => { if (!conversationId || !conversationUrl) { return null; } if (isPreloadingHistory) { return null; } return buildWebSocketUrl(conversationId, conversationUrl); }, [conversationId, conversationUrl, isPreloadingHistory]); // Derived from `subConversationIds` (the pre-filtered, tag-verified id // list) rather than the resolved `subConversations` entry, which lands a // tick later — routing on it would leave a window where the first prompt // sent in plan mode fell through to the parent (the code agent) instead of // the planner. const planningConversationId = useMemo( () => subConversationIds?.[0] ?? null, [subConversationIds], ); const planningAgentWsUrl = useMemo(() => { if (!subConversations?.length) { return null; } // Currently, there is only one sub-conversation and it uses the planning agent. const planningAgentConversation = subConversations[0]; if ( !planningAgentConversation?.id || !planningAgentConversation.conversation_url ) { return null; } return buildWebSocketUrl( planningAgentConversation.id, planningAgentConversation.conversation_url, ); }, [subConversations]); // Merged connection state - reflects combined status of both connections const connectionState = useMemo(() => { // If planning agent connection doesn't exist, use main connection state if (!planningAgentWsUrl) { return mainConnectionState; } // If either is connecting, merged state is connecting if ( mainConnectionState === "CONNECTING" || planningConnectionState === "CONNECTING" ) { return "CONNECTING"; } // If both are open, merged state is open if (mainConnectionState === "OPEN" && planningConnectionState === "OPEN") { return "OPEN"; } // If both are closed, merged state is closed if ( mainConnectionState === "CLOSED" && planningConnectionState === "CLOSED" ) { return "CLOSED"; } // If either is closing, merged state is closing if ( mainConnectionState === "CLOSING" || planningConnectionState === "CLOSING" ) { return "CLOSING"; } // Default to closed if states don't match expected patterns return "CLOSED"; }, [mainConnectionState, planningConnectionState, planningAgentWsUrl]); useEffect(() => { if ( expectedEventCountPlanning !== null && receivedEventCountRefPlanning.current >= expectedEventCountPlanning && isLoadingHistoryPlanning ) { setIsLoadingHistoryPlanning(false); } }, [ expectedEventCountPlanning, isLoadingHistoryPlanning, receivedEventCountRefPlanning, ]); // Call API once after history loading completes if we tracked any PlanningFileEditorObservation events useEffect(() => { if (!isLoadingHistoryPlanning && latestPlanningFileEventRef.current) { const { path, conversationId: currentPlanningConversationId } = latestPlanningFileEventRef.current; readConversationFile( { conversationId: currentPlanningConversationId, filePath: path, }, { onSuccess: (fileContent) => { setPlanContent(fileContent); }, onError: (error) => { console.warn("Failed to read conversation file:", error); }, }, ); // Clear the ref after calling the API latestPlanningFileEventRef.current = null; } }, [isLoadingHistoryPlanning, readConversationFile, setPlanContent]); useEffect(() => { hasConnectedRefMain.current = false; setIsLoadingHistoryPlanning(!!subConversationIds?.length); setExpectedEventCountPlanning(null); receivedEventCountRefPlanning.current = 0; // Reset the tracked event ref when sub-conversations change latestPlanningFileEventRef.current = null; }, [subConversationIds]); // Reset hasConnected flags when the conversation changes. useEffect(() => { hasConnectedRefMain.current = false; hasConnectedRefPlanning.current = false; // Reset the tracked event ref when conversation changes latestPlanningFileEventRef.current = null; }, [conversationId]); // Drop buffered deltas on conversation switch/unmount: the store is cleared on // switch, so flushing them would leak into the next conversation. useEffect(() => { const mainBatcher = mainDeltaBatcherRef.current; const planningBatcher = planningDeltaBatcherRef.current; return () => { mainBatcher?.reset(); planningBatcher?.reset(); }; }, [conversationId]); // Merged loading history state - true if either connection is still loading const isLoadingHistory = useMemo( () => isLoadingHistoryMain || isLoadingHistoryPlanning, [isLoadingHistoryMain, isLoadingHistoryPlanning], ); // Separate message handlers for each connection const handleMainMessage = useCallback( (messageEvent: MessageEvent) => { try { const event = JSON.parse(messageEvent.data); // History loading for the main conversation is REST-driven now; // every WS message is a new event we add to the store. // Use type guard to validate v1 event structure if (isAgentServerEvent(event)) { // Buffer deltas; nothing else in this handler applies to them. if (isStreamingDeltaEvent(event)) { mainDeltaBatcherRef.current?.enqueue(event); return; } // Flush buffered deltas before this event so it can't overtake them. mainDeltaBatcherRef.current?.flush(); // A reconnect replays the backlog from a stale anchor. The store // dedups by id, but the side-effects below aren't idempotent, so skip // them for replayed events (#1656). const isDuplicateEvent = useEventStore .getState() .eventIds.has(event.id ?? ""); const switchLLMObservation = isSwitchLLMObservationEvent(event) ? event : null; addEvent(event); if (isDuplicateEvent) { return; } // Handle displayable error events - show error banner // AgentErrorEvent errors are displayed inline in the chat, not as banners if (isDisplayableErrorEvent(event)) { const errorEvent = event as | ConversationErrorEvent | ServerErrorEvent; const classification = "classification" in errorEvent ? errorEvent.classification : null; trackError({ source: "conversation", metadata: { eventId: errorEvent.id, errorCode: errorEvent.code, }, classification, }); setErrorMessage( errorEvent.detail, "conversation", errorEvent.code, classification, ); } else { handleNonErrorEvent(); } // LLM errors render inline in the chat (see ErrorEventMessage); track // them for analytics but keep them out of the banner above the chat box. if (isAgentErrorEvent(event)) { trackError({ source: "agent", metadata: { eventId: event.id, toolName: event.tool_name, toolCallId: event.tool_call_id, }, classification: event.classification, }); } // Clear optimistic user message when a user message is confirmed. // We match by the echoed text content (with FIFO fallback inside the // store), so an echo for "second" pops "second" — not whichever // pending entry happens to be oldest — protecting against any // out-of-order delivery between conversations or sub-agents. if (isUserMessageEvent(event)) { if (conversationId) { consumeMatchingPendingMessage( conversationId, extractMessageEventText(event), ); // Clear draft from localStorage - message was successfully delivered setConversationState(conversationId, { draftMessage: null }); } } // Handle cache invalidation for ActionEvent if (isActionEvent(event)) { const currentConversationId = conversationId || "test-conversation-id"; // TODO: Get from context handleActionEventCacheInvalidation( event, currentConversationId, queryClient, ); } // Handle conversation state updates // TODO: Tests if (isConversationStateUpdateEvent(event)) { if ( isFullStateConversationStateUpdateEvent(event) && conversationId ) { setExecutionStatus(conversationId, event.value.execution_status); } if ( isAgentStatusConversationStateUpdateEvent(event) && conversationId ) { setExecutionStatus(conversationId, event.value); } if (isStatsConversationStateUpdateEvent(event)) { updateMetricsFromStats(event); } // Mirror goal status into the store. Intentionally duplicated across // the main and planning WebSocket handlers (like the execution_status // and stats branches above), not a merge artifact. if (isGoalConversationStateUpdateEvent(event) && conversationId) { useGoalStore.getState().setStatus(conversationId, event.value); } } // Handle ExecuteBashAction events - add command as input to terminal if (isExecuteBashActionEvent(event)) { appendInput(event.action.command); } // Handle ExecuteBashObservation events - add output to terminal if (isExecuteBashObservationEvent(event)) { // Extract text content from the observation content array const textContent = event.observation.content .filter((c) => c.type === "text") .map((c) => c.text) .join("\n"); appendOutput(textContent); } // Handle BrowserObservation events - update browser store with screenshot if (isBrowserObservationEvent(event)) { const { screenshot_data: screenshotData } = event.observation; if (screenshotData) { const screenshotSrc = screenshotData.startsWith("data:") ? screenshotData : `data:image/png;base64,${screenshotData}`; useBrowserStore.getState().setScreenshotSrc(screenshotSrc); } } // Handle BrowserNavigateAction events - update browser store with URL if (isBrowserNavigateActionEvent(event)) { useBrowserStore.getState().setUrl(event.action.url); } if ( conversationId && switchLLMObservation && !switchLLMObservation.observation.is_error ) { recordModelSwitchMessage( conversationId, switchLLMObservation.observation.profile_name, ); // Mirror the user-driven `/model` path: persist the profile so the // chat-header switcher shows the right name after a reload, even // when several profiles share a model (#1082). Stamp with the // observation's own timestamp so a later history seed of this same // event can't roll it back (or needlessly rewrite it). stampActiveLlmProfile( conversationId, switchLLMObservation.observation.profile_name, switchLLMObservation.timestamp, ); if (switchLLMObservation.observation.active_model) { updateConversationLlmModelInCache( queryClient, conversationId, switchLLMObservation.observation.active_model, ); } invalidateConversationQueries(queryClient, conversationId); } // Handle canvas_ui ActionEvents from both the legacy Python tool and // the client-defined JSON tool. The server acknowledges immediately; // the actual UI change happens here on the client. if (isCanvasUIActionEvent(event)) { handleCanvasUIAction(event.action, conversationId ?? null); } // Same client-tool pattern, but the work is a network call: launch // the requested child conversation and post the outcome back so the // agent learns the id the server-side acknowledgement can't carry. if (conversationId && isLaunchChildConversationActionEvent(event)) { void handleLaunchChildConversationAction( event.action, conversationId, event.tool_call_id, ); } } } catch (error) { console.warn("Failed to parse WebSocket message as JSON:", error); } }, [ addEvent, setErrorMessage, consumeMatchingPendingMessage, queryClient, conversationId, setExecutionStatus, appendInput, appendOutput, updateMetricsFromStats, handleNonErrorEvent, ], ); const handlePlanningMessage = useCallback( (messageEvent: MessageEvent) => { try { const event = JSON.parse(messageEvent.data); // Track received events for history loading (count ALL events from WebSocket) // Always count when loading, even if we don't have the expected count yet if (isLoadingHistoryPlanning) { receivedEventCountRefPlanning.current += 1; if ( expectedEventCountPlanning !== null && receivedEventCountRefPlanning.current >= expectedEventCountPlanning ) { setIsLoadingHistoryPlanning(false); } } // Use type guard to validate v1 event structure if (isAgentServerEvent(event)) { // Buffer deltas (the commit re-applies the planning flag). if (isStreamingDeltaEvent(event)) { planningDeltaBatcherRef.current?.enqueue(event); return; } // Flush buffered deltas before this event so it can't overtake them. planningDeltaBatcherRef.current?.flush(); // Skip non-idempotent side-effects for replayed events, as in the // main handler (#1656). const isDuplicateEvent = useEventStore .getState() .eventIds.has(event.id ?? ""); // Mark this event as coming from the planning agent const eventWithPlanningFlag = { ...event, isFromPlanningAgent: true, }; addEvent(eventWithPlanningFlag); if (isDuplicateEvent) { return; } // Handle displayable error events - show error banner // AgentErrorEvent errors are displayed inline in the chat, not as banners if (isDisplayableErrorEvent(event)) { const errorEvent = event as | ConversationErrorEvent | ServerErrorEvent; const classification = "classification" in errorEvent ? errorEvent.classification : null; trackError({ source: "planning_conversation", metadata: { eventId: errorEvent.id, errorCode: errorEvent.code, }, classification, }); setErrorMessage( errorEvent.detail, "conversation", errorEvent.code, classification, ); } else { handleNonErrorEvent(); } // LLM errors render inline in the chat (see ErrorEventMessage); track // them for analytics but keep them out of the banner above the chat box. if (isAgentErrorEvent(event)) { trackError({ source: "planning_agent", metadata: { eventId: event.id, toolName: event.tool_name, toolCallId: event.tool_call_id, }, classification: event.classification, }); } // Clear optimistic user message when a user message is confirmed. // Always scope to the main `conversationId` (where the user types) // and match on the echoed content so the planning sub-agent's own // events can never consume a main-conversation pending entry. if (isUserMessageEvent(event)) { if (conversationId) { consumeMatchingPendingMessage( conversationId, extractMessageEventText(event), ); setConversationState(conversationId, { draftMessage: null }); } } // Handle cache invalidation for ActionEvent if (isActionEvent(event)) { const planningAgentConversation = subConversations?.[0]; const currentConversationId = planningAgentConversation?.id || "test-conversation-id"; // TODO: Get from context handleActionEventCacheInvalidation( event, currentConversationId, queryClient, ); } // Handle conversation state updates // TODO: Tests if (isConversationStateUpdateEvent(event)) { // Scope to the planning agent's own conversation id, not the main // `conversationId` — this socket reports the planning helper // conversation's run/idle transitions, which must never overwrite // the main conversation's status in the shared store. if ( isFullStateConversationStateUpdateEvent(event) && planningConversationId ) { setExecutionStatus( planningConversationId, event.value.execution_status, ); } if ( isAgentStatusConversationStateUpdateEvent(event) && planningConversationId ) { setExecutionStatus(planningConversationId, event.value); } if (isStatsConversationStateUpdateEvent(event)) { updateMetricsFromStats(event); } // Mirror goal status into the store. Intentionally duplicated across // the main and planning WebSocket handlers (like the execution_status // and stats branches above), not a merge artifact. if (isGoalConversationStateUpdateEvent(event) && conversationId) { useGoalStore.getState().setStatus(conversationId, event.value); } } // Handle ExecuteBashAction events - add command as input to terminal if (isExecuteBashActionEvent(event)) { appendInput(event.action.command); } // Handle ExecuteBashObservation events - add output to terminal if (isExecuteBashObservationEvent(event)) { // Extract text content from the observation content array const textContent = event.observation.content .filter((c) => c.type === "text") .map((c) => c.text) .join("\n"); appendOutput(textContent); } // Handle PlanningFileEditorObservation - only update plan for Plan.md if (isPlanningFileEditorObservationEvent(event)) { const { path } = event.observation; if (isPlanFilePath(path)) { const planningAgentConversation = subConversations?.[0]; const planningConversationId = planningAgentConversation?.id; if (planningConversationId && path) { if (isLoadingHistoryPlanning) { latestPlanningFileEventRef.current = { path, conversationId: planningConversationId, }; } else { readConversationFile( { conversationId: planningConversationId, filePath: path, }, { onSuccess: (fileContent) => { setPlanContent(fileContent); }, onError: (error) => { console.warn( "Failed to read conversation file:", error, ); }, }, ); } } } } } } catch (error) { console.warn("Failed to parse WebSocket message as JSON:", error); } }, [ addEvent, isLoadingHistoryPlanning, expectedEventCountPlanning, setErrorMessage, consumeMatchingPendingMessage, queryClient, subConversations, planningConversationId, conversationId, setExecutionStatus, appendInput, appendOutput, readConversationFile, setPlanContent, updateMetricsFromStats, handleNonErrorEvent, ], ); // Separate WebSocket options for main connection const mainWebsocketOptions: WebSocketHookOptions = useMemo(() => { // History was already loaded over REST (`useConversationHistory`). // Subscribe with `resend_mode='since'` so the server only resends events // strictly after the latest one we already have. If REST returned no // events at all (brand-new conversation), fall back to `'all'` so any // events that may have been written between the REST call and the WS // handshake still show up. Dedup in the event store handles overlap. const queryParams: Record = initialAfterTimestamp ? { resend_mode: "since", after_timestamp: initialAfterTimestamp } : { resend_mode: "all" }; return { queryParams, sessionApiKey, reconnect: { enabled: true }, onOpen: () => { setMainConnectionState("OPEN"); hasConnectedRefMain.current = true; // Mark that we've successfully connected clearConnectionError(); // Clear a previous connection error; keep sticky conversation errors }, onClose: () => { setMainConnectionState("CLOSED"); }, onError: () => { setMainConnectionState("CLOSED"); // Only show error message if we've previously connected successfully if (hasConnectedRefMain.current) { setErrorMessage(SERVER_CONNECTION_ERROR_MESSAGE, "connection"); } }, onMessage: handleMainMessage, }; }, [ handleMainMessage, setErrorMessage, clearConnectionError, sessionApiKey, initialAfterTimestamp, ]); // Separate WebSocket options for planning agent connection const planningWebsocketOptions: WebSocketHookOptions = useMemo(() => { const queryParams: Record = { resend_all: true, }; const planningAgentConversation = subConversations?.[0]; const planningApiKey = planningAgentConversation?.session_api_key ?? sessionApiKey; return { queryParams, sessionApiKey: planningApiKey, reconnect: { enabled: true }, onOpen: async () => { setPlanningConnectionState("OPEN"); hasConnectedRefPlanning.current = true; // Mark that we've successfully connected clearConnectionError(); // Clear a previous connection error; keep sticky conversation errors // Fetch expected event count for history loading detection if ( planningAgentConversation?.id && planningAgentConversation.conversation_url ) { try { const count = await EventService.getEventCount( planningAgentConversation.id, planningAgentConversation.conversation_url, planningAgentConversation.session_api_key, ); setExpectedEventCountPlanning(count); // If no events expected, mark as loaded immediately if (count === 0) { setIsLoadingHistoryPlanning(false); } } catch (error) { // Fall back to marking as loaded to avoid infinite loading state setIsLoadingHistoryPlanning(false); } } }, onClose: () => { setPlanningConnectionState("CLOSED"); }, onError: () => { setPlanningConnectionState("CLOSED"); // Only show error message if we've previously connected successfully if (hasConnectedRefPlanning.current) { setErrorMessage(SERVER_CONNECTION_ERROR_MESSAGE, "connection"); } }, onMessage: handlePlanningMessage, }; }, [ handlePlanningMessage, setErrorMessage, clearConnectionError, sessionApiKey, subConversations, ]); // Only attempt WebSocket connection when we have a valid URL // This prevents connection attempts during task polling phase const websocketUrl = wsUrl; const { socket: mainSocket, reconnect: reconnectMain } = useWebSocket( websocketUrl || "", mainWebsocketOptions, ); const { socket: planningAgentSocket, reconnect: reconnectPlanning } = useWebSocket(planningAgentWsUrl || "", planningWebsocketOptions); const reconnect = useCallback(() => { removeErrorMessage(); const currentMode = useConversationStore.getState().conversationMode; if (currentMode === "plan" && planningAgentWsUrl) { reconnectPlanning(); return; } reconnectMain(); }, [ planningAgentWsUrl, reconnectMain, reconnectPlanning, removeErrorMessage, ]); // V1 send message function via WebSocket // Falls back to REST API queue when WebSocket is not connected const sendMessage = useCallback( async (message: SendMessageRequest): Promise => { const currentMode = useConversationStore.getState().conversationMode; const currentSocket = currentMode === "plan" ? planningAgentSocket : mainSocket; const targetConversationId = currentMode === "plan" ? planningConversationId : conversationId; if (currentSocket?.readyState !== WebSocket.OPEN) { // WebSocket not connected - queue message via REST API // Message will be delivered automatically when conversation becomes ready if (!targetConversationId) { // Never fall back to the parent in plan mode: without a planner // target the message would run in the code agent, which is exactly // the boundary plan mode exists to enforce. const error = new Error( currentMode === "plan" ? "Planning conversation is not ready yet" : "No conversation ID available", ); setErrorMessage(error.message); throw error; } try { await new ConversationClient(getAgentServerClientOptions()).sendEvent( targetConversationId, { role: "user", content: message.content, }, { run: true }, ); // Message queued successfully - it will be delivered when ready // Return queued: true so caller knows not to show optimistic UI return { queued: true }; } catch (error) { const errorMessage = error instanceof Error ? error.message : "Failed to queue message for delivery"; setErrorMessage(errorMessage); throw error; } } try { // Send message through WebSocket as JSON with run: true so the // agent loop starts automatically in async mode. currentSocket.send(JSON.stringify({ ...message, run: true })); return { queued: false }; } catch (error) { const errorMessage = error instanceof Error ? error.message : "Failed to send message"; setErrorMessage(errorMessage); throw error; } }, [ mainSocket, planningAgentSocket, setErrorMessage, conversationId, planningConversationId, ], ); // Track main socket state changes useEffect(() => { // Only process socket updates if we have a valid URL and socket if (mainSocket && wsUrl) { // Update state based on socket readyState const updateState = () => { switch (mainSocket.readyState) { case WebSocket.CONNECTING: setMainConnectionState("CONNECTING"); break; case WebSocket.OPEN: setMainConnectionState("OPEN"); break; case WebSocket.CLOSING: setMainConnectionState("CLOSING"); break; case WebSocket.CLOSED: setMainConnectionState("CLOSED"); break; default: setMainConnectionState("CLOSED"); break; } }; updateState(); } }, [mainSocket, wsUrl]); // Track planning agent socket state changes useEffect(() => { // Only process socket updates if we have a valid URL and socket if (planningAgentSocket && planningAgentWsUrl) { // Update state based on socket readyState const updateState = () => { switch (planningAgentSocket.readyState) { case WebSocket.CONNECTING: setPlanningConnectionState("CONNECTING"); break; case WebSocket.OPEN: setPlanningConnectionState("OPEN"); break; case WebSocket.CLOSING: setPlanningConnectionState("CLOSING"); break; case WebSocket.CLOSED: setPlanningConnectionState("CLOSED"); break; default: setPlanningConnectionState("CLOSED"); break; } }; updateState(); } }, [planningAgentSocket, planningAgentWsUrl]); const contextValue = useMemo( () => ({ connectionState, mainConnectionState, sendMessage, isLoadingHistory, reconnect, }), [ connectionState, mainConnectionState, sendMessage, isLoadingHistory, reconnect, ], ); return ( {children} ); } export const useConversationWebSocket = (): ConversationWebSocketContextType | null => { const context = useContext(ConversationWebSocketContext); // Return null instead of throwing when not in provider // This allows the hook to be called conditionally based on conversation version return context || null; };