@@ -9,9 +9,9 @@ import crypto from 'crypto';
9
9
import { toast } from 'sonner' ;
10
10
import { useSearchParams } from 'next/navigation' ;
11
11
import { getSuggestions } from '@/lib/actions' ;
12
- import Error from 'next/error' ;
13
12
import { Settings } from 'lucide-react' ;
14
13
import SettingsDialog from './SettingsDialog' ;
14
+ import NextError from 'next/error' ;
15
15
16
16
export type Message = {
17
17
messageId : string ;
@@ -34,11 +34,24 @@ const useSocket = (
34
34
setIsWSReady : ( ready : boolean ) => void ,
35
35
setError : ( error : boolean ) => void ,
36
36
) => {
37
- const [ ws , setWs ] = useState < WebSocket | null > ( null ) ;
37
+ const wsRef = useRef < WebSocket | null > ( null ) ;
38
+ const reconnectTimeoutRef = useRef < NodeJS . Timeout > ( ) ;
39
+ const retryCountRef = useRef ( 0 ) ;
40
+ const isCleaningUpRef = useRef ( false ) ;
41
+ const MAX_RETRIES = 3 ;
42
+ const INITIAL_BACKOFF = 1000 ; // 1 second
43
+
44
+ const getBackoffDelay = ( retryCount : number ) => {
45
+ return Math . min ( INITIAL_BACKOFF * Math . pow ( 2 , retryCount ) , 10000 ) ; // Cap at 10 seconds
46
+ } ;
38
47
39
48
useEffect ( ( ) => {
40
- if ( ! ws ) {
41
- const connectWs = async ( ) => {
49
+ const connectWs = async ( ) => {
50
+ if ( wsRef . current ?. readyState === WebSocket . OPEN ) {
51
+ wsRef . current . close ( ) ;
52
+ }
53
+
54
+ try {
42
55
let chatModel = localStorage . getItem ( 'chatModel' ) ;
43
56
let chatModelProvider = localStorage . getItem ( 'chatModelProvider' ) ;
44
57
let embeddingModel = localStorage . getItem ( 'embeddingModel' ) ;
@@ -61,7 +74,13 @@ const useSocket = (
61
74
'Content-Type' : 'application/json' ,
62
75
} ,
63
76
} ,
64
- ) . then ( async ( res ) => await res . json ( ) ) ;
77
+ ) . then ( async ( res ) => {
78
+ if ( ! res . ok )
79
+ throw new Error (
80
+ `Failed to fetch models: ${ res . status } ${ res . statusText } ` ,
81
+ ) ;
82
+ return res . json ( ) ;
83
+ } ) ;
65
84
66
85
if (
67
86
! chatModel ||
@@ -204,6 +223,7 @@ const useSocket = (
204
223
wsURL . search = searchParams . toString ( ) ;
205
224
206
225
const ws = new WebSocket ( wsURL . toString ( ) ) ;
226
+ wsRef . current = ws ;
207
227
208
228
const timeoutId = setTimeout ( ( ) => {
209
229
if ( ws . readyState !== 1 ) {
@@ -219,11 +239,16 @@ const useSocket = (
219
239
const interval = setInterval ( ( ) => {
220
240
if ( ws . readyState === 1 ) {
221
241
setIsWSReady ( true ) ;
242
+ setError ( false ) ;
243
+ if ( retryCountRef . current > 0 ) {
244
+ toast . success ( 'Connection restored.' ) ;
245
+ }
246
+ retryCountRef . current = 0 ;
222
247
clearInterval ( interval ) ;
223
248
}
224
249
} , 5 ) ;
225
250
clearTimeout ( timeoutId ) ;
226
- console . log ( '[DEBUG] opened ') ;
251
+ console . debug ( new Date ( ) , 'ws:connected ') ;
227
252
}
228
253
if ( data . type === 'error' ) {
229
254
toast . error ( data . data ) ;
@@ -232,24 +257,68 @@ const useSocket = (
232
257
233
258
ws . onerror = ( ) => {
234
259
clearTimeout ( timeoutId ) ;
235
- setError ( true ) ;
260
+ setIsWSReady ( false ) ;
236
261
toast . error ( 'WebSocket connection error.' ) ;
237
262
} ;
238
263
239
264
ws . onclose = ( ) => {
240
265
clearTimeout ( timeoutId ) ;
241
- setError ( true ) ;
242
- console . log ( '[DEBUG] closed' ) ;
266
+ setIsWSReady ( false ) ;
267
+ console . debug ( new Date ( ) , 'ws:disconnected' ) ;
268
+ if ( ! isCleaningUpRef . current ) {
269
+ toast . error ( 'Connection lost. Attempting to reconnect...' ) ;
270
+ attemptReconnect ( ) ;
271
+ }
243
272
} ;
273
+ } catch ( error ) {
274
+ console . debug ( new Date ( ) , 'ws:error' , error ) ;
275
+ setIsWSReady ( false ) ;
276
+ attemptReconnect ( ) ;
277
+ }
278
+ } ;
244
279
245
- setWs ( ws ) ;
246
- } ;
280
+ const attemptReconnect = ( ) => {
281
+ retryCountRef . current += 1 ;
247
282
248
- connectWs ( ) ;
249
- }
250
- } , [ ws , url , setIsWSReady , setError ] ) ;
283
+ if ( retryCountRef . current > MAX_RETRIES ) {
284
+ console . debug ( new Date ( ) , 'ws:max_retries' ) ;
285
+ setError ( true ) ;
286
+ toast . error (
287
+ 'Unable to connect to server after multiple attempts. Please refresh the page to try again.' ,
288
+ ) ;
289
+ return ;
290
+ }
251
291
252
- return ws ;
292
+ const backoffDelay = getBackoffDelay ( retryCountRef . current ) ;
293
+ console . debug (
294
+ new Date ( ) ,
295
+ `ws:retry attempt=${ retryCountRef . current } /${ MAX_RETRIES } delay=${ backoffDelay } ms` ,
296
+ ) ;
297
+
298
+ if ( reconnectTimeoutRef . current ) {
299
+ clearTimeout ( reconnectTimeoutRef . current ) ;
300
+ }
301
+
302
+ reconnectTimeoutRef . current = setTimeout ( ( ) => {
303
+ connectWs ( ) ;
304
+ } , backoffDelay ) ;
305
+ } ;
306
+
307
+ connectWs ( ) ;
308
+
309
+ return ( ) => {
310
+ if ( reconnectTimeoutRef . current ) {
311
+ clearTimeout ( reconnectTimeoutRef . current ) ;
312
+ }
313
+ if ( wsRef . current ?. readyState === WebSocket . OPEN ) {
314
+ wsRef . current . close ( ) ;
315
+ isCleaningUpRef . current = true ;
316
+ console . debug ( new Date ( ) , 'ws:cleanup' ) ;
317
+ }
318
+ } ;
319
+ } , [ url , setIsWSReady , setError ] ) ;
320
+
321
+ return wsRef . current ;
253
322
} ;
254
323
255
324
const loadMessages = async (
@@ -293,7 +362,7 @@ const loadMessages = async (
293
362
return [ msg . role , msg . content ] ;
294
363
} ) as [ string , string ] [ ] ;
295
364
296
- console . log ( '[DEBUG] messages loaded ') ;
365
+ console . debug ( new Date ( ) , 'app:messages_loaded ') ;
297
366
298
367
document . title = messages [ 0 ] . content ;
299
368
@@ -377,7 +446,7 @@ const ChatWindow = ({ id }: { id?: string }) => {
377
446
return ( ) => {
378
447
if ( ws ?. readyState === 1 ) {
379
448
ws . close ( ) ;
380
- console . log ( '[DEBUG] closed ') ;
449
+ console . debug ( new Date ( ) , 'ws:cleanup ') ;
381
450
}
382
451
} ;
383
452
// eslint-disable-next-line react-hooks/exhaustive-deps
@@ -392,12 +461,18 @@ const ChatWindow = ({ id }: { id?: string }) => {
392
461
useEffect ( ( ) => {
393
462
if ( isMessagesLoaded && isWSReady ) {
394
463
setIsReady ( true ) ;
395
- console . log ( '[DEBUG] ready' ) ;
464
+ console . debug ( new Date ( ) , 'app:ready' ) ;
465
+ } else {
466
+ setIsReady ( false ) ;
396
467
}
397
468
} , [ isMessagesLoaded , isWSReady ] ) ;
398
469
399
470
const sendMessage = async ( message : string , messageId ?: string ) => {
400
471
if ( loading ) return ;
472
+ if ( ! ws || ws . readyState !== WebSocket . OPEN ) {
473
+ toast . error ( 'Cannot send message while disconnected' ) ;
474
+ return ;
475
+ }
401
476
402
477
setLoading ( true ) ;
403
478
setMessageAppeared ( false ) ;
@@ -408,7 +483,7 @@ const ChatWindow = ({ id }: { id?: string }) => {
408
483
409
484
messageId = messageId ?? crypto . randomBytes ( 7 ) . toString ( 'hex' ) ;
410
485
411
- ws ? .send (
486
+ ws . send (
412
487
JSON . stringify ( {
413
488
type : 'message' ,
414
489
message : {
@@ -571,7 +646,7 @@ const ChatWindow = ({ id }: { id?: string }) => {
571
646
572
647
return isReady ? (
573
648
notFound ? (
574
- < Error statusCode = { 404 } />
649
+ < NextError statusCode = { 404 } />
575
650
) : (
576
651
< div >
577
652
{ messages . length > 0 ? (
0 commit comments