Refactor model download handling and improve progress tracking

- Rearranged imports for consistency across components.
- Enhanced the ModelManagement component to include detailed logging for download actions and errors.
- Updated the ModelProgress component to connect to SSE only when actively downloading, preventing connection exhaustion.
- Added a downloading state to the model status to indicate ongoing downloads.
- Improved toast notifications for model downloads with completion and error callbacks.
- Refactored the useModelDownloadToast hook to support new callbacks for download completion and error handling.
- Updated backend model status to reflect downloading state during active downloads.
This commit is contained in:
Jamie Pine
2026-01-30 19:53:20 -08:00
parent 77418a52ae
commit 07c0aba883
13 changed files with 223 additions and 102 deletions
+12 -6
View File
@@ -1,14 +1,14 @@
import { useEffect, useRef, useState } from 'react';
import { RouterProvider } from '@tanstack/react-router'; import { RouterProvider } from '@tanstack/react-router';
import { useEffect, useRef, useState } from 'react';
import voiceboxLogo from '@/assets/voicebox-logo.png'; import voiceboxLogo from '@/assets/voicebox-logo.png';
import ShinyText from '@/components/ShinyText'; import ShinyText from '@/components/ShinyText';
import { TitleBarDragRegion } from '@/components/TitleBarDragRegion'; import { TitleBarDragRegion } from '@/components/TitleBarDragRegion';
import { useAutoUpdater } from '@/hooks/useAutoUpdater';
import { TOP_SAFE_AREA_PADDING } from '@/lib/constants/ui'; import { TOP_SAFE_AREA_PADDING } from '@/lib/constants/ui';
import { cn } from '@/lib/utils/cn'; import { cn } from '@/lib/utils/cn';
import { usePlatform } from '@/platform/PlatformContext';
import { router } from '@/router'; import { router } from '@/router';
import { useServerStore } from '@/stores/serverStore'; import { useServerStore } from '@/stores/serverStore';
import { usePlatform } from '@/platform/PlatformContext';
import { useAutoUpdater } from '@/hooks/useAutoUpdater';
const LOADING_MESSAGES = [ const LOADING_MESSAGES = [
'Warming up tensors...', 'Warming up tensors...',
@@ -50,14 +50,18 @@ function App() {
console.error('Failed to sync initial setting to Rust:', error); console.error('Failed to sync initial setting to Rust:', error);
}); });
} }
}, [platform]); // Empty dependency array - platform is stable from context, only run once
// eslint-disable-next-line react-hooks/exhaustive-deps
}, [platform.metadata.isTauri, platform.lifecycle]);
// Setup lifecycle callbacks // Setup lifecycle callbacks
useEffect(() => { useEffect(() => {
platform.lifecycle.onServerReady = () => { platform.lifecycle.onServerReady = () => {
setServerReady(true); setServerReady(true);
}; };
}, [platform]); // Empty dependency array - platform is stable from context, only run once
// eslint-disable-next-line react-hooks/exhaustive-deps
}, [platform.lifecycle]);
// Setup window close handler and auto-start server when running in Tauri (production only) // Setup window close handler and auto-start server when running in Tauri (production only)
useEffect(() => { useEffect(() => {
@@ -115,7 +119,9 @@ function App() {
// Window close event handles server shutdown based on setting // Window close event handles server shutdown based on setting
serverStartingRef.current = false; serverStartingRef.current = false;
}; };
}, [platform]); // Empty dependency array - platform is stable from context, only run once
// eslint-disable-next-line react-hooks/exhaustive-deps
}, [platform.metadata.isTauri, platform.lifecycle]);
// Cycle through loading messages every 3 seconds // Cycle through loading messages every 3 seconds
useEffect(() => { useEffect(() => {
+16 -3
View File
@@ -1,6 +1,13 @@
import { AudioWaveform, Download, FileArchive, Loader2, MoreHorizontal, Play, Trash2 } from 'lucide-react'; import {
AudioWaveform,
Download,
FileArchive,
Loader2,
MoreHorizontal,
Play,
Trash2,
} from 'lucide-react';
import { useEffect, useRef, useState } from 'react'; import { useEffect, useRef, useState } from 'react';
import type { HistoryResponse } from '@/lib/api/types';
import { Button } from '@/components/ui/button'; import { Button } from '@/components/ui/button';
import { import {
Dialog, Dialog,
@@ -19,6 +26,7 @@ import {
import { Textarea } from '@/components/ui/textarea'; import { Textarea } from '@/components/ui/textarea';
import { useToast } from '@/components/ui/use-toast'; import { useToast } from '@/components/ui/use-toast';
import { apiClient } from '@/lib/api/client'; import { apiClient } from '@/lib/api/client';
import type { HistoryResponse } from '@/lib/api/types';
import { BOTTOM_SAFE_AREA_PADDING } from '@/lib/constants/ui'; import { BOTTOM_SAFE_AREA_PADDING } from '@/lib/constants/ui';
import { import {
useDeleteGeneration, useDeleteGeneration,
@@ -48,7 +56,11 @@ export function HistoryTable() {
const limit = 20; const limit = 20;
const { toast } = useToast(); const { toast } = useToast();
const { data: historyData, isLoading, isFetching } = useHistory({ const {
data: historyData,
isLoading,
isFetching,
} = useHistory({
limit, limit,
offset: page * limit, offset: page * limit,
}); });
@@ -265,6 +277,7 @@ export function HistoryTable() {
<Textarea <Textarea
value={gen.text} value={gen.text}
className="flex-1 resize-none text-sm text-muted-foreground select-text" className="flex-1 resize-none text-sm text-muted-foreground select-text"
readOnly
/> />
</div> </div>
@@ -1,6 +1,6 @@
import { useMutation, useQuery, useQueryClient } from '@tanstack/react-query'; import { useMutation, useQuery, useQueryClient } from '@tanstack/react-query';
import { Download, Loader2, Trash2 } from 'lucide-react'; import { Download, Loader2, Trash2 } from 'lucide-react';
import { useState } from 'react'; import { useCallback, useState } from 'react';
import { import {
AlertDialog, AlertDialog,
AlertDialogAction, AlertDialogAction,
@@ -17,7 +17,6 @@ import { Card, CardContent, CardDescription, CardHeader, CardTitle } from '@/com
import { useToast } from '@/components/ui/use-toast'; import { useToast } from '@/components/ui/use-toast';
import { apiClient } from '@/lib/api/client'; import { apiClient } from '@/lib/api/client';
import { useModelDownloadToast } from '@/lib/hooks/useModelDownloadToast'; import { useModelDownloadToast } from '@/lib/hooks/useModelDownloadToast';
import { ModelProgress } from './ModelProgress';
export function ModelManagement() { export function ModelManagement() {
const { toast } = useToast(); const { toast } = useToast();
@@ -27,15 +26,36 @@ export function ModelManagement() {
const { data: modelStatus, isLoading } = useQuery({ const { data: modelStatus, isLoading } = useQuery({
queryKey: ['modelStatus'], queryKey: ['modelStatus'],
queryFn: () => apiClient.getModelStatus(), queryFn: async () => {
console.log('[Query] Fetching model status');
const result = await apiClient.getModelStatus();
console.log('[Query] Model status fetched:', result);
return result;
},
refetchInterval: 5000, // Refresh every 5 seconds refetchInterval: 5000, // Refresh every 5 seconds
}); });
// Callbacks for download completion
const handleDownloadComplete = useCallback(() => {
console.log('[ModelManagement] Download complete, clearing state');
setDownloadingModel(null);
setDownloadingDisplayName(null);
queryClient.invalidateQueries({ queryKey: ['modelStatus'] });
}, [queryClient]);
const handleDownloadError = useCallback(() => {
console.log('[ModelManagement] Download error, clearing state');
setDownloadingModel(null);
setDownloadingDisplayName(null);
}, []);
// Use progress toast hook for the downloading model // Use progress toast hook for the downloading model
useModelDownloadToast({ useModelDownloadToast({
modelName: downloadingModel || '', modelName: downloadingModel || '',
displayName: downloadingDisplayName || '', displayName: downloadingDisplayName || '',
enabled: !!downloadingModel && !!downloadingDisplayName, enabled: !!downloadingModel && !!downloadingDisplayName,
onComplete: handleDownloadComplete,
onError: handleDownloadError,
}); });
const [deleteDialogOpen, setDeleteDialogOpen] = useState(false); const [deleteDialogOpen, setDeleteDialogOpen] = useState(false);
@@ -45,44 +65,69 @@ export function ModelManagement() {
sizeMb?: number; sizeMb?: number;
} | null>(null); } | null>(null);
const downloadMutation = useMutation({ const handleDownload = async (modelName: string) => {
mutationFn: (modelName: string) => { console.log('[Download] Button clicked for:', modelName, 'at', new Date().toISOString());
setDownloadingModel(modelName);
// Find display name from model status // Find display name
const model = modelStatus?.models.find((m) => m.model_name === modelName); const model = modelStatus?.models.find((m) => m.model_name === modelName);
setDownloadingDisplayName(model?.display_name || modelName); const displayName = model?.display_name || modelName;
return apiClient.triggerModelDownload(modelName);
}, try {
onSuccess: () => { // IMPORTANT: Call the API FIRST before setting state
// Download completed - clear state and refetch status // Setting state enables the SSE EventSource in useModelDownloadToast,
setDownloadingModel(null); // which can block/delay the download fetch due to HTTP/1.1 connection limits
setDownloadingDisplayName(null); console.log('[Download] Calling download API for:', modelName);
const result = await apiClient.triggerModelDownload(modelName);
console.log('[Download] Download API responded:', result);
// NOW set state to enable SSE tracking (after download has started on backend)
setDownloadingModel(modelName);
setDownloadingDisplayName(displayName);
// Download initiated successfully - state will be cleared when SSE reports completion
// or by the polling interval detecting the model is downloaded
queryClient.invalidateQueries({ queryKey: ['modelStatus'] }); queryClient.invalidateQueries({ queryKey: ['modelStatus'] });
}, } catch (error) {
onError: (error: Error) => { console.error('[Download] Download failed:', error);
setDownloadingModel(null); setDownloadingModel(null);
setDownloadingDisplayName(null); setDownloadingDisplayName(null);
toast({ toast({
title: 'Download failed', title: 'Download failed',
description: error.message, description: error instanceof Error ? error.message : 'Unknown error',
variant: 'destructive', variant: 'destructive',
}); });
}, }
}); };
const deleteMutation = useMutation({ const deleteMutation = useMutation({
mutationFn: (modelName: string) => apiClient.deleteModel(modelName), mutationFn: async (modelName: string) => {
onSuccess: () => { console.log('[Delete] Deleting model:', modelName);
const result = await apiClient.deleteModel(modelName);
console.log('[Delete] Model deleted successfully:', modelName);
return result;
},
onSuccess: async (_data, _modelName) => {
console.log('[Delete] onSuccess - showing toast and invalidating queries');
toast({ toast({
title: 'Model deleted', title: 'Model deleted',
description: `${modelToDelete?.displayName || 'Model'} has been deleted successfully.`, description: `${modelToDelete?.displayName || 'Model'} has been deleted successfully.`,
}); });
setDeleteDialogOpen(false); setDeleteDialogOpen(false);
setModelToDelete(null); setModelToDelete(null);
// Refetch status to update UI // Invalidate AND explicitly refetch to ensure UI updates
queryClient.invalidateQueries({ queryKey: ['modelStatus'] }); // Using refetchType: 'all' ensures we refetch even if the query is stale
console.log('[Delete] Invalidating modelStatus query');
await queryClient.invalidateQueries({
queryKey: ['modelStatus'],
refetchType: 'all',
});
// Also explicitly refetch to guarantee fresh data
console.log('[Delete] Explicitly refetching modelStatus query');
await queryClient.refetchQueries({ queryKey: ['modelStatus'] });
console.log('[Delete] Query refetched');
}, },
onError: (error: Error) => { onError: (error: Error) => {
console.log('[Delete] onError:', error);
toast({ toast({
title: 'Delete failed', title: 'Delete failed',
description: error.message, description: error.message,
@@ -124,7 +169,7 @@ export function ModelManagement() {
<ModelItem <ModelItem
key={model.model_name} key={model.model_name}
model={model} model={model}
onDownload={() => downloadMutation.mutate(model.model_name)} onDownload={() => handleDownload(model.model_name)}
onDelete={() => { onDelete={() => {
setModelToDelete({ setModelToDelete({
name: model.model_name, name: model.model_name,
@@ -152,7 +197,7 @@ export function ModelManagement() {
<ModelItem <ModelItem
key={model.model_name} key={model.model_name}
model={model} model={model}
onDownload={() => downloadMutation.mutate(model.model_name)} onDownload={() => handleDownload(model.model_name)}
onDelete={() => { onDelete={() => {
setModelToDelete({ setModelToDelete({
name: model.model_name, name: model.model_name,
@@ -168,21 +213,6 @@ export function ModelManagement() {
</div> </div>
</div> </div>
{/* Progress indicators */}
<div className="pt-4 border-t">
<h3 className="text-sm font-semibold mb-3 text-muted-foreground">
Download Progress
</h3>
<div className="space-y-2">
{modelStatus.models.map((model) => (
<ModelProgress
key={model.model_name}
modelName={model.model_name}
displayName={model.display_name}
/>
))}
</div>
</div>
</div> </div>
) : null} ) : null}
</CardContent> </CardContent>
@@ -235,16 +265,20 @@ interface ModelItemProps {
model_name: string; model_name: string;
display_name: string; display_name: string;
downloaded: boolean; downloaded: boolean;
downloading?: boolean; // From server - true if download in progress
size_mb?: number; size_mb?: number;
loaded: boolean; loaded: boolean;
}; };
onDownload: () => void; onDownload: () => void;
onDelete: () => void; onDelete: () => void;
isDownloading: boolean; isDownloading: boolean; // Local state - true if user just clicked download
formatSize: (sizeMb?: number) => string; formatSize: (sizeMb?: number) => string;
} }
function ModelItem({ model, onDownload, onDelete, isDownloading, formatSize }: ModelItemProps) { function ModelItem({ model, onDownload, onDelete, isDownloading, formatSize }: ModelItemProps) {
// Use server's downloading state OR local state (for immediate feedback before server updates)
const showDownloading = model.downloading || isDownloading;
return ( return (
<div className="flex items-center justify-between p-3 border rounded-lg"> <div className="flex items-center justify-between p-3 border rounded-lg">
<div className="flex-1"> <div className="flex-1">
@@ -255,20 +289,21 @@ function ModelItem({ model, onDownload, onDelete, isDownloading, formatSize }: M
Loaded Loaded
</Badge> </Badge>
)} )}
{model.downloaded && !model.loaded && ( {/* Only show Downloaded if actually downloaded AND not downloading */}
{model.downloaded && !model.loaded && !showDownloading && (
<Badge variant="secondary" className="text-xs"> <Badge variant="secondary" className="text-xs">
Downloaded Downloaded
</Badge> </Badge>
)} )}
</div> </div>
{model.downloaded && model.size_mb && ( {model.downloaded && model.size_mb && !showDownloading && (
<div className="text-xs text-muted-foreground mt-1"> <div className="text-xs text-muted-foreground mt-1">
Size: {formatSize(model.size_mb)} Size: {formatSize(model.size_mb)}
</div> </div>
)} )}
</div> </div>
<div className="flex items-center gap-2"> <div className="flex items-center gap-2">
{model.downloaded ? ( {model.downloaded && !showDownloading ? (
<div className="flex items-center gap-2"> <div className="flex items-center gap-2">
<div className="flex items-center gap-1 text-sm text-muted-foreground"> <div className="flex items-center gap-1 text-sm text-muted-foreground">
<span>Ready</span> <span>Ready</span>
@@ -283,19 +318,15 @@ function ModelItem({ model, onDownload, onDelete, isDownloading, formatSize }: M
<Trash2 className="h-4 w-4" /> <Trash2 className="h-4 w-4" />
</Button> </Button>
</div> </div>
) : ( ) : showDownloading ? (
<Button size="sm" onClick={onDownload} disabled={isDownloading} variant="outline"> <Button size="sm" variant="outline" disabled>
{isDownloading ? (
<>
<Loader2 className="h-4 w-4 mr-2 animate-spin" /> <Loader2 className="h-4 w-4 mr-2 animate-spin" />
Downloading... Downloading...
</> </Button>
) : ( ) : (
<> <Button size="sm" onClick={onDownload} variant="outline">
<Download className="h-4 w-4 mr-2" /> <Download className="h-4 w-4 mr-2" />
Download Download
</>
)}
</Button> </Button>
)} )}
</div> </div>
@@ -8,14 +8,23 @@ import { useServerStore } from '@/stores/serverStore';
interface ModelProgressProps { interface ModelProgressProps {
modelName: string; modelName: string;
displayName: string; displayName: string;
/** Only connect to SSE when actively downloading - prevents connection exhaustion */
isDownloading?: boolean;
} }
export function ModelProgress({ modelName, displayName }: ModelProgressProps) { export function ModelProgress({ modelName, displayName, isDownloading = false }: ModelProgressProps) {
const [progress, setProgress] = useState<ModelProgressType | null>(null); const [progress, setProgress] = useState<ModelProgressType | null>(null);
const serverUrl = useServerStore((state) => state.serverUrl); const serverUrl = useServerStore((state) => state.serverUrl);
useEffect(() => { useEffect(() => {
if (!serverUrl) return; // IMPORTANT: Only connect to SSE when this specific model is downloading
// Opening SSE connections for all models exhausts HTTP/1.1 connection limits (6 per origin)
// which causes other fetches (like the download trigger) to be queued/blocked
if (!serverUrl || !isDownloading) {
return;
}
console.log(`[ModelProgress] Connecting SSE for ${modelName}`);
// Subscribe to progress updates via Server-Sent Events // Subscribe to progress updates via Server-Sent Events
const eventSource = new EventSource(`${serverUrl}/models/progress/${modelName}`); const eventSource = new EventSource(`${serverUrl}/models/progress/${modelName}`);
@@ -27,6 +36,7 @@ export function ModelProgress({ modelName, displayName }: ModelProgressProps) {
// Close connection if complete or error // Close connection if complete or error
if (data.status === 'complete' || data.status === 'error') { if (data.status === 'complete' || data.status === 'error') {
console.log(`[ModelProgress] Download ${data.status} for ${modelName}, closing SSE`);
eventSource.close(); eventSource.close();
} }
} catch (error) { } catch (error) {
@@ -35,14 +45,15 @@ export function ModelProgress({ modelName, displayName }: ModelProgressProps) {
}; };
eventSource.onerror = (error) => { eventSource.onerror = (error) => {
console.error('SSE error:', error); console.error(`[ModelProgress] SSE error for ${modelName}:`, error);
eventSource.close(); eventSource.close();
}; };
return () => { return () => {
console.log(`[ModelProgress] Cleanup - closing SSE for ${modelName}`);
eventSource.close(); eventSource.close();
}; };
}, [serverUrl, modelName]); }, [serverUrl, modelName, isDownloading]);
// Don't render if no progress or if complete/error and some time has passed // Don't render if no progress or if complete/error and some time has passed
if ( if (
+16 -10
View File
@@ -1,4 +1,4 @@
import { useCallback, useEffect, useState } from 'react'; import { useCallback, useEffect, useRef, useState } from 'react';
import { usePlatform } from '@/platform/PlatformContext'; import { usePlatform } from '@/platform/PlatformContext';
import type { UpdateStatus } from '@/platform/types'; import type { UpdateStatus } from '@/platform/types';
@@ -7,9 +7,8 @@ export type { UpdateStatus };
export function useAutoUpdater(checkOnMount = false) { export function useAutoUpdater(checkOnMount = false) {
const platform = usePlatform(); const platform = usePlatform();
const [status, setStatus] = useState<UpdateStatus>( const [status, setStatus] = useState<UpdateStatus>(platform.updater.getStatus());
platform.updater.getStatus(), const hasCheckedRef = useRef(false);
);
// Subscribe to updater status changes // Subscribe to updater status changes
useEffect(() => { useEffect(() => {
@@ -17,25 +16,32 @@ export function useAutoUpdater(checkOnMount = false) {
setStatus(newStatus); setStatus(newStatus);
}); });
return unsubscribe; return unsubscribe;
}, [platform]); // Empty dependency array - platform is stable from context
// eslint-disable-next-line react-hooks/exhaustive-deps
}, [platform.updater.subscribe]);
const checkForUpdates = useCallback(async () => { const checkForUpdates = useCallback(async () => {
await platform.updater.checkForUpdates(); await platform.updater.checkForUpdates();
}, [platform]); // eslint-disable-next-line react-hooks/exhaustive-deps
}, [platform.updater.checkForUpdates]);
const downloadAndInstall = useCallback(async () => { const downloadAndInstall = useCallback(async () => {
await platform.updater.downloadAndInstall(); await platform.updater.downloadAndInstall();
}, [platform]); // eslint-disable-next-line react-hooks/exhaustive-deps
}, [platform.updater.downloadAndInstall]);
const restartAndInstall = useCallback(async () => { const restartAndInstall = useCallback(async () => {
await platform.updater.restartAndInstall(); await platform.updater.restartAndInstall();
}, [platform]); // eslint-disable-next-line react-hooks/exhaustive-deps
}, [platform.updater.restartAndInstall]);
useEffect(() => { useEffect(() => {
if (checkOnMount && platform.metadata.isTauri) { if (checkOnMount && platform.metadata.isTauri && !hasCheckedRef.current) {
hasCheckedRef.current = true;
checkForUpdates(); checkForUpdates();
} }
}, [checkOnMount, checkForUpdates, platform.metadata.isTauri]); // eslint-disable-next-line react-hooks/exhaustive-deps
}, [platform.metadata.isTauricheckOnMountcheckForUpdates]);
return { return {
status, status,
+12 -5
View File
@@ -44,19 +44,24 @@ export function useAutoUpdater(options: boolean | UseAutoUpdaterOptions = false)
setStatus(newStatus); setStatus(newStatus);
}); });
return unsubscribe; return unsubscribe;
}, [platform]); // Empty dependency array - platform is stable from context
// eslint-disable-next-line react-hooks/exhaustive-deps
}, [platform.updater.subscribe]);
const checkForUpdates = useCallback(async () => { const checkForUpdates = useCallback(async () => {
await platform.updater.checkForUpdates(); await platform.updater.checkForUpdates();
}, [platform]); // eslint-disable-next-line react-hooks/exhaustive-deps
}, [platform.updater.checkForUpdates]);
const downloadAndInstall = useCallback(async () => { const downloadAndInstall = useCallback(async () => {
await platform.updater.downloadAndInstall(); await platform.updater.downloadAndInstall();
}, [platform]); // eslint-disable-next-line react-hooks/exhaustive-deps
}, [platform.updater.downloadAndInstall]);
const restartAndInstall = useCallback(async () => { const restartAndInstall = useCallback(async () => {
await platform.updater.restartAndInstall(); await platform.updater.restartAndInstall();
}, [platform]); // eslint-disable-next-line react-hooks/exhaustive-deps
}, [platform.updater.restartAndInstall]);
// Check for updates on mount // Check for updates on mount
useEffect(() => { useEffect(() => {
@@ -66,7 +71,9 @@ export function useAutoUpdater(options: boolean | UseAutoUpdaterOptions = false)
console.error('Auto update check failed:', error); console.error('Auto update check failed:', error);
}); });
} }
}, [checkOnMount, checkForUpdates, platform.metadata.isTauri]); // Empty dependency array - only run once on mount
// eslint-disable-next-line react-hooks/exhaustive-deps
}, [platform.metadata.isTauricheckOnMountcheckForUpdates]);
// Show toast when update is available // Show toast when update is available
useEffect(() => { useEffect(() => {
+4 -1
View File
@@ -310,10 +310,13 @@ class ApiClient {
} }
async triggerModelDownload(modelName: string): Promise<{ message: string }> { async triggerModelDownload(modelName: string): Promise<{ message: string }> {
return this.request<{ message: string }>('/models/download', { console.log('[API] triggerModelDownload called for:', modelName, 'at', new Date().toISOString());
const result = await this.request<{ message: string }>('/models/download', {
method: 'POST', method: 'POST',
body: JSON.stringify({ model_name: modelName } as ModelDownloadRequest), body: JSON.stringify({ model_name: modelName } as ModelDownloadRequest),
}); });
console.log('[API] triggerModelDownload response:', result);
return result;
} }
async deleteModel(modelName: string): Promise<{ message: string }> { async deleteModel(modelName: string): Promise<{ message: string }> {
+1
View File
@@ -9,6 +9,7 @@ export type ModelStatus = {
model_name: string; model_name: string;
display_name: string; display_name: string;
downloaded: boolean; downloaded: boolean;
downloading?: boolean; // True if download is in progress
size_mb?: number | null; size_mb?: number | null;
loaded?: boolean; loaded?: boolean;
}; };
+1
View File
@@ -96,6 +96,7 @@ export interface ModelStatus {
model_name: string; model_name: string;
display_name: string; display_name: string;
downloaded: boolean; downloaded: boolean;
downloading: boolean; // True if download is in progress
size_mb?: number; size_mb?: number;
loaded: boolean; loaded: boolean;
} }
+42 -17
View File
@@ -1,4 +1,4 @@
import { useEffect, useRef } from 'react'; import { useCallback, useEffect, useRef } from 'react';
import { useToast } from '@/components/ui/use-toast'; import { useToast } from '@/components/ui/use-toast';
import { useServerStore } from '@/stores/serverStore'; import { useServerStore } from '@/stores/serverStore';
import { Progress } from '@/components/ui/progress'; import { Progress } from '@/components/ui/progress';
@@ -9,6 +9,8 @@ interface UseModelDownloadToastOptions {
modelName: string; modelName: string;
displayName: string; displayName: string;
enabled?: boolean; enabled?: boolean;
onComplete?: () => void;
onError?: () => void;
} }
/** /**
@@ -19,47 +21,59 @@ export function useModelDownloadToast({
modelName, modelName,
displayName, displayName,
enabled = false, enabled = false,
onComplete,
onError,
}: UseModelDownloadToastOptions) { }: UseModelDownloadToastOptions) {
const { toast } = useToast(); const { toast } = useToast();
const serverUrl = useServerStore((state) => state.serverUrl); const serverUrl = useServerStore((state) => state.serverUrl);
const toastIdRef = useRef<string | null>(null); const toastIdRef = useRef<string | null>(null);
const toastUpdateRef = useRef< // biome-ignore lint: Using any for toast update ref to handle complex toast types
((props: { const toastUpdateRef = useRef<any>(null);
title?: React.ReactNode;
description?: React.ReactNode;
duration?: number;
variant?: 'default' | 'destructive';
open?: boolean;
}) => void) | null
>(null);
const eventSourceRef = useRef<EventSource | null>(null); const eventSourceRef = useRef<EventSource | null>(null);
const formatBytes = (bytes: number): string => { const formatBytes = useCallback((bytes: number): string => {
if (bytes === 0) return '0 B'; if (bytes === 0) return '0 B';
const k = 1024; const k = 1024;
const sizes = ['B', 'KB', 'MB', 'GB']; const sizes = ['B', 'KB', 'MB', 'GB'];
const i = Math.floor(Math.log(bytes) / Math.log(k)); const i = Math.floor(Math.log(bytes) / Math.log(k));
return `${(bytes / Math.pow(k, i)).toFixed(1)} ${sizes[i]}`; return `${(bytes / Math.pow(k, i)).toFixed(1)} ${sizes[i]}`;
}; }, []);
useEffect(() => { useEffect(() => {
console.log('[useModelDownloadToast] useEffect triggered', { enabled, serverUrl, modelName, displayName });
if (!enabled || !serverUrl || !modelName) { if (!enabled || !serverUrl || !modelName) {
console.log('[useModelDownloadToast] Not enabled, skipping');
return; return;
} }
console.log('[useModelDownloadToast] Creating toast and EventSource for:', modelName);
// Create initial toast // Create initial toast
const toastResult = toast({ const toastResult = toast({
title: displayName, title: displayName,
description: 'Starting download...', description: (
<div className="flex items-center gap-2">
<Loader2 className="h-4 w-4 animate-spin" />
<span>Connecting to download...</span>
</div>
),
duration: Infinity, // Don't auto-dismiss, we'll handle it manually duration: Infinity, // Don't auto-dismiss, we'll handle it manually
}); });
toastIdRef.current = toastResult.id; toastIdRef.current = toastResult.id;
toastUpdateRef.current = toastResult.update; toastUpdateRef.current = toastResult.update;
// Subscribe to progress updates via Server-Sent Events // Subscribe to progress updates via Server-Sent Events
const eventSource = new EventSource(`${serverUrl}/models/progress/${modelName}`); 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) => { eventSource.onmessage = (event) => {
console.log('[useModelDownloadToast] Received SSE message:', event.data);
try { try {
const progress = JSON.parse(event.data) as ModelProgress; const progress = JSON.parse(event.data) as ModelProgress;
@@ -86,7 +100,7 @@ export function useModelDownloadToast({
break; break;
case 'downloading': case 'downloading':
statusIcon = <Loader2 className="h-4 w-4 animate-spin" />; statusIcon = <Loader2 className="h-4 w-4 animate-spin" />;
statusText = progress.filename ? `Downloading ${progress.filename}...` : 'Downloading...'; statusText = progress.filename || 'Downloading...';
break; break;
case 'extracting': case 'extracting':
statusIcon = <Loader2 className="h-4 w-4 animate-spin" />; statusIcon = <Loader2 className="h-4 w-4 animate-spin" />;
@@ -121,6 +135,15 @@ export function useModelDownloadToast({
eventSource.close(); eventSource.close();
eventSourceRef.current = null; eventSourceRef.current = null;
// Call callbacks
if (progress.status === 'complete' && onComplete) {
console.log('[useModelDownloadToast] Download complete, calling onComplete callback');
onComplete();
} else if (progress.status === 'error' && onError) {
console.log('[useModelDownloadToast] Download error, calling onError callback');
onError();
}
// Auto-dismiss on completion after delay // Auto-dismiss on completion after delay
if (progress.status === 'complete') { if (progress.status === 'complete') {
setTimeout(() => { setTimeout(() => {
@@ -141,7 +164,8 @@ export function useModelDownloadToast({
}; };
eventSource.onerror = (error) => { eventSource.onerror = (error) => {
console.error('SSE error:', error); console.error('[useModelDownloadToast] SSE error for:', modelName, error);
console.log('[useModelDownloadToast] EventSource readyState:', eventSource.readyState);
eventSource.close(); eventSource.close();
eventSourceRef.current = null; eventSourceRef.current = null;
@@ -162,13 +186,14 @@ export function useModelDownloadToast({
// Cleanup on unmount or when disabled // Cleanup on unmount or when disabled
return () => { return () => {
console.log('[useModelDownloadToast] Cleanup - closing EventSource for:', modelName);
if (eventSourceRef.current) { if (eventSourceRef.current) {
eventSourceRef.current.close(); eventSourceRef.current.close();
eventSourceRef.current = null; eventSourceRef.current = null;
} }
// Note: We don't dismiss the toast here as it might still be showing completion state // Note: We don't dismiss the toast here as it might still be showing completion state
}; };
}, [enabled, serverUrl, modelName, displayName, toast]); }, [enabled, serverUrl, modelName, displayName, toast, formatBytes, onComplete, onError]);
return { return {
isTracking: enabled && eventSourceRef.current !== null, isTracking: enabled && eventSourceRef.current !== null,
+16
View File
@@ -1161,6 +1161,10 @@ async def get_model_status():
import os import os
backend_type = get_backend_type() backend_type = get_backend_type()
task_manager = get_task_manager()
# Get set of currently downloading models
active_downloads = {task.model_name for task in task_manager.get_active_downloads()}
# Try to import scan_cache_dir (might not be available in older versions) # Try to import scan_cache_dir (might not be available in older versions)
try: try:
@@ -1328,10 +1332,18 @@ async def get_model_status():
except Exception: except Exception:
loaded = False loaded = False
# Check if this model is currently being downloaded
is_downloading = config["model_name"] in active_downloads
# If downloading, don't report as downloaded (partial files exist)
if is_downloading:
downloaded = False
statuses.append(models.ModelStatus( statuses.append(models.ModelStatus(
model_name=config["model_name"], model_name=config["model_name"],
display_name=config["display_name"], display_name=config["display_name"],
downloaded=downloaded, downloaded=downloaded,
downloading=is_downloading,
size_mb=size_mb, size_mb=size_mb,
loaded=loaded, loaded=loaded,
)) ))
@@ -1342,10 +1354,14 @@ async def get_model_status():
except Exception: except Exception:
loaded = False loaded = False
# Check if this model is currently being downloaded
is_downloading = config["model_name"] in active_downloads
statuses.append(models.ModelStatus( statuses.append(models.ModelStatus(
model_name=config["model_name"], model_name=config["model_name"],
display_name=config["display_name"], display_name=config["display_name"],
downloaded=False, # Assume not downloaded if check failed downloaded=False, # Assume not downloaded if check failed
downloading=is_downloading,
size_mb=None, size_mb=None,
loaded=loaded, loaded=loaded,
)) ))
+1
View File
@@ -134,6 +134,7 @@ class ModelStatus(BaseModel):
model_name: str model_name: str
display_name: str display_name: str
downloaded: bool downloaded: bool
downloading: bool = False # True if download is in progress
size_mb: Optional[float] = None size_mb: Optional[float] = None
loaded: bool = False loaded: bool = False
Binary file not shown.