/
/
/
1import { useState, useEffect, useRef, useCallback } from 'react';
2import type { WebSocketMessage } from '../api/types';
3import { ConnectionStatus, WebSocketError } from '../api/types';
4
5interface UseWebSocketOptions {
6 url?: string;
7 protocols?: string | string[];
8 reconnectAttempts?: number;
9 reconnectInterval?: number;
10 heartbeatInterval?: number;
11 onMessage?: (message: WebSocketMessage) => void;
12 onError?: (error: WebSocketError) => void;
13 onConnectionChange?: (status: ConnectionStatus) => void;
14}
15
16interface UseWebSocketReturn {
17 connectionStatus: ConnectionStatus;
18 sendMessage: (message: WebSocketMessage) => void;
19 disconnect: () => void;
20 reconnect: () => void;
21 lastMessage: WebSocketMessage | null;
22 error: WebSocketError | null;
23}
24
25export function useWebSocket(options: UseWebSocketOptions = {}): UseWebSocketReturn {
26 const {
27 url = `${window.location.protocol === 'https:' ? 'wss:' : 'ws:'}//${window.location.host}/ws`,
28 protocols,
29 reconnectAttempts = 5,
30 reconnectInterval = 3000,
31 heartbeatInterval = 30000,
32 onMessage,
33 onError,
34 onConnectionChange,
35 } = options;
36
37 const [connectionStatus, setConnectionStatus] = useState<ConnectionStatus>(ConnectionStatus.DISCONNECTED);
38 const [lastMessage, setLastMessage] = useState<WebSocketMessage | null>(null);
39 const [error, setError] = useState<WebSocketError | null>(null);
40
41 const wsRef = useRef<WebSocket | null>(null);
42 const reconnectTimeoutRef = useRef<NodeJS.Timeout | null>(null);
43 const heartbeatTimeoutRef = useRef<NodeJS.Timeout | null>(null);
44 const reconnectAttemptsRef = useRef(0);
45 const isManuallyClosedRef = useRef(false);
46
47 const updateConnectionStatus = useCallback((status: ConnectionStatus) => {
48 setConnectionStatus(status);
49 onConnectionChange?.(status);
50 }, [onConnectionChange]);
51
52 const handleError = useCallback((wsError: WebSocketError) => {
53 setError(wsError);
54 onError?.(wsError);
55 }, [onError]);
56
57 const startHeartbeat = useCallback(() => {
58 if (heartbeatTimeoutRef.current) {
59 clearTimeout(heartbeatTimeoutRef.current);
60 }
61
62 heartbeatTimeoutRef.current = setTimeout(() => {
63 if (wsRef.current?.readyState === WebSocket.OPEN) {
64 const pingMessage: WebSocketMessage = {
65 type: 'ping',
66 data: {},
67 timestamp: new Date().toISOString(),
68 };
69 wsRef.current.send(JSON.stringify(pingMessage));
70 startHeartbeat(); // Schedule next heartbeat
71 }
72 }, heartbeatInterval);
73 }, [heartbeatInterval]);
74
75 const stopHeartbeat = useCallback(() => {
76 if (heartbeatTimeoutRef.current) {
77 clearTimeout(heartbeatTimeoutRef.current);
78 heartbeatTimeoutRef.current = null;
79 }
80 }, []);
81
82 const connect = useCallback(() => {
83 if (wsRef.current?.readyState === WebSocket.OPEN) {
84 return;
85 }
86
87 try {
88 updateConnectionStatus(ConnectionStatus.CONNECTING);
89 wsRef.current = new WebSocket(url, protocols);
90
91 wsRef.current.onopen = () => {
92 updateConnectionStatus(ConnectionStatus.CONNECTED);
93 reconnectAttemptsRef.current = 0;
94 setError(null);
95 startHeartbeat();
96 };
97
98 wsRef.current.onmessage = (event) => {
99 try {
100 const message: WebSocketMessage = JSON.parse(event.data);
101 setLastMessage(message);
102 onMessage?.(message);
103
104 // Handle pong responses
105 if (message.type === 'pong') {
106 // Reset heartbeat timer on pong
107 startHeartbeat();
108 }
109 } catch (err) {
110 handleError(new WebSocketError('Failed to parse message', 0, 'Invalid JSON'));
111 }
112 };
113
114 wsRef.current.onclose = (event) => {
115 stopHeartbeat();
116
117 if (isManuallyClosedRef.current) {
118 updateConnectionStatus(ConnectionStatus.DISCONNECTED);
119 return;
120 }
121
122 updateConnectionStatus(ConnectionStatus.ERROR);
123
124 if (reconnectAttemptsRef.current < reconnectAttempts) {
125 reconnectAttemptsRef.current++;
126 reconnectTimeoutRef.current = setTimeout(() => {
127 connect();
128 }, reconnectInterval * reconnectAttemptsRef.current);
129 } else {
130 handleError(new WebSocketError('Max reconnection attempts reached', event.code, event.reason));
131 }
132 };
133
134 wsRef.current.onerror = (_event) => {
135 handleError(new WebSocketError('WebSocket connection error', 0, 'Connection failed'));
136 };
137
138 } catch (err) {
139 handleError(new WebSocketError('Failed to create WebSocket connection', 0, 'Connection creation failed'));
140 }
141 }, [url, protocols, reconnectAttempts, reconnectInterval, updateConnectionStatus, handleError, onMessage, startHeartbeat, stopHeartbeat]);
142
143 const disconnect = useCallback(() => {
144 isManuallyClosedRef.current = true;
145
146 if (reconnectTimeoutRef.current) {
147 clearTimeout(reconnectTimeoutRef.current);
148 reconnectTimeoutRef.current = null;
149 }
150
151 stopHeartbeat();
152
153 if (wsRef.current) {
154 wsRef.current.close();
155 wsRef.current = null;
156 }
157
158 updateConnectionStatus(ConnectionStatus.DISCONNECTED);
159 }, [stopHeartbeat, updateConnectionStatus]);
160
161 const reconnect = useCallback(() => {
162 disconnect();
163 isManuallyClosedRef.current = false;
164 reconnectAttemptsRef.current = 0;
165 setTimeout(connect, 100);
166 }, [disconnect, connect]);
167
168 const sendMessage = useCallback((message: WebSocketMessage) => {
169 if (wsRef.current?.readyState === WebSocket.OPEN) {
170 try {
171 const messageWithTimestamp: WebSocketMessage = {
172 ...message,
173 timestamp: message.timestamp || new Date().toISOString(),
174 };
175 wsRef.current.send(JSON.stringify(messageWithTimestamp));
176 } catch (err) {
177 handleError(new WebSocketError('Failed to send message', 0, 'Send failed'));
178 }
179 } else {
180 handleError(new WebSocketError('WebSocket is not connected', 0, 'Not connected'));
181 }
182 }, [handleError]);
183
184 // Connect on mount
185 useEffect(() => {
186 isManuallyClosedRef.current = false;
187 connect();
188
189 return () => {
190 isManuallyClosedRef.current = true;
191 disconnect();
192 };
193 }, [connect, disconnect]);
194
195 // Cleanup timeouts on unmount
196 useEffect(() => {
197 return () => {
198 if (reconnectTimeoutRef.current) {
199 clearTimeout(reconnectTimeoutRef.current);
200 }
201 stopHeartbeat();
202 };
203 }, [stopHeartbeat]);
204
205 return {
206 connectionStatus,
207 sendMessage,
208 disconnect,
209 reconnect,
210 lastMessage,
211 error,
212 };
213}