diff --git a/frontend/src/context/ws-client-provider.tsx b/frontend/src/context/ws-client-provider.tsx index bfecefadd00a..97aad6ca6061 100644 --- a/frontend/src/context/ws-client-provider.tsx +++ b/frontend/src/context/ws-client-provider.tsx @@ -5,6 +5,8 @@ import ActionType from "#/types/ActionType"; import EventLogger from "#/utils/event-logger"; import AgentState from "#/types/AgentState"; +const RECONNECT_RETRIES = 5; + export enum WsClientProviderStatus { STOPPED, OPENING, @@ -46,6 +48,7 @@ export function WsClientProvider({ const closeRef = React.useRef | null>(null); const [status, setStatus] = React.useState(WsClientProviderStatus.STOPPED); const [events, setEvents] = React.useState[]>([]); + const [retryCount, setRetryCount] = React.useState(RECONNECT_RETRIES); function send(event: Record) { if (!wsRef.current) { @@ -56,6 +59,7 @@ export function WsClientProvider({ } function handleOpen() { + setRetryCount(RECONNECT_RETRIES); setStatus(WsClientProviderStatus.OPENING); const initEvent = { action: ActionType.INIT, @@ -79,8 +83,14 @@ export function WsClientProvider({ } function handleClose() { - setStatus(WsClientProviderStatus.STOPPED); - setEvents([]); + if (retryCount) { + setTimeout(() => { + setRetryCount(retryCount - 1); + }, 1000); + } else { + setStatus(WsClientProviderStatus.STOPPED); + setEvents([]); + } wsRef.current = null; } @@ -95,7 +105,7 @@ export function WsClientProvider({ let ws = wsRef.current; // If disabled close any existing websockets... - if (!enabled) { + if (!enabled || !retryCount) { if (ws) { ws.close(); } @@ -116,7 +126,11 @@ export function WsClientProvider({ const baseUrl = import.meta.env.VITE_BACKEND_BASE_URL || window?.location.host; const protocol = window.location.protocol === "https:" ? "wss:" : "ws:"; - ws = new WebSocket(`${protocol}//${baseUrl}/ws`, [ + let wsUrl = `${protocol}//${baseUrl}/ws`; + if (events.length) { + wsUrl += `?latest_event_id=${events[events.length - 1].id}`; + } + ws = new WebSocket(wsUrl, [ "openhands", token || "NO_JWT", ghToken || "NO_GITHUB", @@ -136,7 +150,7 @@ export function WsClientProvider({ ws.removeEventListener("error", handleError); ws.removeEventListener("close", handleClose); }; - }, [enabled, token, ghToken]); + }, [enabled, token, ghToken, retryCount]); // Strict mode mounts and unmounts each component twice, so we have to wait in the destructor // before actually closing the socket and cancel the operation if the component gets remounted. @@ -148,7 +162,11 @@ export function WsClientProvider({ return () => { closeRef.current = setTimeout(() => { - wsRef.current?.close(); + const ws = wsRef.current; + if (ws) { + ws.removeEventListener("close", handleClose); + ws.close(); + } }, 100); }; }, []);