Spaces:
Running
Running
File size: 5,398 Bytes
af1ae43 |
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 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 |
import { WebSocketEvent } from '@/types/agent';
import { useCallback, useEffect, useRef, useState } from 'react';
interface UseWebSocketProps {
url: string;
onMessage: (event: WebSocketEvent) => void;
onError?: (error: Event) => void;
}
export const useWebSocket = ({ url, onMessage, onError }: UseWebSocketProps) => {
const [isConnected, setIsConnected] = useState(false);
const [connectionState, setConnectionState] = useState<'connecting' | 'connected' | 'disconnected' | 'error'>('disconnected');
const wsRef = useRef<WebSocket | null>(null);
const reconnectTimeoutRef = useRef<NodeJS.Timeout>();
const reconnectAttemptsRef = useRef(0);
const maxReconnectAttempts = 3; // Only try three times, then stop
const baseReconnectDelay = 3000; // Start with 3 seconds
const maxReconnectDelay = 5000; // Max 5 seconds
const lastErrorTimeRef = useRef(0);
const errorThrottleMs = 5000; // Only show error toast once every 5 seconds
const isInitialConnectionRef = useRef(true); // Track if this is the first connection attempt
const getReconnectDelay = () => {
// Exponential backoff with jitter
const delay = Math.min(
baseReconnectDelay * Math.pow(2, reconnectAttemptsRef.current),
maxReconnectDelay
);
return delay + Math.random() * 1000; // Add jitter
};
const connect = useCallback(() => {
if (wsRef.current?.readyState === WebSocket.OPEN || wsRef.current?.readyState === WebSocket.CONNECTING) {
return; // Already connected or connecting
}
try {
setConnectionState('connecting');
const ws = new WebSocket(url);
ws.onopen = () => {
console.log('WebSocket connected');
setIsConnected(true);
setConnectionState('connected');
reconnectAttemptsRef.current = 0; // Reset attempts on successful connection
isInitialConnectionRef.current = false; // Mark that we've had a successful connection
};
ws.onmessage = (event) => {
try {
const data = JSON.parse(event.data) as WebSocketEvent;
onMessage(data);
} catch (error) {
console.error('Failed to parse WebSocket message:', error);
}
};
ws.onerror = (error) => {
console.error('WebSocket error:', error);
setConnectionState('error');
// Don't show error toasts on initial connection failure
// Only show toasts after we've had a successful connection before
if (!isInitialConnectionRef.current) {
// Throttle error notifications
const now = Date.now();
if (now - lastErrorTimeRef.current > errorThrottleMs) {
lastErrorTimeRef.current = now;
onError?.(error);
}
}
};
ws.onclose = (event) => {
console.log('WebSocket disconnected', { code: event.code, reason: event.reason });
setIsConnected(false);
setConnectionState('disconnected');
// Only attempt to reconnect if it wasn't a manual close (code 1000) and we haven't exceeded max attempts
if (event.code !== 1000 && reconnectAttemptsRef.current < maxReconnectAttempts) {
const delay = getReconnectDelay();
console.log(`Attempting to reconnect in ${Math.round(delay)}ms (attempt ${reconnectAttemptsRef.current + 1}/${maxReconnectAttempts})`);
reconnectTimeoutRef.current = setTimeout(() => {
reconnectAttemptsRef.current++;
connect();
}, delay);
} else if (reconnectAttemptsRef.current >= maxReconnectAttempts) {
console.log('Max reconnection attempts reached');
setConnectionState('error');
} else if (event.code === 1000) {
// Normal closure - don't reconnect
setConnectionState('disconnected');
console.log('WebSocket closed normally, not reconnecting');
}
};
wsRef.current = ws;
} catch (error) {
console.error('Failed to create WebSocket connection:', error);
setConnectionState('error');
}
}, [url, onMessage, onError]);
const disconnect = useCallback(() => {
if (reconnectTimeoutRef.current) {
clearTimeout(reconnectTimeoutRef.current);
}
if (wsRef.current) {
wsRef.current.close(1000, 'Manual disconnect');
wsRef.current = null;
}
setIsConnected(false);
setConnectionState('disconnected');
reconnectAttemptsRef.current = 0;
}, []);
const manualReconnect = useCallback(() => {
console.log('Manual reconnect requested');
disconnect();
reconnectAttemptsRef.current = 0;
isInitialConnectionRef.current = false; // Allow error toasts on manual reconnect
setTimeout(() => connect(), 1000); // Small delay before reconnecting
}, [disconnect, connect]);
const sendMessage = (message: unknown) => {
if (wsRef.current?.readyState === WebSocket.OPEN) {
try {
wsRef.current.send(JSON.stringify(message));
} catch (error) {
console.error('Failed to send WebSocket message:', error);
}
} else {
console.warn('WebSocket is not connected');
}
};
useEffect(() => {
connect();
return () => {
disconnect();
};
}, [url]); // Only depend on url, not the functions
return {
isConnected,
connectionState,
sendMessage,
reconnect: connect,
disconnect,
manualReconnect
};
};
|