Reconnecting websockets (#4954)

This commit is contained in:
tofarr 2024-11-13 00:38:26 -07:00 committed by GitHub
parent 207df9dd30
commit 87c02177d7
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -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<ReturnType<typeof setTimeout> | null>(null);
const [status, setStatus] = React.useState(WsClientProviderStatus.STOPPED);
const [events, setEvents] = React.useState<Record<string, unknown>[]>([]);
const [retryCount, setRetryCount] = React.useState(RECONNECT_RETRIES);
function send(event: Record<string, unknown>) {
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);
};
}, []);