Files
platform/src/workspaces/hooks/src/useChatWebSocket.ts
T
2026-02-16 19:34:35 +00:00

79 lines
2.2 KiB
TypeScript

import { useState, useEffect, useRef } from 'react';
type UseChatWebSocketParams = {
url: string;
onMessage: (data: unknown) => void;
};
export const useChatWebSocket = ({ url, onMessage }: UseChatWebSocketParams) => {
const [isConnected, setIsConnected] = useState(false);
const socketRef = useRef<WebSocket | null>(null);
const retryRef = useRef(0);
const retryTimeoutRef = useRef<number | null>(null);
const isCleaningUpRef = useRef(false);
const onMessageRef = useRef(onMessage);
onMessageRef.current = onMessage;
const connect = () => {
if (isCleaningUpRef.current) return;
if (socketRef.current && socketRef.current.readyState === WebSocket.OPEN) return;
const socket = new WebSocket(url);
socketRef.current = socket;
socket.addEventListener('open', () => {
if (socketRef.current !== socket) return;
setIsConnected(true);
retryRef.current = 0;
});
socket.addEventListener('message', (ev) => {
try {
const data = JSON.parse(typeof ev.data === 'string' ? ev.data : new TextDecoder().decode(ev.data));
onMessageRef.current(data);
} catch {
// ignore unparseable messages
}
});
socket.addEventListener('close', () => {
if (isCleaningUpRef.current) return;
if (socketRef.current !== socket) return;
setIsConnected(false);
const retry = (retryRef.current ?? 0) + 1;
retryRef.current = retry;
const delay = Math.min(5000, 300 * retry);
retryTimeoutRef.current = window.setTimeout(connect, delay);
});
socket.addEventListener('error', () => {
socket.close();
});
};
useEffect(() => {
isCleaningUpRef.current = false;
connect();
return () => {
isCleaningUpRef.current = true;
if (retryTimeoutRef.current !== null) {
clearTimeout(retryTimeoutRef.current);
retryTimeoutRef.current = null;
}
if (socketRef.current) {
socketRef.current.close();
socketRef.current = null;
}
};
}, [url]);
const send = (data: Record<string, unknown>) => {
const socket = socketRef.current;
if (!socket || socket.readyState !== WebSocket.OPEN) return;
socket.send(JSON.stringify(data));
};
return { isConnected, send };
};