Fix new Frontend 11

frontend/src/api/websocket/hooks.ts

@@ -98,72 +98,102 @@ export const useWebSocket = (config: WebSocketConfig) => {
 export const useTrainingWebSocket = (jobId: string, tenantId?: string) => {
   const [jobUpdates, setJobUpdates] = useState<any[]>([]);
 
-  // 🚀 FIX 1: Construct URL with actual tenantId (get from localStorage if not provided)
+  // Get tenant ID reliably
   const actualTenantId = tenantId || (() => {
-    const userData = localStorage.getItem('user_data');
-    if (userData) {
-      try {
+    try {
+      const userData = localStorage.getItem('user_data');
+      if (userData) {
         const parsed = JSON.parse(userData);
         return parsed.current_tenant_id || parsed.tenant_id;
-      } catch (e) {
-        console.error('Failed to parse user data for tenant ID:', e);
       }
+      
+      const authData = localStorage.getItem('auth_data');
+      if (authData) {
+        const parsed = JSON.parse(authData);
+        return parsed.tenant_id;
+      }
+    } catch (e) {
+      console.error('Failed to parse tenant ID from storage:', e);
     }
     return null;
   })();
 
-  const config: WebSocketConfig = {
-    // 🚀 FIX: Use actual tenant ID instead of placeholder
+  const config = {
     url: actualTenantId 
       ? `ws://localhost:8002/api/v1/ws/tenants/${actualTenantId}/training/jobs/${jobId}/live`
       : `ws://localhost:8002/api/v1/ws/tenants/unknown/training/jobs/${jobId}/live`,
     reconnect: true,
+    reconnectInterval: 3000,
+    maxReconnectAttempts: 10
   };
 
-  const { status, connect, disconnect, addMessageHandler, isConnected, lastMessage } = useWebSocket(config);
+  const { 
+    status, 
+    connect, 
+    disconnect, 
+    addMessageHandler, 
+    isConnected, 
+    lastMessage,
+    sendMessage 
+  } = useWebSocket(config);
 
-  // 🚀 FIX 2: Handle ALL message types, not just specific ones
-  useEffect(() => {
-    addMessageHandler((message) => {
-      console.log('🔥 WebSocket message received:', message);
-      
-      // Handle all training-related message types
-      const trainingMessageTypes = [
-        'progress', 'training_progress',
-        'completed', 'training_completed', 
-        'failed', 'training_failed',
-        'error', 'training_error',
-        'started', 'training_started',
-        'heartbeat', 'initial_status'
-      ];
-      
-      if (trainingMessageTypes.includes(message.type)) {
-        console.log('✅ Processing training message:', message.type, message.data);
-        setJobUpdates(prev => [message, ...prev.slice(0, 99)]); // Keep full message object
-      } else {
-        console.log('ℹ️ Unhandled message type:', message.type);
-        // Still add to updates for debugging
-        setJobUpdates(prev => [message, ...prev.slice(0, 99)]);
-      }
-    });
-  }, [addMessageHandler]);
-
-  // 🚀 FIX 3: Log connection attempts and status
-  useEffect(() => {
-    console.log('🔌 WebSocket config:', {
-      url: config.url,
-      jobId,
-      tenantId: actualTenantId,
-      status
-    });
-  }, [config.url, jobId, actualTenantId, status]);
-
-  // 🚀 FIX 4: Debug latest message
-  useEffect(() => {
-    if (lastMessage) {
-      console.log('📨 Latest WebSocket message:', lastMessage);
+  // Enhanced message handler
+  const handleWebSocketMessage = useCallback((message: any) => {
+    // Handle different message structures
+    let processedMessage = message;
+    
+    // If message has nested data, flatten it for easier processing
+    if (message.data && typeof message.data === 'object') {
+      processedMessage = {
+        ...message,
+        // Merge data properties to root level for backward compatibility
+        ...message.data
+      };
     }
-  }, [lastMessage]);
+
+    // Comprehensive message type handling
+    const trainingMessageTypes = [
+      'progress', 'training_progress',
+      'completed', 'training_completed', 
+      'failed', 'training_failed',
+      'error', 'training_error',
+      'started', 'training_started',
+      'heartbeat', 'initial_status',
+      'status_update'
+    ];
+    
+    if (trainingMessageTypes.includes(message.type)) {
+      // Add to updates array with processed message
+      setJobUpdates(prev => {
+        const newUpdates = [processedMessage, ...prev.slice(0, 49)]; // Keep last 50 messages
+        return newUpdates;
+      });
+    } else {
+      // Still add to updates for debugging purposes
+      setJobUpdates(prev => [processedMessage, ...prev.slice(0, 49)]);
+    }
+  }, []);
+
+  // Set up message handler when hook initializes
+  useEffect(() => {
+    addMessageHandler(handleWebSocketMessage);
+  }, [addMessageHandler, handleWebSocketMessage]);
+
+  // Send periodic ping to keep connection alive
+  useEffect(() => {
+    if (isConnected) {
+      const pingInterval = setInterval(() => {
+        sendMessage({
+          type: 'ping',
+          data: undefined
+        });
+      }, 30000); // Every 30 seconds
+
+      return () => {
+        clearInterval(pingInterval);
+      };
+    }
+  }, [isConnected, sendMessage]);
 
   return {
     status,
@@ -171,9 +201,16 @@ export const useTrainingWebSocket = (jobId: string, tenantId?: string) => {
     connect,
     disconnect,
     isConnected,
-    lastMessage, // Expose for debugging
-    tenantId: actualTenantId, // Expose for debugging
-    wsUrl: config.url, // Expose for debugging
+    lastMessage,
+    tenantId: actualTenantId,
+    wsUrl: config.url,
+    // Manual refresh function
+    refreshConnection: useCallback(() => {
+      disconnect();
+      setTimeout(() => {
+        connect();
+      }, 1000);
+    }, [connect, disconnect])
   };
 };