/** * Custom hook for WebSocket connection. * * Manages WebSocket connection for real-time updates. */ import { useState, useEffect, useCallback, useRef } from 'react'; // Determine WebSocket URL based on environment const getWebSocketURL = () => { // If VITE_WS_URL is explicitly set, use it if (import.meta.env.VITE_WS_URL) { return import.meta.env.VITE_WS_URL; } // Otherwise, build WebSocket URL from current location const protocol = window.location.protocol === 'https:' ? 'wss:' : 'ws:'; const host = window.location.host; return `${protocol}//${host}`; }; const WS_URL = getWebSocketURL(); interface UseWebSocketOptions { onMessage?: (data: any) => void; onConnect?: () => void; onDisconnect?: () => void; reconnect?: boolean; reconnectInterval?: number; } export const useWebSocket = (endpoint: string, options: UseWebSocketOptions = {}) => { const { onMessage, onConnect, onDisconnect, reconnect = true, reconnectInterval = 3000, } = options; const [isConnected, setIsConnected] = useState(false); const [lastMessage, setLastMessage] = useState(null); const wsRef = useRef(null); const reconnectTimeoutRef = useRef(null); // Store callbacks in refs to avoid recreating connect function const onMessageRef = useRef(onMessage); const onConnectRef = useRef(onConnect); const onDisconnectRef = useRef(onDisconnect); // Update refs when callbacks change useEffect(() => { onMessageRef.current = onMessage; onConnectRef.current = onConnect; onDisconnectRef.current = onDisconnect; }, [onMessage, onConnect, onDisconnect]); const connect = useCallback(() => { // Close existing connection if any if (wsRef.current) { wsRef.current.close(); wsRef.current = null; } try { // Get authentication token from localStorage const token = localStorage.getItem('auth_token'); // Build WebSocket URL with token as query parameter const wsUrl = token ? `${WS_URL}${endpoint}?token=${token}` : `${WS_URL}${endpoint}`; const ws = new WebSocket(wsUrl); ws.onopen = () => { console.log('WebSocket connected'); setIsConnected(true); wsRef.current = ws; onConnectRef.current?.(); }; ws.onmessage = (event) => { try { const data = JSON.parse(event.data); setLastMessage(data); onMessageRef.current?.(data); } catch (err) { console.error('Failed to parse WebSocket message:', err); } }; ws.onclose = () => { console.log('WebSocket disconnected'); setIsConnected(false); wsRef.current = null; onDisconnectRef.current?.(); // Attempt to reconnect if (reconnect) { reconnectTimeoutRef.current = setTimeout(() => { console.log('Attempting to reconnect WebSocket...'); connect(); }, reconnectInterval); } }; ws.onerror = (error) => { console.error('WebSocket error:', error); }; } catch (err) { console.error('Failed to create WebSocket connection:', err); } }, [endpoint, reconnect, reconnectInterval]); const disconnect = useCallback(() => { if (reconnectTimeoutRef.current) { clearTimeout(reconnectTimeoutRef.current); } if (wsRef.current) { wsRef.current.close(); wsRef.current = null; } setIsConnected(false); }, []); const sendMessage = useCallback((data: any) => { if (wsRef.current && isConnected) { const message = typeof data === 'string' ? data : JSON.stringify(data); wsRef.current.send(message); } else { console.warn('Cannot send message: WebSocket not connected'); } }, [isConnected]); useEffect(() => { connect(); return () => { disconnect(); }; // eslint-disable-next-line react-hooks/exhaustive-deps }, [endpoint]); // Only reconnect if endpoint changes return { isConnected, lastMessage, sendMessage, disconnect, reconnect: connect, }; };