708 lines
26 KiB
TypeScript
708 lines
26 KiB
TypeScript
/**
|
||
* Training React Query hooks
|
||
* Provides data fetching, caching, and state management for training operations
|
||
*/
|
||
|
||
import * as React from 'react';
|
||
import { useMutation, useQuery, useQueryClient, UseQueryOptions, UseMutationOptions } from '@tanstack/react-query';
|
||
import { trainingService } from '../services/training';
|
||
import { ApiError, apiClient } from '../client/apiClient';
|
||
import { useAuthStore } from '../../stores/auth.store';
|
||
import {
|
||
HTTP_POLLING_INTERVAL_MS,
|
||
HTTP_POLLING_DEBOUNCE_MS,
|
||
WEBSOCKET_HEARTBEAT_INTERVAL_MS,
|
||
WEBSOCKET_MAX_RECONNECT_ATTEMPTS,
|
||
WEBSOCKET_RECONNECT_INITIAL_DELAY_MS,
|
||
WEBSOCKET_RECONNECT_MAX_DELAY_MS,
|
||
PROGRESS_DATA_ANALYSIS,
|
||
PROGRESS_TRAINING_RANGE_START,
|
||
PROGRESS_TRAINING_RANGE_END
|
||
} from '../../constants/training';
|
||
import type {
|
||
TrainingJobRequest,
|
||
TrainingJobResponse,
|
||
TrainingJobStatus,
|
||
SingleProductTrainingRequest,
|
||
ModelMetricsResponse,
|
||
TrainedModelResponse,
|
||
} from '../types/training';
|
||
|
||
// Query Keys Factory
|
||
export const trainingKeys = {
|
||
all: ['training'] as const,
|
||
jobs: {
|
||
all: () => [...trainingKeys.all, 'jobs'] as const,
|
||
status: (tenantId: string, jobId: string) =>
|
||
[...trainingKeys.jobs.all(), 'status', tenantId, jobId] as const,
|
||
},
|
||
models: {
|
||
all: () => [...trainingKeys.all, 'models'] as const,
|
||
lists: () => [...trainingKeys.models.all(), 'list'] as const,
|
||
list: (tenantId: string, params?: any) =>
|
||
[...trainingKeys.models.lists(), tenantId, params] as const,
|
||
details: () => [...trainingKeys.models.all(), 'detail'] as const,
|
||
detail: (tenantId: string, modelId: string) =>
|
||
[...trainingKeys.models.details(), tenantId, modelId] as const,
|
||
active: (tenantId: string, inventoryProductId: string) =>
|
||
[...trainingKeys.models.all(), 'active', tenantId, inventoryProductId] as const,
|
||
metrics: (tenantId: string, modelId: string) =>
|
||
[...trainingKeys.models.all(), 'metrics', tenantId, modelId] as const,
|
||
performance: (tenantId: string, modelId: string) =>
|
||
[...trainingKeys.models.all(), 'performance', tenantId, modelId] as const,
|
||
},
|
||
statistics: (tenantId: string) =>
|
||
[...trainingKeys.all, 'statistics', tenantId] as const,
|
||
} as const;
|
||
|
||
// Training Job Queries
|
||
export const useTrainingJobStatus = (
|
||
tenantId: string,
|
||
jobId: string,
|
||
options?: Omit<UseQueryOptions<TrainingJobStatus, ApiError>, 'queryKey' | 'queryFn'> & {
|
||
isWebSocketConnected?: boolean;
|
||
}
|
||
) => {
|
||
const { isWebSocketConnected, ...queryOptions } = options || {};
|
||
const [enablePolling, setEnablePolling] = React.useState(false);
|
||
|
||
// Debounce HTTP polling activation: wait after WebSocket disconnects
|
||
// This prevents race conditions where both WebSocket and HTTP are briefly active
|
||
React.useEffect(() => {
|
||
if (!isWebSocketConnected) {
|
||
const debounceTimer = setTimeout(() => {
|
||
setEnablePolling(true);
|
||
console.log(`🔄 HTTP polling enabled after ${HTTP_POLLING_DEBOUNCE_MS}ms debounce (WebSocket disconnected)`);
|
||
}, HTTP_POLLING_DEBOUNCE_MS);
|
||
|
||
return () => clearTimeout(debounceTimer);
|
||
} else {
|
||
setEnablePolling(false);
|
||
console.log('❌ HTTP polling disabled (WebSocket connected)');
|
||
}
|
||
}, [isWebSocketConnected]);
|
||
|
||
// Completely disable the query when WebSocket is connected or during debounce period
|
||
const isEnabled = !!tenantId && !!jobId && !isWebSocketConnected && enablePolling;
|
||
|
||
console.log('🔄 Training status query:', {
|
||
tenantId: !!tenantId,
|
||
jobId: !!jobId,
|
||
isWebSocketConnected,
|
||
enablePolling,
|
||
queryEnabled: isEnabled
|
||
});
|
||
|
||
return useQuery<TrainingJobStatus, ApiError>({
|
||
queryKey: trainingKeys.jobs.status(tenantId, jobId),
|
||
queryFn: () => {
|
||
console.log('📡 Executing HTTP training status query (WebSocket disconnected)');
|
||
return trainingService.getTrainingJobStatus(tenantId, jobId);
|
||
},
|
||
enabled: isEnabled, // Completely disable when WebSocket connected
|
||
refetchInterval: isEnabled ? (query) => {
|
||
// Only set up refetch interval if the query is enabled
|
||
const data = query.state.data;
|
||
|
||
// Stop polling if we get auth errors or training is completed
|
||
if (query.state.error && (query.state.error as any)?.status === 401) {
|
||
console.log('🚫 Stopping status polling due to auth error');
|
||
return false;
|
||
}
|
||
if (data?.status === 'completed' || data?.status === 'failed') {
|
||
console.log('🏁 Training completed - stopping HTTP polling');
|
||
return false; // Stop polling when training is done
|
||
}
|
||
|
||
console.log(`📊 HTTP fallback polling active (WebSocket disconnected) - ${HTTP_POLLING_INTERVAL_MS}ms interval`);
|
||
return HTTP_POLLING_INTERVAL_MS; // Poll while training (fallback when WebSocket unavailable)
|
||
} : false, // Completely disable interval when WebSocket connected
|
||
staleTime: 1000, // Consider data stale after 1 second
|
||
retry: (failureCount, error) => {
|
||
// Don't retry on auth errors
|
||
if ((error as any)?.status === 401) {
|
||
console.log('🚫 Not retrying due to auth error');
|
||
return false;
|
||
}
|
||
return failureCount < 3;
|
||
},
|
||
...queryOptions,
|
||
});
|
||
};
|
||
|
||
// Model Queries
|
||
export const useActiveModel = (
|
||
tenantId: string,
|
||
inventoryProductId: string,
|
||
options?: Omit<UseQueryOptions<any, ApiError>, 'queryKey' | 'queryFn'>
|
||
) => {
|
||
return useQuery<any, ApiError>({
|
||
queryKey: trainingKeys.models.active(tenantId, inventoryProductId),
|
||
queryFn: () => trainingService.getActiveModel(tenantId, inventoryProductId),
|
||
enabled: !!tenantId && !!inventoryProductId,
|
||
staleTime: 5 * 60 * 1000, // 5 minutes
|
||
...options,
|
||
});
|
||
};
|
||
|
||
export const useModels = (
|
||
tenantId: string,
|
||
queryParams?: any,
|
||
options?: Omit<UseQueryOptions<any, ApiError>, 'queryKey' | 'queryFn'>
|
||
) => {
|
||
return useQuery<any, ApiError>({
|
||
queryKey: trainingKeys.models.list(tenantId, queryParams),
|
||
queryFn: () => trainingService.getModels(tenantId, queryParams),
|
||
enabled: !!tenantId,
|
||
staleTime: 2 * 60 * 1000, // 2 minutes
|
||
...options,
|
||
});
|
||
};
|
||
|
||
export const useModelMetrics = (
|
||
tenantId: string,
|
||
modelId: string,
|
||
options?: Omit<UseQueryOptions<ModelMetricsResponse, ApiError>, 'queryKey' | 'queryFn'>
|
||
) => {
|
||
return useQuery<ModelMetricsResponse, ApiError>({
|
||
queryKey: trainingKeys.models.metrics(tenantId, modelId),
|
||
queryFn: () => trainingService.getModelMetrics(tenantId, modelId),
|
||
enabled: !!tenantId && !!modelId,
|
||
staleTime: 10 * 60 * 1000, // 10 minutes
|
||
...options,
|
||
});
|
||
};
|
||
|
||
export const useModelPerformance = (
|
||
tenantId: string,
|
||
modelId: string,
|
||
options?: Omit<UseQueryOptions<any, ApiError>, 'queryKey' | 'queryFn'>
|
||
) => {
|
||
return useQuery<any, ApiError>({
|
||
queryKey: trainingKeys.models.performance(tenantId, modelId),
|
||
queryFn: () => trainingService.getModelPerformance(tenantId, modelId),
|
||
enabled: !!tenantId && !!modelId,
|
||
staleTime: 10 * 60 * 1000, // 10 minutes
|
||
...options,
|
||
});
|
||
};
|
||
|
||
// Statistics Queries
|
||
export const useTenantTrainingStatistics = (
|
||
tenantId: string,
|
||
options?: Omit<UseQueryOptions<any, ApiError>, 'queryKey' | 'queryFn'>
|
||
) => {
|
||
return useQuery<any, ApiError>({
|
||
queryKey: trainingKeys.statistics(tenantId),
|
||
queryFn: () => trainingService.getTenantStatistics(tenantId),
|
||
enabled: !!tenantId,
|
||
staleTime: 5 * 60 * 1000, // 5 minutes
|
||
...options,
|
||
});
|
||
};
|
||
|
||
// Training Job Mutations
|
||
export const useCreateTrainingJob = (
|
||
options?: UseMutationOptions<
|
||
TrainingJobResponse,
|
||
ApiError,
|
||
{ tenantId: string; request: TrainingJobRequest }
|
||
>
|
||
) => {
|
||
const queryClient = useQueryClient();
|
||
|
||
return useMutation<
|
||
TrainingJobResponse,
|
||
ApiError,
|
||
{ tenantId: string; request: TrainingJobRequest }
|
||
>({
|
||
mutationFn: ({ tenantId, request }) => trainingService.createTrainingJob(tenantId, request),
|
||
onSuccess: (data, { tenantId }) => {
|
||
// Add the job status to cache
|
||
queryClient.setQueryData(
|
||
trainingKeys.jobs.status(tenantId, data.job_id),
|
||
{
|
||
job_id: data.job_id,
|
||
status: data.status,
|
||
progress: 0,
|
||
}
|
||
);
|
||
|
||
// Invalidate statistics to reflect the new training job
|
||
queryClient.invalidateQueries({ queryKey: trainingKeys.statistics(tenantId) });
|
||
},
|
||
...options,
|
||
});
|
||
};
|
||
|
||
export const useTrainSingleProduct = (
|
||
options?: UseMutationOptions<
|
||
TrainingJobResponse,
|
||
ApiError,
|
||
{ tenantId: string; inventoryProductId: string; request: SingleProductTrainingRequest }
|
||
>
|
||
) => {
|
||
const queryClient = useQueryClient();
|
||
|
||
return useMutation<
|
||
TrainingJobResponse,
|
||
ApiError,
|
||
{ tenantId: string; inventoryProductId: string; request: SingleProductTrainingRequest }
|
||
>({
|
||
mutationFn: ({ tenantId, inventoryProductId, request }) =>
|
||
trainingService.trainSingleProduct(tenantId, inventoryProductId, request),
|
||
onSuccess: (data, { tenantId, inventoryProductId }) => {
|
||
// Add the job status to cache
|
||
queryClient.setQueryData(
|
||
trainingKeys.jobs.status(tenantId, data.job_id),
|
||
{
|
||
job_id: data.job_id,
|
||
status: data.status,
|
||
progress: 0,
|
||
}
|
||
);
|
||
|
||
// Invalidate active model for this product
|
||
queryClient.invalidateQueries({
|
||
queryKey: trainingKeys.models.active(tenantId, inventoryProductId)
|
||
});
|
||
|
||
// Invalidate statistics
|
||
queryClient.invalidateQueries({ queryKey: trainingKeys.statistics(tenantId) });
|
||
},
|
||
...options,
|
||
});
|
||
};
|
||
|
||
// Admin Mutations
|
||
export const useDeleteAllTenantModels = (
|
||
options?: UseMutationOptions<{ message: string }, ApiError, { tenantId: string }>
|
||
) => {
|
||
const queryClient = useQueryClient();
|
||
|
||
return useMutation<{ message: string }, ApiError, { tenantId: string }>({
|
||
mutationFn: ({ tenantId }) => trainingService.deleteAllTenantModels(tenantId),
|
||
onSuccess: (_, { tenantId }) => {
|
||
// Invalidate all model-related queries for this tenant
|
||
queryClient.invalidateQueries({ queryKey: trainingKeys.models.all() });
|
||
queryClient.invalidateQueries({ queryKey: trainingKeys.statistics(tenantId) });
|
||
},
|
||
...options,
|
||
});
|
||
};
|
||
|
||
// WebSocket Hook for Real-time Training Updates
|
||
export const useTrainingWebSocket = (
|
||
tenantId: string,
|
||
jobId: string,
|
||
token?: string,
|
||
options?: {
|
||
onProgress?: (data: any) => void;
|
||
onCompleted?: (data: any) => void;
|
||
onError?: (error: any) => void;
|
||
onStarted?: (data: any) => void;
|
||
onCancelled?: (data: any) => void;
|
||
}
|
||
) => {
|
||
const queryClient = useQueryClient();
|
||
const [isConnected, setIsConnected] = React.useState(false);
|
||
const [connectionError, setConnectionError] = React.useState<string | null>(null);
|
||
const [connectionAttempts, setConnectionAttempts] = React.useState(0);
|
||
|
||
// Memoize options to prevent unnecessary effect re-runs
|
||
const memoizedOptions = React.useMemo(() => options, [
|
||
options?.onProgress,
|
||
options?.onCompleted,
|
||
options?.onError,
|
||
options?.onStarted,
|
||
options?.onCancelled
|
||
]);
|
||
|
||
React.useEffect(() => {
|
||
if (!tenantId || !jobId || !memoizedOptions) {
|
||
return;
|
||
}
|
||
|
||
let ws: WebSocket | null = null;
|
||
let reconnectTimer: NodeJS.Timeout | null = null;
|
||
let isManuallyDisconnected = false;
|
||
let reconnectAttempts = 0;
|
||
const maxReconnectAttempts = WEBSOCKET_MAX_RECONNECT_ATTEMPTS;
|
||
|
||
const connect = async () => {
|
||
try {
|
||
setConnectionError(null);
|
||
setConnectionAttempts(prev => prev + 1);
|
||
|
||
// Use centralized token management from apiClient
|
||
let effectiveToken: string | null;
|
||
|
||
try {
|
||
// Always use the apiClient's token management
|
||
effectiveToken = await apiClient.ensureValidToken();
|
||
|
||
if (!effectiveToken) {
|
||
throw new Error('No valid token available');
|
||
}
|
||
} catch (error) {
|
||
console.error('❌ Failed to get valid token for WebSocket:', error);
|
||
setConnectionError('Authentication failed. Please log in again.');
|
||
return;
|
||
}
|
||
|
||
console.log(`🔄 Attempting WebSocket connection (attempt ${reconnectAttempts + 1}/${maxReconnectAttempts + 1}):`, {
|
||
tenantId,
|
||
jobId,
|
||
hasToken: !!effectiveToken,
|
||
tokenFromApiClient: true
|
||
});
|
||
|
||
ws = trainingService.createWebSocketConnection(tenantId, jobId, effectiveToken);
|
||
|
||
ws.onopen = () => {
|
||
console.log('✅ Training WebSocket connected successfully', {
|
||
readyState: ws?.readyState,
|
||
url: ws?.url,
|
||
jobId
|
||
});
|
||
// Track connection time for debugging
|
||
(ws as any)._connectTime = Date.now();
|
||
setIsConnected(true);
|
||
reconnectAttempts = 0; // Reset on successful connection
|
||
|
||
// Request current status on connection
|
||
try {
|
||
ws?.send('get_status');
|
||
console.log('📤 Requested current training status');
|
||
} catch (e) {
|
||
console.warn('Failed to request status on connection:', e);
|
||
}
|
||
|
||
// Helper function to check if tokens represent different auth users/sessions
|
||
const isNewAuthSession = (oldToken: string, newToken: string): boolean => {
|
||
if (!oldToken || !newToken) return !!oldToken !== !!newToken;
|
||
|
||
try {
|
||
const oldPayload = JSON.parse(atob(oldToken.split('.')[1]));
|
||
const newPayload = JSON.parse(atob(newToken.split('.')[1]));
|
||
|
||
// Compare by user ID - different user means new auth session
|
||
// If user_id is same, it's just a token refresh, no need to reconnect
|
||
return oldPayload.user_id !== newPayload.user_id ||
|
||
oldPayload.sub !== newPayload.sub;
|
||
} catch (e) {
|
||
console.warn('Failed to parse token for session comparison:', e);
|
||
// On parse error, don't reconnect (assume same session)
|
||
return false;
|
||
}
|
||
};
|
||
|
||
// Set up periodic ping and check for auth session changes
|
||
const heartbeatInterval = setInterval(async () => {
|
||
if (ws?.readyState === WebSocket.OPEN && !isManuallyDisconnected) {
|
||
try {
|
||
// Check token validity (this may refresh if needed)
|
||
const currentToken = await apiClient.ensureValidToken();
|
||
|
||
// Only reconnect if user changed (new auth session)
|
||
if (currentToken && effectiveToken && isNewAuthSession(effectiveToken, currentToken)) {
|
||
console.log('🔄 Auth session changed (different user) - reconnecting WebSocket');
|
||
ws?.close(1000, 'Auth session changed - reconnecting');
|
||
clearInterval(heartbeatInterval);
|
||
return;
|
||
}
|
||
|
||
// Token may have been refreshed but it's the same user - continue
|
||
if (currentToken && currentToken !== effectiveToken) {
|
||
console.log('ℹ️ Token refreshed (same user) - updating reference');
|
||
effectiveToken = currentToken;
|
||
}
|
||
|
||
// Send ping
|
||
ws?.send('ping');
|
||
console.log('💓 Sent ping to server');
|
||
} catch (e) {
|
||
console.warn('Failed to send ping or validate token:', e);
|
||
clearInterval(heartbeatInterval);
|
||
}
|
||
} else {
|
||
clearInterval(heartbeatInterval);
|
||
}
|
||
}, WEBSOCKET_HEARTBEAT_INTERVAL_MS); // Check for auth changes and send ping
|
||
|
||
// Store interval for cleanup
|
||
(ws as any).heartbeatInterval = heartbeatInterval;
|
||
};
|
||
|
||
ws.onmessage = (event) => {
|
||
try {
|
||
// Handle non-JSON messages (like pong responses)
|
||
if (typeof event.data === 'string' && event.data === 'pong') {
|
||
console.log('🏓 Pong received from server');
|
||
return;
|
||
}
|
||
|
||
const message = JSON.parse(event.data);
|
||
|
||
console.log('🔔 Training WebSocket message received:', message);
|
||
|
||
// Handle initial state message to restore the latest known state
|
||
if (message.type === 'initial_state') {
|
||
console.log('📥 Received initial state:', message.data);
|
||
const initialData = message.data;
|
||
const initialEventData = initialData.data || {};
|
||
let initialProgress = initialEventData.progress || 0;
|
||
|
||
// Calculate progress for product_completed events
|
||
if (initialData.type === 'product_completed') {
|
||
const productsCompleted = initialEventData.products_completed || 0;
|
||
const totalProducts = initialEventData.total_products || 1;
|
||
const trainingRangeWidth = PROGRESS_TRAINING_RANGE_END - PROGRESS_DATA_ANALYSIS;
|
||
initialProgress = PROGRESS_DATA_ANALYSIS + Math.floor((productsCompleted / totalProducts) * trainingRangeWidth);
|
||
console.log('📦 Product training completed in initial state',
|
||
`${productsCompleted}/${totalProducts}`,
|
||
`progress: ${initialProgress}%`);
|
||
}
|
||
|
||
// Update job status in cache with initial state
|
||
queryClient.setQueryData(
|
||
trainingKeys.jobs.status(tenantId, jobId),
|
||
(oldData: TrainingJobStatus | undefined) => ({
|
||
...oldData,
|
||
job_id: jobId,
|
||
status: initialData.type === 'completed' ? 'completed' :
|
||
initialData.type === 'failed' ? 'failed' :
|
||
initialData.type === 'started' ? 'running' :
|
||
initialData.type === 'progress' ? 'running' :
|
||
initialData.type === 'product_completed' ? 'running' :
|
||
initialData.type === 'step_completed' ? 'running' :
|
||
oldData?.status || 'running',
|
||
progress: typeof initialProgress === 'number' ? initialProgress : oldData?.progress || 0,
|
||
current_step: initialEventData.current_step || initialEventData.step_name || oldData?.current_step,
|
||
})
|
||
);
|
||
return; // Initial state messages are only for state restoration, don't process as regular events
|
||
}
|
||
|
||
// Extract data from backend message structure
|
||
const eventData = message.data || {};
|
||
let progress = eventData.progress || 0;
|
||
const currentStep = eventData.current_step || eventData.step_name || '';
|
||
const stepDetails = eventData.step_details || '';
|
||
|
||
// Handle product_completed events - calculate progress dynamically
|
||
if (message.type === 'product_completed') {
|
||
const productsCompleted = eventData.products_completed || 0;
|
||
const totalProducts = eventData.total_products || 1;
|
||
|
||
// Calculate progress: DATA_ANALYSIS% base + (completed/total * (TRAINING_RANGE_END - DATA_ANALYSIS)%)
|
||
const trainingRangeWidth = PROGRESS_TRAINING_RANGE_END - PROGRESS_DATA_ANALYSIS;
|
||
progress = PROGRESS_DATA_ANALYSIS + Math.floor((productsCompleted / totalProducts) * trainingRangeWidth);
|
||
|
||
console.log('📦 Product training completed',
|
||
`${productsCompleted}/${totalProducts}`,
|
||
`progress: ${progress}%`);
|
||
}
|
||
|
||
// Update job status in cache
|
||
queryClient.setQueryData(
|
||
trainingKeys.jobs.status(tenantId, jobId),
|
||
(oldData: TrainingJobStatus | undefined) => ({
|
||
...oldData,
|
||
job_id: jobId,
|
||
status: message.type === 'completed' ? 'completed' :
|
||
message.type === 'failed' ? 'failed' :
|
||
message.type === 'started' ? 'running' :
|
||
oldData?.status || 'running',
|
||
progress: typeof progress === 'number' ? progress : oldData?.progress || 0,
|
||
current_step: currentStep || oldData?.current_step,
|
||
})
|
||
);
|
||
|
||
// Call appropriate callback based on message type
|
||
switch (message.type) {
|
||
case 'connected':
|
||
console.log('🔗 WebSocket connected');
|
||
break;
|
||
|
||
case 'started':
|
||
console.log('🚀 Training started');
|
||
memoizedOptions?.onStarted?.(message);
|
||
break;
|
||
|
||
case 'progress':
|
||
console.log('📊 Training progress update', `${progress}%`);
|
||
memoizedOptions?.onProgress?.(message);
|
||
break;
|
||
|
||
case 'product_completed':
|
||
console.log('✅ Product training completed');
|
||
// Treat as progress update
|
||
memoizedOptions?.onProgress?.({
|
||
...message,
|
||
data: {
|
||
...eventData,
|
||
progress, // Use calculated progress
|
||
}
|
||
});
|
||
break;
|
||
|
||
case 'step_completed':
|
||
console.log('📋 Step completed');
|
||
memoizedOptions?.onProgress?.(message);
|
||
break;
|
||
|
||
case 'completed':
|
||
console.log('✅ Training completed successfully');
|
||
memoizedOptions?.onCompleted?.(message);
|
||
// Invalidate models and statistics
|
||
queryClient.invalidateQueries({ queryKey: trainingKeys.models.all() });
|
||
queryClient.invalidateQueries({ queryKey: trainingKeys.statistics(tenantId) });
|
||
isManuallyDisconnected = true;
|
||
break;
|
||
|
||
case 'failed':
|
||
console.log('❌ Training failed');
|
||
memoizedOptions?.onError?.(message);
|
||
isManuallyDisconnected = true;
|
||
break;
|
||
|
||
default:
|
||
console.log(`🔍 Unknown message type: ${message.type}`);
|
||
break;
|
||
}
|
||
} catch (error) {
|
||
console.error('Error parsing WebSocket message:', error);
|
||
setConnectionError('Error parsing message from server');
|
||
}
|
||
};
|
||
|
||
ws.onerror = (error) => {
|
||
console.error('Training WebSocket error:', error);
|
||
setConnectionError('WebSocket connection error');
|
||
setIsConnected(false);
|
||
};
|
||
|
||
ws.onclose = (event) => {
|
||
console.log(`❌ Training WebSocket disconnected. Code: ${event.code}, Reason: "${event.reason}"`, {
|
||
wasClean: event.wasClean,
|
||
jobId,
|
||
timeConnected: ws ? `${Date.now() - (ws as any)._connectTime || 0}ms` : 'unknown',
|
||
reconnectAttempts
|
||
});
|
||
setIsConnected(false);
|
||
|
||
// Detailed logging for different close codes
|
||
switch (event.code) {
|
||
case 1000:
|
||
if (event.reason === 'Auth session changed - reconnecting') {
|
||
console.log('🔄 WebSocket closed for auth session change - will reconnect immediately');
|
||
} else {
|
||
console.log('🔒 WebSocket closed normally');
|
||
}
|
||
break;
|
||
case 1006:
|
||
console.log('⚠️ WebSocket closed abnormally (1006) - likely server-side issue or network problem');
|
||
break;
|
||
case 1001:
|
||
console.log('🔄 WebSocket endpoint going away');
|
||
break;
|
||
case 1003:
|
||
console.log('❌ WebSocket unsupported data received');
|
||
break;
|
||
default:
|
||
console.log(`❓ WebSocket closed with code ${event.code}`);
|
||
}
|
||
|
||
// Handle auth session change reconnection (immediate reconnect)
|
||
if (event.code === 1000 && event.reason === 'Auth session changed - reconnecting') {
|
||
console.log('🔄 Reconnecting immediately due to auth session change...');
|
||
reconnectTimer = setTimeout(() => {
|
||
connect(); // Reconnect immediately with new session token
|
||
}, WEBSOCKET_RECONNECT_INITIAL_DELAY_MS); // Short delay to allow cleanup
|
||
return;
|
||
}
|
||
|
||
// Try to reconnect if not manually disconnected and haven't exceeded max attempts
|
||
if (!isManuallyDisconnected && event.code !== 1000 && reconnectAttempts < maxReconnectAttempts) {
|
||
const delay = Math.min(
|
||
WEBSOCKET_RECONNECT_INITIAL_DELAY_MS * Math.pow(2, reconnectAttempts),
|
||
WEBSOCKET_RECONNECT_MAX_DELAY_MS
|
||
); // Exponential backoff
|
||
console.log(`🔄 Attempting to reconnect WebSocket in ${delay/1000}s... (attempt ${reconnectAttempts + 1}/${maxReconnectAttempts})`);
|
||
|
||
reconnectTimer = setTimeout(() => {
|
||
reconnectAttempts++;
|
||
connect();
|
||
}, delay);
|
||
} else if (reconnectAttempts >= maxReconnectAttempts) {
|
||
console.log(`❌ Max reconnection attempts (${maxReconnectAttempts}) reached. Giving up.`);
|
||
setConnectionError(`Connection failed after ${maxReconnectAttempts} attempts. The training job may not exist or the server may be unavailable.`);
|
||
}
|
||
};
|
||
|
||
} catch (error) {
|
||
console.error('Error creating WebSocket connection:', error);
|
||
setConnectionError('Failed to create WebSocket connection');
|
||
}
|
||
};
|
||
|
||
// Connect immediately to avoid missing early progress updates
|
||
console.log('🚀 Starting immediate WebSocket connection...');
|
||
connect();
|
||
|
||
// Cleanup function
|
||
return () => {
|
||
isManuallyDisconnected = true;
|
||
|
||
if (reconnectTimer) {
|
||
clearTimeout(reconnectTimer);
|
||
}
|
||
|
||
if (ws) {
|
||
ws.close(1000, 'Component unmounted');
|
||
}
|
||
|
||
setIsConnected(false);
|
||
};
|
||
}, [tenantId, jobId, queryClient, memoizedOptions]);
|
||
|
||
return {
|
||
isConnected,
|
||
connectionError
|
||
};
|
||
};
|
||
|
||
// Utility Hooks
|
||
export const useIsTrainingInProgress = (
|
||
tenantId: string,
|
||
jobId?: string,
|
||
isWebSocketConnected?: boolean
|
||
) => {
|
||
const { data: jobStatus } = useTrainingJobStatus(tenantId, jobId || '', {
|
||
enabled: !!jobId,
|
||
isWebSocketConnected,
|
||
});
|
||
|
||
return jobStatus?.status === 'running' || jobStatus?.status === 'pending';
|
||
};
|
||
|
||
export const useTrainingProgress = (
|
||
tenantId: string,
|
||
jobId?: string,
|
||
isWebSocketConnected?: boolean
|
||
) => {
|
||
const { data: jobStatus } = useTrainingJobStatus(tenantId, jobId || '', {
|
||
enabled: !!jobId,
|
||
isWebSocketConnected,
|
||
});
|
||
|
||
return {
|
||
progress: jobStatus?.progress || 0,
|
||
currentStep: jobStatus?.current_step,
|
||
isComplete: jobStatus?.status === 'completed',
|
||
isFailed: jobStatus?.status === 'failed',
|
||
isRunning: jobStatus?.status === 'running',
|
||
};
|
||
};
|