feat(websocket): enhance WebSocket error handling and connection mana…

apps/the_monkeys/src/app/edit/[blogId]/page.tsx

@@ -35,6 +35,8 @@ const EditPage = ({ params }: { params: { blogId: string } }) => {
   const blogTopicsRef = useRef<string[]>([]);
   const webSocketRef = useRef<WebSocket | null>(null);
   const reconnectTimeoutRef = useRef<NodeJS.Timeout | null>(null);
+  const heartbeatIntervalRef = useRef<NodeJS.Timeout | null>(null);
+  const lastHeartbeatRef = useRef<number>(Date.now());
 
   // State variables
   const [data, setData] = useState<OutputData | null>(null);
@@ -44,6 +46,7 @@ const EditPage = ({ params }: { params: { blogId: string } }) => {
   const [token, setToken] = useState<string | null>(null);
   const [isConnected, setIsConnected] = useState(false);
   const [connectionStatus, setConnectionStatus] = useState('Connecting...');
+  const [heartbeatEnabled, setHeartbeatEnabled] = useState(true);
 
   const accountId = session?.account_id;
   const username = session?.username;
@@ -122,67 +125,182 @@ const EditPage = ({ params }: { params: { blogId: string } }) => {
       }
 
       setConnectionStatus('Connecting...');
-      const ws = new WebSocket(
-        `${WSS_URL_V2}/blog/draft/${blogId}?token=${token}`
-      );
-
-      ws.onopen = () => {
-        if (!isMounted) return;
-        console.log('WebSocket connected 🟢');
-        setIsConnected(true);
-        setConnectionStatus('Connected');
-        retryCount = 0;
-
-        // Send latest data on reconnect
-        if (dataRef.current) {
-          const formattedData = formatData(
-            dataRef.current,
-            accountIdRef.current,
-            blogTopicsRef.current
-          );
-          ws.send(JSON.stringify(formattedData));
-          setIsSaving(true);
-        }
-      };
-
-      ws.onmessage = (event) => {
-        setIsSaving(false);
-        // Optional: Handle any incoming messages from server
-      };
 
-      ws.onclose = (event) => {
-        if (!isMounted) return;
-        console.log('WebSocket closed 🔴', event.code, event.reason);
+      try {
+        const ws = new WebSocket(
+          `${WSS_URL_V2}/blog/draft/${blogId}?token=${token}`
+        );
+
+        ws.onopen = () => {
+          if (!isMounted) return;
+          console.log('WebSocket connected 🟢');
+          setIsConnected(true);
+          setConnectionStatus('Connected');
+          retryCount = 0;
+          lastHeartbeatRef.current = Date.now();
+
+          // Clear any existing heartbeat interval before setting a new one
+          if (heartbeatIntervalRef.current) {
+            clearInterval(heartbeatIntervalRef.current);
+          }
+
+          // Start heartbeat to monitor connection health (reduce frequency to avoid server issues)
+          if (heartbeatEnabled) {
+            heartbeatIntervalRef.current = setInterval(() => {
+              if (ws.readyState === WebSocket.OPEN && isMounted) {
+                // Send a ping message to keep connection alive
+                try {
+                  // Use a simpler ping format that the server might handle better
+                  ws.send(JSON.stringify({ type: 'ping' }));
+                  lastHeartbeatRef.current = Date.now();
+                  console.log('Heartbeat sent 💓');
+                } catch (error) {
+                  console.warn('Failed to send heartbeat:', error);
+                  // Disable heartbeat if it's causing issues
+                  setHeartbeatEnabled(false);
+                  if (heartbeatIntervalRef.current) {
+                    clearInterval(heartbeatIntervalRef.current);
+                  }
+                }
+              }
+            }, 60000); // Increase to 60 seconds to be more conservative
+          }
+
+          // Send latest data on reconnect
+          if (dataRef.current) {
+            const formattedData = formatData(
+              dataRef.current,
+              accountIdRef.current,
+              blogTopicsRef.current
+            );
+            ws.send(JSON.stringify(formattedData));
+            setIsSaving(true);
+          }
+        };
+
+        ws.onmessage = (event) => {
+          setIsSaving(false);
+          lastHeartbeatRef.current = Date.now();
+
+          try {
+            const message = JSON.parse(event.data);
+            // Handle pong responses to keep connection alive
+            if (message.type === 'pong') {
+              console.log('Received pong - connection healthy 💚');
+              return;
+            }
+            // Handle other server messages here if needed
+            console.log('Received message:', message);
+          } catch (error) {
+            // Message might not be JSON, that's okay
+            console.log('Received non-JSON message:', event.data);
+          }
+        };
+
+        ws.onclose = (event) => {
+          if (!isMounted) return;
+          console.log('WebSocket closed 🔴', event.code, event.reason);
+          setIsConnected(false);
+          setIsSaving(false);
+
+          // Clear heartbeat interval when connection closes
+          if (heartbeatIntervalRef.current) {
+            clearInterval(heartbeatIntervalRef.current);
+            heartbeatIntervalRef.current = null;
+          }
+
+          // Handle different close codes with immediate vs delayed reconnection
+          let statusMessage = 'Disconnected';
+          let shouldReconnectImmediately = false;
+
+          if (event.code === 1006) {
+            statusMessage = 'Connection lost unexpectedly';
+            shouldReconnectImmediately = true; // Network issues - reconnect immediately
+          } else if (event.code === 1000) {
+            statusMessage = 'Disconnected normally';
+            console.log('Normal WebSocket closure - not reconnecting');
+            return; // Don't reconnect for normal closure
+          } else if (event.code === 1001) {
+            statusMessage = 'Server going away';
+            shouldReconnectImmediately = true; // Server restart - reconnect immediately
+          } else if (event.code === 1008 || event.code === 1002) {
+            statusMessage = 'Connection failed - invalid token';
+            setConnectionStatus(
+              'Authentication failed. Please refresh the page.'
+            );
+            return; // Don't reconnect for auth issues
+          } else if (event.code === 1011) {
+            statusMessage = 'Server error';
+            shouldReconnectImmediately = false; // Server error - use backoff
+          } else if (event.code === 1005) {
+            // No status code - might be a normal close
+            console.log(
+              'WebSocket closed without status code - not reconnecting'
+            );
+            return;
+          } else {
+            statusMessage = 'Connection closed';
+            shouldReconnectImmediately = true; // Unknown reason - try immediately
+          }
+
+          setConnectionStatus(statusMessage);
+
+          // Add a small delay before attempting reconnection to avoid rapid reconnection loops
+          const baseDelay = shouldReconnectImmediately ? 1000 : 2000;
+
+          // Reconnect logic - immediate for certain scenarios, exponential backoff for others
+          if (retryCount < MAX_RETRIES) {
+            let delay = 0;
+
+            if (shouldReconnectImmediately && retryCount === 0) {
+              // First retry for network/server issues - small delay
+              delay = baseDelay;
+              setConnectionStatus('Reconnecting...');
+            } else {
+              // Use exponential backoff for subsequent retries or server errors
+              delay = Math.min(baseDelay * Math.pow(2, retryCount), 15000);
+              setConnectionStatus(
+                `Reconnecting in ${Math.round(delay / 1000)}s...`
+              );
+            }
+
+            reconnectTimeoutRef.current = setTimeout(() => {
+              if (isMounted) {
+                retryCount++;
+                connectWebSocket();
+              }
+            }, delay);
+          } else {
+            setConnectionStatus('Connection failed. Please refresh the page.');
+          }
+        };
+
+        ws.onerror = (error) => {
+          console.error('WebSocket error ⭕', error);
+          setConnectionStatus('Connection error occurred');
+          // Don't immediately close, let onclose handle the reconnection
+        };
+
+        webSocketRef.current = ws;
+      } catch (error) {
+        console.error('Failed to create WebSocket connection:', error);
+        setConnectionStatus('Failed to establish connection');
         setIsConnected(false);
-        setIsSaving(false);
-        setConnectionStatus('Disconnected');
 
-        // Reconnect logic with exponential backoff
+        // Retry after a delay if we haven't exceeded max retries
         if (retryCount < MAX_RETRIES) {
           const delay = Math.min(1000 * 2 ** retryCount, 30000);
-          setConnectionStatus(
-            `Reconnecting in ${Math.round(delay / 1000)}s...`
-          );
-          reconnectTimeoutRef.current = setTimeout(() => {
+          setTimeout(() => {
             retryCount++;
             connectWebSocket();
           }, delay);
-        } else {
-          setConnectionStatus('Connection failed. Please refresh the page.');
         }
-      };
-
-      ws.onerror = (error) => {
-        console.error('WebSocket error ⭕', error);
-        ws.close();
-      };
-
-      webSocketRef.current = ws;
+      }
     };
 
     connectWebSocket();
 
-    // Handle tab visibility changes
+    // Handle tab visibility changes and network connectivity
     const handleVisibilityChange = () => {
       if (document.visibilityState === 'visible' && !isConnected && token) {
         // Reconnect immediately when tab becomes visible
@@ -193,7 +311,25 @@ const EditPage = ({ params }: { params: { blogId: string } }) => {
       }
     };
 
+    const handleOnline = () => {
+      if (!isConnected && token) {
+        console.log('Network back online - attempting reconnection');
+        setConnectionStatus('Network restored - reconnecting...');
+        if (reconnectTimeoutRef.current) {
+          clearTimeout(reconnectTimeoutRef.current);
+        }
+        connectWebSocket();
+      }
+    };
+
+    const handleOffline = () => {
+      setConnectionStatus('Network offline');
+      setIsConnected(false);
+    };
+
     document.addEventListener('visibilitychange', handleVisibilityChange);
+    window.addEventListener('online', handleOnline);
+    window.addEventListener('offline', handleOffline);
 
     return () => {
       isMounted = false;
@@ -203,28 +339,76 @@ const EditPage = ({ params }: { params: { blogId: string } }) => {
       if (reconnectTimeoutRef.current) {
         clearTimeout(reconnectTimeoutRef.current);
       }
+      if (heartbeatIntervalRef.current) {
+        clearInterval(heartbeatIntervalRef.current);
+      }
       document.removeEventListener('visibilitychange', handleVisibilityChange);
+      window.removeEventListener('online', handleOnline);
+      window.removeEventListener('offline', handleOffline);
     };
+    // eslint-disable-next-line react-hooks/exhaustive-deps
   }, [token, blogId, formatData]);
 
   // Auto-save when data changes
   useEffect(() => {
-    if (!data || !isConnected) return;
+    if (!data || !isConnected || !webSocketRef.current) return;
 
     const formattedData = formatData(data, accountId, blogTopics);
 
     try {
-      if (webSocketRef.current?.readyState === WebSocket.OPEN) {
+      if (webSocketRef.current.readyState === WebSocket.OPEN) {
         webSocketRef.current.send(JSON.stringify(formattedData));
         setIsSaving(true);
+      } else if (webSocketRef.current.readyState === WebSocket.CONNECTING) {
+        // Wait for connection to open, then send
+        console.log('WebSocket still connecting, will send when ready');
+      } else {
+        console.log(
+          'WebSocket not ready, state:',
+          webSocketRef.current.readyState
+        );
+        setIsConnected(false);
       }
     } catch (error) {
       console.error('Error sending data:', error);
       setIsSaving(false);
+      setIsConnected(false);
+      toast({
+        variant: 'destructive',
+        title: 'Connection Error',
+        description: 'Failed to save changes. Please check your connection.',
+      });
     }
   }, [data, blogTopics, isConnected, accountId, formatData]);
 
-  // Handle blog publishing
+  // Manual reconnect function
+  const handleManualReconnect = useCallback(() => {
+    if (reconnectTimeoutRef.current) {
+      clearTimeout(reconnectTimeoutRef.current);
+    }
+
+    // Reset connection state and try to reconnect
+    setIsConnected(false);
+    setConnectionStatus('Reconnecting...');
+
+    // Close existing connection if any
+    if (webSocketRef.current) {
+      webSocketRef.current.close();
+    }
+
+    // Trigger reconnection by updating token (this will restart the WebSocket effect)
+    if (session) {
+      axiosInstance
+        .get('/auth/ws-token')
+        .then((response) => {
+          setToken(response.data.token);
+        })
+        .catch((error) => {
+          console.error('Failed to refresh token:', error);
+          setConnectionStatus('Failed to refresh connection token');
+        });
+    }
+  }, [session]);
   const handlePublishStep = useCallback(async () => {
     if (!data || data.blocks.length === 0) {
       toast({
@@ -289,10 +473,41 @@ const EditPage = ({ params }: { params: { blogId: string } }) => {
             <div className='flex items-center gap-2 text-sm'>
               <span
                 className={`inline-block w-3 h-3 rounded-full ${
-                  isConnected ? 'bg-green-500' : 'bg-yellow-500'
+                  isConnected
+                    ? 'bg-green-500'
+                    : connectionStatus.includes('Reconnecting')
+                      ? 'bg-yellow-500 animate-pulse'
+                      : connectionStatus.includes('failed') ||
+                          connectionStatus.includes('Authentication')
+                        ? 'bg-red-500'
+                        : 'bg-yellow-500'
                 }`}
               />
-              {connectionStatus}
+              <span
+                className={`${
+                  connectionStatus.includes('failed') ||
+                  connectionStatus.includes('Authentication')
+                    ? 'text-red-600 dark:text-red-400'
+                    : connectionStatus.includes('Reconnecting')
+                      ? 'text-yellow-600 dark:text-yellow-400'
+                      : isConnected
+                        ? 'text-green-600 dark:text-green-400'
+                        : 'text-gray-600 dark:text-gray-400'
+                }`}
+              >
+                {connectionStatus}
+              </span>
+
+              {/* Manual reconnect button for failed connections */}
+              {(connectionStatus.includes('failed') ||
+                connectionStatus.includes('Authentication')) && (
+                <button
+                  onClick={handleManualReconnect}
+                  className='ml-2 px-2 py-1 text-xs bg-blue-500 text-white rounded hover:bg-blue-600 transition-colors'
+                >
+                  Reconnect
+                </button>
+              )}
             </div>
 
             <div className='flex gap-2'>