HyperCluster/Fara-BrowserUse

1import { WebSocketEvent } from '@/types/agent';
2import { useCallback, useEffect, useRef, useState } from 'react';
3
4interface UseWebSocketProps {
5  url: string;
6  onMessage: (event: WebSocketEvent) => void;
7  onError?: (error: Event) => void;
8}
9
10export const useWebSocket = ({ url, onMessage, onError }: UseWebSocketProps) => {
11  const [isConnected, setIsConnected] = useState(false);
12  const [connectionState, setConnectionState] = useState<'connecting' | 'connected' | 'disconnected' | 'error'>('disconnected');
13  const wsRef = useRef<WebSocket | null>(null);
14  const reconnectTimeoutRef = useRef<NodeJS.Timeout>();
15  const reconnectAttemptsRef = useRef(0);
16  const maxReconnectAttempts = 3; // Only try three times, then stop
17  const baseReconnectDelay = 3000; // Start with 3 seconds
18  const maxReconnectDelay = 5000; // Max 5 seconds
19  const lastErrorTimeRef = useRef(0);
20  const errorThrottleMs = 5000; // Only show error toast once every 5 seconds
21  const isInitialConnectionRef = useRef(true); // Track if this is the first connection attempt
22
23  const getReconnectDelay = () => {
24    // Exponential backoff with jitter
25    const delay = Math.min(
26      baseReconnectDelay * Math.pow(2, reconnectAttemptsRef.current),
27      maxReconnectDelay
28    );
29    return delay + Math.random() * 1000; // Add jitter
30  };
31
32  const connect = useCallback(() => {
33    if (wsRef.current?.readyState === WebSocket.OPEN || wsRef.current?.readyState === WebSocket.CONNECTING) {
34      return; // Already connected or connecting
35    }
36
37    try {
38      setConnectionState('connecting');
39      const ws = new WebSocket(url);
40
41      ws.onopen = () => {
42        console.log('WebSocket connected');
43        setIsConnected(true);
44        setConnectionState('connected');
45        reconnectAttemptsRef.current = 0; // Reset attempts on successful connection
46        isInitialConnectionRef.current = false; // Mark that we've had a successful connection
47      };
48
49      ws.onmessage = (event) => {
50        try {
51          const data = JSON.parse(event.data) as WebSocketEvent;
52          onMessage(data);
53        } catch (error) {
54          console.error('Failed to parse WebSocket message:', error);
55        }
56      };
57
58      ws.onerror = (error) => {
59        console.error('WebSocket error:', error);
60        setConnectionState('error');
61
62        // Don't show error toasts on initial connection failure
63        // Only show toasts after we've had a successful connection before
64        if (!isInitialConnectionRef.current) {
65          // Throttle error notifications
66          const now = Date.now();
67          if (now - lastErrorTimeRef.current > errorThrottleMs) {
68            lastErrorTimeRef.current = now;
69            onError?.(error);
70          }
71        }
72      };
73
74      ws.onclose = (event) => {
75        console.log('WebSocket disconnected', { code: event.code, reason: event.reason });
76        setIsConnected(false);
77        setConnectionState('disconnected');
78
79        // Only attempt to reconnect if it wasn't a manual close (code 1000) and we haven't exceeded max attempts
80        if (event.code !== 1000 && reconnectAttemptsRef.current < maxReconnectAttempts) {
81          const delay = getReconnectDelay();
82          console.log(`Attempting to reconnect in ${Math.round(delay)}ms (attempt ${reconnectAttemptsRef.current + 1}/${maxReconnectAttempts})`);
83
84          reconnectTimeoutRef.current = setTimeout(() => {
85            reconnectAttemptsRef.current++;
86            connect();
87          }, delay);
88        } else if (reconnectAttemptsRef.current >= maxReconnectAttempts) {
89          console.log('Max reconnection attempts reached');
90          setConnectionState('error');
91        } else if (event.code === 1000) {
92          // Normal closure - don't reconnect
93          setConnectionState('disconnected');
94          console.log('WebSocket closed normally, not reconnecting');
95        }
96      };
97
98      wsRef.current = ws;
99    } catch (error) {
100      console.error('Failed to create WebSocket connection:', error);
101      setConnectionState('error');
102    }
103  }, [url, onMessage, onError]);
104
105  const disconnect = useCallback(() => {
106    if (reconnectTimeoutRef.current) {
107      clearTimeout(reconnectTimeoutRef.current);
108    }
109    if (wsRef.current) {
110      wsRef.current.close(1000, 'Manual disconnect');
111      wsRef.current = null;
112    }
113    setIsConnected(false);
114    setConnectionState('disconnected');
115    reconnectAttemptsRef.current = 0;
116  }, []);
117
118  const manualReconnect = useCallback(() => {
119    console.log('Manual reconnect requested');
120    disconnect();
121    reconnectAttemptsRef.current = 0;
122    isInitialConnectionRef.current = false; // Allow error toasts on manual reconnect
123    setTimeout(() => connect(), 1000); // Small delay before reconnecting
124  }, [disconnect, connect]);
125
126  const sendMessage = (message: unknown) => {
127    if (wsRef.current?.readyState === WebSocket.OPEN) {
128      try {
129        wsRef.current.send(JSON.stringify(message));
130      } catch (error) {
131        console.error('Failed to send WebSocket message:', error);
132      }
133    } else {
134      console.warn('WebSocket is not connected');
135    }
136  };
137
138  useEffect(() => {
139    connect();
140
141    return () => {
142      disconnect();
143    };
144  }, [url]); // Only depend on url, not the functions
145
146  return {
147    isConnected,
148    connectionState,
149    sendMessage,
150    reconnect: connect,
151    disconnect,
152    manualReconnect
153  };
154};
155