Simplify the onboardinf flow components 2

frontend/src/components/domain/onboarding/steps/MLTrainingStep.tsx

@@ -1,4 +1,4 @@
-import React, { useState, useEffect } from 'react';
+import React, { useState, useCallback } from 'react';
 import { Button } from '../../../ui/Button';
 import { useCurrentTenant } from '../../../../stores/tenant.store';
 import { useCreateTrainingJob, useTrainingWebSocket } from '../../../../api/hooks/training';
@@ -32,50 +32,59 @@ export const MLTrainingStep: React.FC<MLTrainingStepProps> = ({
   const currentTenant = useCurrentTenant();
   const createTrainingJob = useCreateTrainingJob();
 
-  // WebSocket for real-time training progress
-  const trainingWebSocket = useTrainingWebSocket(
+  // Memoized WebSocket callbacks to prevent reconnections
+  const handleProgress = useCallback((data: any) => {
+    setTrainingProgress({
+      stage: 'training',
+      progress: data.data?.progress || 0,
+      message: data.data?.message || 'Entrenando modelo...',
+      currentStep: data.data?.current_step,
+      estimatedTimeRemaining: data.data?.estimated_time_remaining
+    });
+  }, []);
+
+  const handleCompleted = useCallback((_data: any) => {
+    setTrainingProgress({
+      stage: 'completed',
+      progress: 100,
+      message: 'Entrenamiento completado exitosamente'
+    });
+    setIsTraining(false);
+    
+    setTimeout(() => {
+      onComplete({
+        jobId: jobId,
+        success: true,
+        message: 'Modelo entrenado correctamente'
+      });
+    }, 2000);
+  }, [onComplete, jobId]);
+
+  const handleError = useCallback((data: any) => {
+    setError(data.data?.error || data.error || 'Error durante el entrenamiento');
+    setIsTraining(false);
+    setTrainingProgress(null);
+  }, []);
+
+  const handleStarted = useCallback((_data: any) => {
+    setTrainingProgress({
+      stage: 'starting',
+      progress: 5,
+      message: 'Iniciando entrenamiento del modelo...'
+    });
+  }, []);
+
+  // WebSocket for real-time training progress - only connect when we have a jobId
+  const { isConnected, connectionError } = useTrainingWebSocket(
     currentTenant?.id || '',
     jobId || '',
     undefined, // token will be handled by the service
-    {
-      onProgress: (data) => {
-        setTrainingProgress({
-          stage: 'training',
-          progress: data.progress?.percentage || 0,
-          message: data.message || 'Entrenando modelo...',
-          currentStep: data.progress?.current_step,
-          estimatedTimeRemaining: data.progress?.estimated_time_remaining
-        });
-      },
-      onCompleted: (data) => {
-        setTrainingProgress({
-          stage: 'completed',
-          progress: 100,
-          message: 'Entrenamiento completado exitosamente'
-        });
-        setIsTraining(false);
-        
-        setTimeout(() => {
-          onComplete({
-            jobId: jobId,
-            success: true,
-            message: 'Modelo entrenado correctamente'
-          });
-        }, 2000);
-      },
-      onError: (data) => {
-        setError(data.error || 'Error durante el entrenamiento');
-        setIsTraining(false);
-        setTrainingProgress(null);
-      },
-      onStarted: (data) => {
-        setTrainingProgress({
-          stage: 'starting',
-          progress: 5,
-          message: 'Iniciando entrenamiento del modelo...'
-        });
-      }
-    }
+    jobId ? {
+      onProgress: handleProgress,
+      onCompleted: handleCompleted,
+      onError: handleError,
+      onStarted: handleStarted
+    } : undefined
   );
 
   const handleStartTraining = async () => {
@@ -201,9 +210,16 @@ export const MLTrainingStep: React.FC<MLTrainingStepProps> = ({
 
                     <div className="flex justify-between text-xs text-[var(--text-tertiary)]">
                       <span>{trainingProgress.currentStep || 'Procesando...'}</span>
-                      {trainingProgress.estimatedTimeRemaining && (
-                        <span>Tiempo estimado: {formatTime(trainingProgress.estimatedTimeRemaining)}</span>
-                      )}
+                      <div className="flex items-center gap-2">
+                        {jobId && (
+                          <span className={`text-xs ${isConnected ? 'text-green-500' : 'text-red-500'}`}>
+                            {isConnected ? '🟢 Conectado' : '🔴 Desconectado'}
+                          </span>
+                        )}
+                        {trainingProgress.estimatedTimeRemaining && (
+                          <span>Tiempo estimado: {formatTime(trainingProgress.estimatedTimeRemaining)}</span>
+                        )}
+                      </div>
                     </div>
                   </div>
                 )}
@@ -224,9 +240,9 @@ export const MLTrainingStep: React.FC<MLTrainingStepProps> = ({
         </ul>
       </div>
 
-      {error && (
+      {(error || connectionError) && (
         <div className="bg-[var(--color-error)]/10 border border-[var(--color-error)]/20 rounded-lg p-4">
-          <p className="text-[var(--color-error)]">{error}</p>
+          <p className="text-[var(--color-error)]">{error || connectionError}</p>
         </div>
       )}