diff --git a/web/src/lib/api/process.ts b/web/src/lib/api/process.ts index 4bd31ab..89713cc 100644 --- a/web/src/lib/api/process.ts +++ b/web/src/lib/api/process.ts @@ -7,6 +7,8 @@ import { type UseQueryResult, } from "@tanstack/react-query"; +import { useCallback, useEffect, useRef, useState } from "react"; + import { apiClient } from "@/lib/api/client"; export interface ProcessStatus { @@ -78,3 +80,146 @@ export function useProcessStop(): UseMutationResult { export function useProcessRestart(): UseMutationResult { return useProcessActionMutation("restart"); } + +export type ProcessWebSocketStatus = + | "idle" + | "connecting" + | "open" + | "closed" + | "error"; + +export interface ProcessWebSocket { + status: ProcessWebSocketStatus; + send: (data: string) => void; + close: () => void; + onData: (cb: (data: string) => void) => () => void; +} + +export function useProcessWebSocket( + workspaceId: string | null, +): ProcessWebSocket { + const [status, setStatus] = useState(() => + workspaceId ? "connecting" : "idle", + ); + const statusRef = useRef(status); + useEffect(() => { + statusRef.current = status; + }, [status]); + + const socketRef = useRef(null); + const reconnectAttemptsRef = useRef(0); + const reconnectTimeoutRef = useRef(null); + const intentionalCloseRef = useRef(false); + const lastFrameWasExitRef = useRef(false); + const callbacksRef = useRef(new Set<(data: string) => void>()); + + const clearReconnectTimeout = useCallback(() => { + if (reconnectTimeoutRef.current !== null) { + window.clearTimeout(reconnectTimeoutRef.current); + reconnectTimeoutRef.current = null; + } + }, []); + + const closeSocket = useCallback(() => { + intentionalCloseRef.current = true; + clearReconnectTimeout(); + const socket = socketRef.current; + if (socket) { + socket.close(); + socketRef.current = null; + } + setStatus("closed"); + }, [clearReconnectTimeout]); + + useEffect(() => { + if (!workspaceId) { + setStatus("idle"); + return; + } + + reconnectAttemptsRef.current = 0; + lastFrameWasExitRef.current = false; + intentionalCloseRef.current = false; + + const openSocket = () => { + const wsUrl = + (window.location.protocol === "https:" ? "wss:" : "ws:") + + "//" + + window.location.host + + "/api/workspaces/" + + encodeURIComponent(workspaceId) + + "/process/ws"; + + const ws = new WebSocket(wsUrl); + socketRef.current = ws; + setStatus("connecting"); + + ws.onopen = () => { + reconnectAttemptsRef.current = 0; + lastFrameWasExitRef.current = false; + setStatus("open"); + }; + + ws.onmessage = (event) => { + const data = typeof event.data === "string" ? event.data : ""; + if ( + data.startsWith("[process exited") || + data.startsWith("\r\n[process exited") + ) { + lastFrameWasExitRef.current = true; + } + callbacksRef.current.forEach((cb) => cb(data)); + }; + + ws.onclose = () => { + socketRef.current = null; + if (intentionalCloseRef.current || lastFrameWasExitRef.current) { + setStatus("closed"); + return; + } + + const attempts = reconnectAttemptsRef.current; + if (attempts < 3) { + setStatus("connecting"); + reconnectTimeoutRef.current = window.setTimeout(() => { + reconnectTimeoutRef.current = null; + openSocket(); + }, Math.pow(2, attempts) * 1000); + reconnectAttemptsRef.current = attempts + 1; + } else { + setStatus("error"); + } + }; + + ws.onerror = () => { + // The browser fires onclose after an error; handle retries there. + }; + }; + + openSocket(); + + return () => { + closeSocket(); + }; + }, [workspaceId, closeSocket]); + + const send = useCallback((data: string) => { + const socket = socketRef.current; + if ( + statusRef.current === "open" && + socket && + socket.readyState === WebSocket.OPEN + ) { + socket.send(data); + } + }, []); + + const onData = useCallback((cb: (data: string) => void) => { + callbacksRef.current.add(cb); + return () => { + callbacksRef.current.delete(cb); + }; + }, []); + + return { status, send, close: closeSocket, onData }; +}