From 5526d5f60f31ec76e01e8f8482cf873ff6522202 Mon Sep 17 00:00:00 2001 From: realies <5107843+realies@users.noreply.github.com> Date: Sun, 5 Jan 2025 17:29:53 +0000 Subject: [PATCH 1/5] fix(ws-error): add exponential reconnect mechanism --- ui/components/ChatWindow.tsx | 102 ++++++++++++++++++++++++++++------- 1 file changed, 82 insertions(+), 20 deletions(-) diff --git a/ui/components/ChatWindow.tsx b/ui/components/ChatWindow.tsx index 62fa9dc..3968d92 100644 --- a/ui/components/ChatWindow.tsx +++ b/ui/components/ChatWindow.tsx @@ -9,7 +9,7 @@ import crypto from 'crypto'; import { toast } from 'sonner'; import { useSearchParams } from 'next/navigation'; import { getSuggestions } from '@/lib/actions'; -import Error from 'next/error'; +import NextError from 'next/error'; export type Message = { messageId: string; @@ -32,11 +32,24 @@ const useSocket = ( setIsWSReady: (ready: boolean) => void, setError: (error: boolean) => void, ) => { - const [ws, setWs] = useState(null); + const wsRef = useRef(null); + const reconnectTimeoutRef = useRef(); + const retryCountRef = useRef(0); + const isCleaningUpRef = useRef(false); + const MAX_RETRIES = 3; + const INITIAL_BACKOFF = 1000; // 1 second + + const getBackoffDelay = (retryCount: number) => { + return Math.min(INITIAL_BACKOFF * Math.pow(2, retryCount), 10000); // Cap at 10 seconds + }; useEffect(() => { - if (!ws) { - const connectWs = async () => { + const connectWs = async () => { + if (wsRef.current?.readyState === WebSocket.OPEN) { + wsRef.current.close(); + } + + try { let chatModel = localStorage.getItem('chatModel'); let chatModelProvider = localStorage.getItem('chatModelProvider'); let embeddingModel = localStorage.getItem('embeddingModel'); @@ -59,7 +72,10 @@ const useSocket = ( 'Content-Type': 'application/json', }, }, - ).then(async (res) => await res.json()); + ).then(async (res) => { + if (!res.ok) throw new Error(`Failed to fetch models: ${res.status} ${res.statusText}`); + return res.json(); + }); if ( !chatModel || @@ -202,6 +218,7 @@ const useSocket = ( wsURL.search = searchParams.toString(); const ws = new WebSocket(wsURL.toString()); + wsRef.current = ws; const timeoutId = setTimeout(() => { if (ws.readyState !== 1) { @@ -217,11 +234,14 @@ const useSocket = ( const interval = setInterval(() => { if (ws.readyState === 1) { setIsWSReady(true); + retryCountRef.current = 0; + setError(false); + toast.success('Connection restored'); clearInterval(interval); } }, 5); clearTimeout(timeoutId); - console.log('[DEBUG] opened'); + console.debug(new Date(), 'ws:connected'); } if (data.type === 'error') { toast.error(data.data); @@ -230,24 +250,62 @@ const useSocket = ( ws.onerror = () => { clearTimeout(timeoutId); - setError(true); + setIsWSReady(false); toast.error('WebSocket connection error.'); }; ws.onclose = () => { clearTimeout(timeoutId); - setError(true); - console.log('[DEBUG] closed'); + setIsWSReady(false); + console.debug(new Date(), 'ws:disconnected'); + if (!isCleaningUpRef.current) { + toast.error('Connection lost. Attempting to reconnect...'); + attemptReconnect(); + } }; - setWs(ws); - }; + } catch (error) { + console.debug(new Date(), 'ws:error', error); + attemptReconnect(); + } + }; - connectWs(); - } - }, [ws, url, setIsWSReady, setError]); + const attemptReconnect = () => { + retryCountRef.current += 1; + if (retryCountRef.current > MAX_RETRIES) { + console.debug(new Date(), 'ws:max_retries'); + setError(true); + toast.error('Unable to connect to server after multiple attempts. Please refresh the page to try again.'); + return; + } - return ws; + const backoffDelay = getBackoffDelay(retryCountRef.current); + console.debug(new Date(), `ws:retry attempt=${retryCountRef.current}/${MAX_RETRIES} delay=${backoffDelay}ms`); + + if (reconnectTimeoutRef.current) { + clearTimeout(reconnectTimeoutRef.current); + } + + reconnectTimeoutRef.current = setTimeout(() => { + connectWs(); + }, backoffDelay); + }; + + connectWs(); + + return () => { + isCleaningUpRef.current = true; + if (reconnectTimeoutRef.current) { + clearTimeout(reconnectTimeoutRef.current); + } + if (wsRef.current?.readyState === WebSocket.OPEN) { + wsRef.current.close(); + } + console.debug(new Date(), 'ws:cleanup'); + }; + }, [url, setIsWSReady, setError]); + + return wsRef.current; }; const loadMessages = async ( @@ -291,7 +349,7 @@ const loadMessages = async ( return [msg.role, msg.content]; }) as [string, string][]; - console.log('[DEBUG] messages loaded'); + console.debug(new Date(), 'app:messages_loaded'); document.title = messages[0].content; @@ -373,7 +431,7 @@ const ChatWindow = ({ id }: { id?: string }) => { return () => { if (ws?.readyState === 1) { ws.close(); - console.log('[DEBUG] closed'); + console.debug(new Date(), 'ws:cleanup'); } }; // eslint-disable-next-line react-hooks/exhaustive-deps @@ -388,12 +446,16 @@ const ChatWindow = ({ id }: { id?: string }) => { useEffect(() => { if (isMessagesLoaded && isWSReady) { setIsReady(true); - console.log('[DEBUG] ready'); + console.debug(new Date(), 'app:ready'); } }, [isMessagesLoaded, isWSReady]); const sendMessage = async (message: string, messageId?: string) => { if (loading) return; + if (!ws || ws.readyState !== WebSocket.OPEN) { + toast.error('Cannot send message while disconnected'); + return; + } setLoading(true); setMessageAppeared(false); @@ -404,7 +466,7 @@ const ChatWindow = ({ id }: { id?: string }) => { messageId = messageId ?? crypto.randomBytes(7).toString('hex'); - ws?.send( + ws.send( JSON.stringify({ type: 'message', message: { @@ -558,7 +620,7 @@ const ChatWindow = ({ id }: { id?: string }) => { return isReady ? ( notFound ? ( - + ) : (
{messages.length > 0 ? ( From 0ec54fe6c05270dcf03bfc3cc01b07c19b4bd41a Mon Sep 17 00:00:00 2001 From: ItzCrazyKns <95534749+ItzCrazyKns@users.noreply.github.com> Date: Tue, 7 Jan 2025 11:43:54 +0530 Subject: [PATCH 2/5] feat(chat-window): remove toast --- ui/components/ChatWindow.tsx | 1 - 1 file changed, 1 deletion(-) diff --git a/ui/components/ChatWindow.tsx b/ui/components/ChatWindow.tsx index 3968d92..b6f1b30 100644 --- a/ui/components/ChatWindow.tsx +++ b/ui/components/ChatWindow.tsx @@ -236,7 +236,6 @@ const useSocket = ( setIsWSReady(true); retryCountRef.current = 0; setError(false); - toast.success('Connection restored'); clearInterval(interval); } }, 5); From b7f7d25f549fc213b6dc63f17ea9204e41ba9806 Mon Sep 17 00:00:00 2001 From: ItzCrazyKns <95534749+ItzCrazyKns@users.noreply.github.com> Date: Tue, 7 Jan 2025 11:44:19 +0530 Subject: [PATCH 3/5] feat(chat-window): lint & beautify --- ui/components/ChatWindow.tsx | 15 +++++++++++---- 1 file changed, 11 insertions(+), 4 deletions(-) diff --git a/ui/components/ChatWindow.tsx b/ui/components/ChatWindow.tsx index b6f1b30..ed3594c 100644 --- a/ui/components/ChatWindow.tsx +++ b/ui/components/ChatWindow.tsx @@ -73,7 +73,10 @@ const useSocket = ( }, }, ).then(async (res) => { - if (!res.ok) throw new Error(`Failed to fetch models: ${res.status} ${res.statusText}`); + if (!res.ok) + throw new Error( + `Failed to fetch models: ${res.status} ${res.statusText}`, + ); return res.json(); }); @@ -262,7 +265,6 @@ const useSocket = ( attemptReconnect(); } }; - } catch (error) { console.debug(new Date(), 'ws:error', error); attemptReconnect(); @@ -274,12 +276,17 @@ const useSocket = ( if (retryCountRef.current > MAX_RETRIES) { console.debug(new Date(), 'ws:max_retries'); setError(true); - toast.error('Unable to connect to server after multiple attempts. Please refresh the page to try again.'); + toast.error( + 'Unable to connect to server after multiple attempts. Please refresh the page to try again.', + ); return; } const backoffDelay = getBackoffDelay(retryCountRef.current); - console.debug(new Date(), `ws:retry attempt=${retryCountRef.current}/${MAX_RETRIES} delay=${backoffDelay}ms`); + console.debug( + new Date(), + `ws:retry attempt=${retryCountRef.current}/${MAX_RETRIES} delay=${backoffDelay}ms`, + ); if (reconnectTimeoutRef.current) { clearTimeout(reconnectTimeoutRef.current); From 99cae076a79f993ba8da5f91d0166f06d67b1ad1 Mon Sep 17 00:00:00 2001 From: ItzCrazyKns <95534749+ItzCrazyKns@users.noreply.github.com> Date: Tue, 7 Jan 2025 11:49:40 +0530 Subject: [PATCH 4/5] feat(chat-window): display toast when retried --- ui/components/ChatWindow.tsx | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/ui/components/ChatWindow.tsx b/ui/components/ChatWindow.tsx index ed3594c..a44c550 100644 --- a/ui/components/ChatWindow.tsx +++ b/ui/components/ChatWindow.tsx @@ -237,8 +237,11 @@ const useSocket = ( const interval = setInterval(() => { if (ws.readyState === 1) { setIsWSReady(true); - retryCountRef.current = 0; setError(false); + if (retryCountRef.current > 0) { + toast.success('Connection restored.'); + } + retryCountRef.current = 0; clearInterval(interval); } }, 5); From 6d9d71279072112243ca45a887830ba2666000e8 Mon Sep 17 00:00:00 2001 From: ItzCrazyKns <95534749+ItzCrazyKns@users.noreply.github.com> Date: Tue, 7 Jan 2025 12:26:38 +0530 Subject: [PATCH 5/5] feat(chat-window): correctly handle server side WS closure --- ui/components/ChatWindow.tsx | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/ui/components/ChatWindow.tsx b/ui/components/ChatWindow.tsx index a44c550..ac15b37 100644 --- a/ui/components/ChatWindow.tsx +++ b/ui/components/ChatWindow.tsx @@ -270,12 +270,14 @@ const useSocket = ( }; } catch (error) { console.debug(new Date(), 'ws:error', error); + setIsWSReady(false); attemptReconnect(); } }; const attemptReconnect = () => { retryCountRef.current += 1; + if (retryCountRef.current > MAX_RETRIES) { console.debug(new Date(), 'ws:max_retries'); setError(true); @@ -303,14 +305,14 @@ const useSocket = ( connectWs(); return () => { - isCleaningUpRef.current = true; if (reconnectTimeoutRef.current) { clearTimeout(reconnectTimeoutRef.current); } if (wsRef.current?.readyState === WebSocket.OPEN) { wsRef.current.close(); + isCleaningUpRef.current = true; + console.debug(new Date(), 'ws:cleanup'); } - console.debug(new Date(), 'ws:cleanup'); }; }, [url, setIsWSReady, setError]); @@ -456,6 +458,8 @@ const ChatWindow = ({ id }: { id?: string }) => { if (isMessagesLoaded && isWSReady) { setIsReady(true); console.debug(new Date(), 'app:ready'); + } else { + setIsReady(false); } }, [isMessagesLoaded, isWSReady]);