Spaces:
Sleeping
Sleeping
File size: 4,159 Bytes
453520f 27f7513 453520f 30fd5f3 453520f 2d64894 453520f 2d64894 453520f 2d64894 453520f 2d64894 453520f 2d64894 453520f 2d64894 453520f 2d64894 453520f | 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 | /**
* 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<any>(null);
const wsRef = useRef<WebSocket | null>(null);
const reconnectTimeoutRef = useRef<number | null>(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,
};
};
|