File size: 4,883 Bytes
973d6a7 | 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 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 | import { useCallback, useEffect, useRef, useState } from 'react';
/* ───────────────── Types ───────────────── */
export type WSEventType =
| 'thinking'
| 'status'
| 'tool_start'
| 'stream'
| 'chunk' // old backend sends 'chunk', we normalize to 'stream'
| 'plot'
| 'video'
| 'complete'
| 'error'
| 'clear'
| 'keys_configured'
| 'request_keys'
| 'arraylake_snippet';
export interface WSEvent {
type: WSEventType;
content?: string;
ready?: boolean;
reason?: string;
data?: string; // base64 payload for plot/video
path?: string; // file path for plot/video
code?: string; // code that generated the plot
mimetype?: string; // video mimetype
[key: string]: unknown;
}
export type WSStatus = 'connecting' | 'connected' | 'disconnected';
interface UseWebSocketReturn {
status: WSStatus;
send: (payload: Record<string, unknown>) => void;
sendMessage: (message: string) => void;
configureKeys: (keys: { openai_api_key?: string; arraylake_api_key?: string; hf_token?: string }) => void;
lastEvent: WSEvent | null;
}
/* ─────────────── Hook ─────────────── */
export function useWebSocket(onEvent?: (event: WSEvent) => void): UseWebSocketReturn {
const wsRef = useRef<WebSocket | null>(null);
const [status, setStatus] = useState<WSStatus>('disconnected');
const [lastEvent, setLastEvent] = useState<WSEvent | null>(null);
const reconnectTimer = useRef<ReturnType<typeof setTimeout> | undefined>(undefined);
const onEventRef = useRef(onEvent);
const connectingRef = useRef(false); // guard against double-connect
onEventRef.current = onEvent;
const connect = useCallback(() => {
// Guard: don't open a second WS if one is OPEN or CONNECTING
if (wsRef.current) {
const rs = wsRef.current.readyState;
if (rs === WebSocket.OPEN || rs === WebSocket.CONNECTING) return;
}
if (connectingRef.current) return;
connectingRef.current = true;
const protocol = window.location.protocol === 'https:' ? 'wss:' : 'ws:';
const ws = new WebSocket(`${protocol}//${window.location.host}/ws/chat`);
wsRef.current = ws;
setStatus('connecting');
ws.onopen = () => {
connectingRef.current = false;
setStatus('connected');
// Auto-resend keys from sessionStorage on reconnect
const saved = sessionStorage.getItem('eurus-keys');
if (saved) {
try {
const keys = JSON.parse(saved);
if (keys.openai_api_key) {
ws.send(JSON.stringify({ type: 'configure_keys', ...keys }));
}
} catch { /* ignore */ }
}
};
ws.onmessage = (e) => {
try {
const event: WSEvent = JSON.parse(e.data);
// Normalize 'chunk' → 'stream' for unified handling
if (event.type === 'chunk') {
event.type = 'stream';
}
setLastEvent(event);
onEventRef.current?.(event);
} catch { /* ignore non-json */ }
};
ws.onclose = () => {
connectingRef.current = false;
setStatus('disconnected');
wsRef.current = null;
// auto-reconnect after 2 s
reconnectTimer.current = setTimeout(connect, 2000);
};
ws.onerror = () => {
connectingRef.current = false;
ws.close();
};
}, []);
useEffect(() => {
connect();
return () => {
clearTimeout(reconnectTimer.current);
connectingRef.current = false;
if (wsRef.current) {
wsRef.current.onclose = null; // prevent reconnect on cleanup close
wsRef.current.close();
wsRef.current = null;
}
};
}, [connect]);
const send = useCallback((payload: Record<string, unknown>) => {
if (wsRef.current?.readyState === WebSocket.OPEN) {
wsRef.current.send(JSON.stringify(payload));
}
}, []);
const sendMessage = useCallback((message: string) => send({ message }), [send]);
const configureKeys = useCallback(
(keys: { openai_api_key?: string; arraylake_api_key?: string; hf_token?: string }) => {
// Save to sessionStorage
sessionStorage.setItem('eurus-keys', JSON.stringify(keys));
send({ type: 'configure_keys', ...keys });
},
[send],
);
return { status, send, sendMessage, configureKeys, lastEvent };
}
|