feat(web): add useProcessWebSocket hook
Co-Authored-By: Claude <noreply@anthropic.com>
This commit is contained in:
@@ -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<void, Error, string> {
|
||||
export function useProcessRestart(): UseMutationResult<void, Error, string> {
|
||||
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<ProcessWebSocketStatus>(() =>
|
||||
workspaceId ? "connecting" : "idle",
|
||||
);
|
||||
const statusRef = useRef(status);
|
||||
useEffect(() => {
|
||||
statusRef.current = status;
|
||||
}, [status]);
|
||||
|
||||
const socketRef = useRef<WebSocket | null>(null);
|
||||
const reconnectAttemptsRef = useRef(0);
|
||||
const reconnectTimeoutRef = useRef<number | null>(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 };
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user