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 ? (