| 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 { |
| |
| } |
| 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 { |
| |
| } |
| }; |
| 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 }; |
| } |
|
|
