import { CheckCircle2, Loader2, XCircle } from 'lucide-react'; import { useCallback, useEffect, useRef } from 'react'; import { Progress } from '@/components/ui/progress'; import { useToast } from '@/components/ui/use-toast'; import type { ModelProgress } from '@/lib/api/types'; import { useServerStore } from '@/stores/serverStore'; interface UseModelDownloadToastOptions { modelName: string; displayName: string; enabled?: boolean; onComplete?: () => void; onError?: (error: string) => void; } /** * Hook to show and update a toast notification with model download progress. * Subscribes to Server-Sent Events for real-time progress updates. */ export function useModelDownloadToast({ modelName, displayName, enabled = false, onComplete, onError, }: UseModelDownloadToastOptions) { const { toast } = useToast(); const serverUrl = useServerStore((state) => state.serverUrl); const toastIdRef = useRef(null); // biome-ignore lint: Using any for toast update ref to handle complex toast types const toastUpdateRef = useRef(null); const eventSourceRef = useRef(null); const formatBytes = useCallback((bytes: number): string => { if (bytes === 0) return '0 B'; const k = 1024; const sizes = ['B', 'KB', 'MB', 'GB']; const i = Math.floor(Math.log(bytes) / Math.log(k)); return `${(bytes / k ** i).toFixed(1)} ${sizes[i]}`; }, []); useEffect(() => { console.log('[useModelDownloadToast] useEffect triggered', { enabled, serverUrl, modelName, displayName, }); if (!enabled || !serverUrl || !modelName) { console.log('[useModelDownloadToast] Not enabled, skipping'); return; } console.log('[useModelDownloadToast] Creating toast and EventSource for:', modelName); // Create initial toast const toastResult = toast({ title: displayName, description: (
Connecting to download...
), duration: Infinity, // Don't auto-dismiss, we'll handle it manually }); toastIdRef.current = toastResult.id; toastUpdateRef.current = toastResult.update; // Subscribe to progress updates via Server-Sent Events const eventSourceUrl = `${serverUrl}/models/progress/${modelName}`; console.log('[useModelDownloadToast] Creating EventSource to:', eventSourceUrl); const eventSource = new EventSource(eventSourceUrl); eventSource.onopen = () => { console.log('[useModelDownloadToast] EventSource connection opened for:', modelName); }; eventSource.onmessage = (event) => { console.log('[useModelDownloadToast] Received SSE message:', event.data); try { const progress = JSON.parse(event.data) as ModelProgress; // Update toast with progress if (toastIdRef.current && toastUpdateRef.current) { const progressPercent = progress.total > 0 ? progress.progress : 0; const progressText = progress.total > 0 ? `${formatBytes(progress.current)} / ${formatBytes(progress.total)} (${progress.progress.toFixed(1)}%)` : ''; // Determine status icon and text let statusIcon: React.ReactNode = null; let statusText = 'Processing...'; switch (progress.status) { case 'complete': statusIcon = ; statusText = 'Download complete'; break; case 'error': statusIcon = ; statusText = 'Download failed. See Problems panel for details.'; break; case 'downloading': statusIcon = ; statusText = progress.filename || 'Downloading...'; break; case 'extracting': statusIcon = ; statusText = 'Extracting...'; break; } toastUpdateRef.current({ title: (
{statusIcon} {displayName}
), description: (
{statusText}
{progress.total > 0 && ( <>
{progressText}
)}
), duration: progress.status === 'complete' || progress.status === 'error' ? 5000 : Infinity, }); // Close connection and dismiss toast on completion or error // Also treat progress >= 100% as complete const isComplete = progress.status === 'complete' || progress.progress >= 100; const isError = progress.status === 'error'; if (isComplete || isError) { console.log('[useModelDownloadToast] Download finished:', { isComplete, isError, progress: progress.progress, }); eventSource.close(); eventSourceRef.current = null; // Update toast to show completion state before callbacks if (isComplete && toastUpdateRef.current) { toastUpdateRef.current({ title: (
{displayName}
), description: 'Download complete', duration: 3000, }); } // Call callbacks if (isComplete && onComplete) { console.log('[useModelDownloadToast] Download complete, calling onComplete callback'); onComplete(); } else if (isError && onError) { console.log('[useModelDownloadToast] Download error, calling onError callback'); onError(progress.error || 'Unknown error'); } } } } catch (error) { console.error('Error parsing progress event:', error); } }; eventSource.onerror = (error) => { console.error('[useModelDownloadToast] SSE error for:', modelName, error); console.log('[useModelDownloadToast] EventSource readyState:', eventSource.readyState); eventSource.close(); eventSourceRef.current = null; // Show error toast if (toastIdRef.current && toastUpdateRef.current) { toastUpdateRef.current({ title: displayName, description: 'Failed to track download progress', variant: 'destructive', duration: 5000, }); toastIdRef.current = null; toastUpdateRef.current = null; } }; eventSourceRef.current = eventSource; // Cleanup on unmount or when disabled return () => { console.log('[useModelDownloadToast] Cleanup - closing EventSource for:', modelName); if (eventSourceRef.current) { eventSourceRef.current.close(); eventSourceRef.current = null; } // Note: We don't dismiss the toast here as it might still be showing completion state }; }, [enabled, serverUrl, modelName, displayName, toast, formatBytes, onComplete, onError]); return { isTracking: enabled && eventSourceRef.current !== null, }; }