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] 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]);