Fix websocket reconnecting dialog page refresh

ui/src/hooks/useEnterpriseWebSocket.ts

@@ -51,6 +51,30 @@ export function useEnterpriseWebSocket(
   const wsRef = useRef<WebSocket | null>(null);
   const reconnectTimeoutRef = useRef<NodeJS.Timeout | null>(null);
   const shouldReconnectRef = useRef(true);
+  const reconnectAttemptsRef = useRef(0);
+
+  // Store callbacks in refs to avoid reconnection when they change
+  const onMessageRef = useRef(onMessage);
+  const onErrorRef = useRef(onError);
+  const onOpenRef = useRef(onOpen);
+  const onCloseRef = useRef(onClose);
+
+  // Update refs when callbacks change
+  useEffect(() => {
+    onMessageRef.current = onMessage;
+  }, [onMessage]);
+
+  useEffect(() => {
+    onErrorRef.current = onError;
+  }, [onError]);
+
+  useEffect(() => {
+    onOpenRef.current = onOpen;
+  }, [onOpen]);
+
+  useEffect(() => {
+    onCloseRef.current = onClose;
+  }, [onClose]);
 
   const getWebSocketUrl = useCallback(() => {
     const protocol = window.location.protocol === 'https:' ? 'wss:' : 'ws:';
@@ -82,10 +106,11 @@ export function useEnterpriseWebSocket(
         // console.log('[WebSocket] Connected to enterprise WebSocket');
         setIsConnected(true);
         setReconnectAttempts(0);
+        reconnectAttemptsRef.current = 0;
         shouldReconnectRef.current = true;
 
-        if (onOpen) {
-          onOpen();
+        if (onOpenRef.current) {
+          onOpenRef.current();
         }
       };
 
@@ -96,8 +121,8 @@ export function useEnterpriseWebSocket(
 
           setLastMessage(message);
 
-          if (onMessage) {
-            onMessage(message);
+          if (onMessageRef.current) {
+            onMessageRef.current(message);
           }
         } catch (error) {
           console.error('[WebSocket] Failed to parse message:', error);
@@ -108,34 +133,36 @@ export function useEnterpriseWebSocket(
         console.error('[WebSocket] Error:', error);
         setIsConnected(false);
 
-        if (onError) {
-          onError(error);
+        if (onErrorRef.current) {
+          onErrorRef.current(error);
         }
       };
 
       wsRef.current.onclose = () => {
         // console.log('[WebSocket] Connection closed');
         setIsConnected(false);
 
-        if (onClose) {
-          onClose();
+        if (onCloseRef.current) {
+          onCloseRef.current();
         }
 
         // Attempt reconnection if enabled and within retry limit
+        const currentAttempts = reconnectAttemptsRef.current;
         if (
           shouldReconnectRef.current &&
           autoReconnect &&
-          reconnectAttempts < maxReconnectAttempts
+          currentAttempts < maxReconnectAttempts
         ) {
           // console.log(
-          //   `[WebSocket] Attempting reconnection in ${reconnectInterval}ms (attempt ${reconnectAttempts + 1}/${maxReconnectAttempts})`
+          //   `[WebSocket] Attempting reconnection in ${reconnectInterval}ms (attempt ${currentAttempts + 1}/${maxReconnectAttempts})`
           // );
 
           reconnectTimeoutRef.current = setTimeout(() => {
+            reconnectAttemptsRef.current += 1;
             setReconnectAttempts((prev) => prev + 1);
             connect();
           }, reconnectInterval);
-        } else if (reconnectAttempts >= maxReconnectAttempts) {
+        } else if (currentAttempts >= maxReconnectAttempts) {
           console.error('[WebSocket] Max reconnection attempts reached');
         }
       };
@@ -145,14 +172,9 @@ export function useEnterpriseWebSocket(
     }
   }, [
     getWebSocketUrl,
-    onOpen,
-    onMessage,
-    onError,
-    onClose,
     autoReconnect,
     reconnectInterval,
     maxReconnectAttempts,
-    reconnectAttempts,
   ]);
 
   const disconnect = useCallback(() => {
@@ -201,9 +223,15 @@ export function useEnterpriseWebSocket(
   }, []); // Empty dependency array - only run on mount/unmount
 
   // Handle page visibility changes (reconnect when page becomes visible)
+  // Store isConnected in ref to avoid effect recreation
+  const isConnectedRef = useRef(isConnected);
+  useEffect(() => {
+    isConnectedRef.current = isConnected;
+  }, [isConnected]);
+
   useEffect(() => {
     const handleVisibilityChange = () => {
-      if (document.visibilityState === 'visible' && !isConnected) {
+      if (document.visibilityState === 'visible' && !isConnectedRef.current) {
         // console.log('[WebSocket] Page visible, attempting reconnection');
         connect();
       }
@@ -214,7 +242,8 @@ export function useEnterpriseWebSocket(
     return () => {
       document.removeEventListener('visibilitychange', handleVisibilityChange);
     };
-  }, [isConnected, connect]);
+    // eslint-disable-next-line react-hooks/exhaustive-deps
+  }, []); // Only set up listener once on mount
 
   return {
     isConnected,

ui/src/hooks/useWebSocket.ts

@@ -40,6 +40,29 @@ export function useWebSocket({
   const shouldReconnectRef = useRef(true);
   const reconnectAttemptsRef = useRef(0);
 
+  // Store callbacks in refs to avoid reconnection when they change
+  const onMessageRef = useRef(onMessage);
+  const onErrorRef = useRef(onError);
+  const onOpenRef = useRef(onOpen);
+  const onCloseRef = useRef(onClose);
+
+  // Update refs when callbacks change
+  useEffect(() => {
+    onMessageRef.current = onMessage;
+  }, [onMessage]);
+
+  useEffect(() => {
+    onErrorRef.current = onError;
+  }, [onError]);
+
+  useEffect(() => {
+    onOpenRef.current = onOpen;
+  }, [onOpen]);
+
+  useEffect(() => {
+    onCloseRef.current = onClose;
+  }, [onClose]);
+
   const connect = useCallback(() => {
     try {
       const ws = new WebSocket(url);
@@ -49,27 +72,27 @@ export function useWebSocket({
         setIsConnected(true);
         setReconnectAttempts(0);
         reconnectAttemptsRef.current = 0;
-        onOpen?.();
+        onOpenRef.current?.();
       };
 
       ws.onmessage = (event) => {
         try {
           const data = JSON.parse(event.data);
-          onMessage(data);
+          onMessageRef.current(data);
         } catch (error) {
           console.error('Failed to parse WebSocket message:', error);
         }
       };
 
       ws.onerror = (error) => {
         console.error('WebSocket error:', error);
-        onError?.(error);
+        onErrorRef.current?.(error);
       };
 
       ws.onclose = () => {
         // console.log(`WebSocket closed: ${url}`);
         setIsConnected(false);
-        onClose?.();
+        onCloseRef.current?.();
 
         // Attempt reconnection with exponential backoff
         const currentAttempts = reconnectAttemptsRef.current;
@@ -91,7 +114,7 @@ export function useWebSocket({
     } catch (error) {
       console.error('Failed to create WebSocket connection:', error);
     }
-  }, [url, onMessage, onError, onOpen, onClose, reconnectInterval, maxReconnectAttempts]);
+  }, [url, reconnectInterval, maxReconnectAttempts]);
 
   const sendMessage = useCallback((message: any) => {
     if (wsRef.current?.readyState === WebSocket.OPEN) {
@@ -122,7 +145,8 @@ export function useWebSocket({
     return () => {
       close();
     };
-  }, [connect, close]);
+    // eslint-disable-next-line react-hooks/exhaustive-deps
+  }, []); // Only run on mount/unmount, not when connect/close change
 
   return {
     isConnected,

ui/src/pages/Dashboard.tsx

@@ -1,4 +1,4 @@
-import { useState, useRef } from 'react';
+import { useState, useRef, useCallback } from 'react';
 import { Grid, Paper, Typography, Box, Card, CardContent, Chip } from '@mui/material';
 import {
   Computer as ComputerIcon,
@@ -72,7 +72,8 @@ export default function Dashboard() {
   const { data: repositories = [], isLoading: reposLoading } = useRepositories();
 
   // Real-time sessions updates via WebSocket with notifications
-  const baseSessionsWs = useSessionsWebSocket((updatedSessions) => {
+  // Wrap callback in useCallback to prevent reconnection loop
+  const handleSessionsUpdate = useCallback((updatedSessions: Session[]) => {
     // Filter to only show current user's sessions
     const userSessions = username
       ? updatedSessions.filter((s: Session) => s.user === username)
@@ -93,15 +94,20 @@ export default function Dashboard() {
     });
 
     setSessions(userSessions);
-  });
+  }, [username, addNotification]);
+
+  const baseSessionsWs = useSessionsWebSocket(handleSessionsUpdate);
 
   // Enhanced WebSocket with connection quality and manual reconnect
   const sessionsWs = useEnhancedWebSocket(baseSessionsWs);
 
   // Real-time metrics updates via WebSocket
-  const metricsWs = useMetricsWebSocket((updatedMetrics) => {
+  // Wrap callback in useCallback to prevent reconnection loop
+  const handleMetricsUpdate = useCallback((updatedMetrics: any) => {
     setMetrics(updatedMetrics);
-  });
+  }, []);
+
+  const metricsWs = useMetricsWebSocket(handleMetricsUpdate);
 
   const stats = [
     {

ui/src/pages/SessionViewer.tsx

@@ -1,4 +1,4 @@
-import { useState, useEffect, useRef } from 'react';
+import { useState, useEffect, useRef, useCallback } from 'react';
 import { useParams, useNavigate } from 'react-router-dom';
 import {
   Box,
@@ -132,7 +132,8 @@ export default function SessionViewer() {
   const { addNotification } = useNotificationQueue();
 
   // Real-time session updates via WebSocket with notifications
-  const baseWebSocket = useSessionsWebSocket((updatedSessions) => {
+  // Wrap callback in useCallback to prevent reconnection loop
+  const handleSessionUpdate = useCallback((updatedSessions: any[]) => {
     if (!sessionId) return;
 
     // Find this session in the update
@@ -157,7 +158,9 @@ export default function SessionViewer() {
       prevStateRef.current = updatedSession.state;
       setSession(updatedSession);
     }
-  });
+  }, [sessionId, session, addNotification]);
+
+  const baseWebSocket = useSessionsWebSocket(handleSessionUpdate);
 
   // Enhanced WebSocket with connection quality and manual reconnect
   const enhanced = useEnhancedWebSocket(baseWebSocket);

ui/src/pages/SharedSessions.tsx

@@ -1,4 +1,4 @@
-import { useState, useEffect, useRef } from 'react';
+import { useState, useEffect, useRef, useCallback } from 'react';
 import {
   Box,
   Typography,
@@ -112,7 +112,8 @@ export default function SharedSessions() {
   const { addNotification } = useNotificationQueue();
 
   // Real-time session updates via WebSocket with notifications
-  const baseWebSocket = useSessionsWebSocket((updatedSessions) => {
+  // Wrap callback in useCallback to prevent reconnection loop
+  const handleSessionsUpdate = useCallback((updatedSessions: any[]) => {
     if (!currentUser?.id || sessions.length === 0) return;
 
     // Update shared sessions with real-time data and show notifications for changes
@@ -144,7 +145,9 @@ export default function SharedSessions() {
     });
 
     setSessions(updatedSharedSessions);
-  });
+  }, [currentUser?.id, sessions, addNotification]);
+
+  const baseWebSocket = useSessionsWebSocket(handleSessionsUpdate);
 
   // Enhanced WebSocket with connection quality and manual reconnect
   const enhanced = useEnhancedWebSocket(baseWebSocket);