File size: 2,090 Bytes
e6a1f55 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 | import { useCallback, useEffect, useRef, useState } from "react";
import { inferenceWebSocketUrl } from "@/lib/api";
import type { WsEvent } from "@/types/inference";
interface UseInferenceWebSocket {
status: "idle" | "connecting" | "open" | "closed" | "error";
events: WsEvent[];
open: (sessionId?: string) => void;
close: () => void;
send: (data: object) => void;
sessionId: string | null;
reset: () => void;
}
export function useInferenceWebSocket(): UseInferenceWebSocket {
const [status, setStatus] = useState<UseInferenceWebSocket["status"]>("idle");
const [events, setEvents] = useState<WsEvent[]>([]);
const [sessionId, setSessionId] = useState<string | null>(null);
const wsRef = useRef<WebSocket | null>(null);
const close = useCallback(() => {
if (wsRef.current) {
try {
wsRef.current.close();
} catch {
// noop
}
wsRef.current = null;
}
setStatus("closed");
}, []);
const reset = useCallback(() => {
close();
setEvents([]);
setSessionId(null);
setStatus("idle");
}, [close]);
const open = useCallback((override?: string) => {
const id =
override ??
(typeof crypto !== "undefined" && "randomUUID" in crypto
? crypto.randomUUID()
: Math.random().toString(36).slice(2, 14));
setSessionId(id);
setEvents([]);
setStatus("connecting");
const ws = new WebSocket(inferenceWebSocketUrl(id));
wsRef.current = ws;
ws.onopen = () => setStatus("open");
ws.onmessage = (msg) => {
try {
const evt = JSON.parse(msg.data) as WsEvent;
setEvents((prev) => [...prev, evt]);
} catch {
// ignore malformed
}
};
ws.onerror = () => setStatus("error");
ws.onclose = () => setStatus("closed");
}, []);
const send = useCallback((data: object) => {
if (wsRef.current?.readyState === WebSocket.OPEN) {
wsRef.current.send(JSON.stringify(data));
}
}, []);
useEffect(() => () => close(), [close]);
return { status, events, open, close, send, sessionId, reset };
}
|