11import { useMutation , useQuery , useQueryClient } from '@tanstack/react-query' ;
22import { Download , Loader2 , Trash2 } from 'lucide-react' ;
3- import { useState } from 'react' ;
3+ import { useCallback , useState } from 'react' ;
44import {
55 AlertDialog ,
66 AlertDialogAction ,
@@ -17,7 +17,6 @@ import { Card, CardContent, CardDescription, CardHeader, CardTitle } from '@/com
1717import { useToast } from '@/components/ui/use-toast' ;
1818import { apiClient } from '@/lib/api/client' ;
1919import { useModelDownloadToast } from '@/lib/hooks/useModelDownloadToast' ;
20- import { ModelProgress } from './ModelProgress' ;
2120
2221export function ModelManagement ( ) {
2322 const { toast } = useToast ( ) ;
@@ -27,15 +26,36 @@ export function ModelManagement() {
2726
2827 const { data : modelStatus , isLoading } = useQuery ( {
2928 queryKey : [ 'modelStatus' ] ,
30- queryFn : ( ) => apiClient . getModelStatus ( ) ,
29+ queryFn : async ( ) => {
30+ console . log ( '[Query] Fetching model status' ) ;
31+ const result = await apiClient . getModelStatus ( ) ;
32+ console . log ( '[Query] Model status fetched:' , result ) ;
33+ return result ;
34+ } ,
3135 refetchInterval : 5000 , // Refresh every 5 seconds
3236 } ) ;
3337
38+ // Callbacks for download completion
39+ const handleDownloadComplete = useCallback ( ( ) => {
40+ console . log ( '[ModelManagement] Download complete, clearing state' ) ;
41+ setDownloadingModel ( null ) ;
42+ setDownloadingDisplayName ( null ) ;
43+ queryClient . invalidateQueries ( { queryKey : [ 'modelStatus' ] } ) ;
44+ } , [ queryClient ] ) ;
45+
46+ const handleDownloadError = useCallback ( ( ) => {
47+ console . log ( '[ModelManagement] Download error, clearing state' ) ;
48+ setDownloadingModel ( null ) ;
49+ setDownloadingDisplayName ( null ) ;
50+ } , [ ] ) ;
51+
3452 // Use progress toast hook for the downloading model
3553 useModelDownloadToast ( {
3654 modelName : downloadingModel || '' ,
3755 displayName : downloadingDisplayName || '' ,
3856 enabled : ! ! downloadingModel && ! ! downloadingDisplayName ,
57+ onComplete : handleDownloadComplete ,
58+ onError : handleDownloadError ,
3959 } ) ;
4060
4161 const [ deleteDialogOpen , setDeleteDialogOpen ] = useState ( false ) ;
@@ -45,44 +65,69 @@ export function ModelManagement() {
4565 sizeMb ?: number ;
4666 } | null > ( null ) ;
4767
48- const downloadMutation = useMutation ( {
49- mutationFn : ( modelName : string ) => {
68+ const handleDownload = async ( modelName : string ) => {
69+ console . log ( '[Download] Button clicked for:' , modelName , 'at' , new Date ( ) . toISOString ( ) ) ;
70+
71+ // Find display name
72+ const model = modelStatus ?. models . find ( ( m ) => m . model_name === modelName ) ;
73+ const displayName = model ?. display_name || modelName ;
74+
75+ try {
76+ // IMPORTANT: Call the API FIRST before setting state
77+ // Setting state enables the SSE EventSource in useModelDownloadToast,
78+ // which can block/delay the download fetch due to HTTP/1.1 connection limits
79+ console . log ( '[Download] Calling download API for:' , modelName ) ;
80+ const result = await apiClient . triggerModelDownload ( modelName ) ;
81+ console . log ( '[Download] Download API responded:' , result ) ;
82+
83+ // NOW set state to enable SSE tracking (after download has started on backend)
5084 setDownloadingModel ( modelName ) ;
51- // Find display name from model status
52- const model = modelStatus ?. models . find ( ( m ) => m . model_name === modelName ) ;
53- setDownloadingDisplayName ( model ?. display_name || modelName ) ;
54- return apiClient . triggerModelDownload ( modelName ) ;
55- } ,
56- onSuccess : ( ) => {
57- // Download completed - clear state and refetch status
58- setDownloadingModel ( null ) ;
59- setDownloadingDisplayName ( null ) ;
85+ setDownloadingDisplayName ( displayName ) ;
86+
87+ // Download initiated successfully - state will be cleared when SSE reports completion
88+ // or by the polling interval detecting the model is downloaded
6089 queryClient . invalidateQueries ( { queryKey : [ 'modelStatus' ] } ) ;
61- } ,
62- onError : ( error : Error ) => {
90+ } catch ( error ) {
91+ console . error ( '[Download] Download failed:' , error ) ;
6392 setDownloadingModel ( null ) ;
6493 setDownloadingDisplayName ( null ) ;
6594 toast ( {
6695 title : 'Download failed' ,
67- description : error . message ,
96+ description : error instanceof Error ? error . message : 'Unknown error' ,
6897 variant : 'destructive' ,
6998 } ) ;
70- } ,
71- } ) ;
99+ }
100+ } ;
72101
73102 const deleteMutation = useMutation ( {
74- mutationFn : ( modelName : string ) => apiClient . deleteModel ( modelName ) ,
75- onSuccess : ( ) => {
103+ mutationFn : async ( modelName : string ) => {
104+ console . log ( '[Delete] Deleting model:' , modelName ) ;
105+ const result = await apiClient . deleteModel ( modelName ) ;
106+ console . log ( '[Delete] Model deleted successfully:' , modelName ) ;
107+ return result ;
108+ } ,
109+ onSuccess : async ( _data , _modelName ) => {
110+ console . log ( '[Delete] onSuccess - showing toast and invalidating queries' ) ;
76111 toast ( {
77112 title : 'Model deleted' ,
78113 description : `${ modelToDelete ?. displayName || 'Model' } has been deleted successfully.` ,
79114 } ) ;
80115 setDeleteDialogOpen ( false ) ;
81116 setModelToDelete ( null ) ;
82- // Refetch status to update UI
83- queryClient . invalidateQueries ( { queryKey : [ 'modelStatus' ] } ) ;
117+ // Invalidate AND explicitly refetch to ensure UI updates
118+ // Using refetchType: 'all' ensures we refetch even if the query is stale
119+ console . log ( '[Delete] Invalidating modelStatus query' ) ;
120+ await queryClient . invalidateQueries ( {
121+ queryKey : [ 'modelStatus' ] ,
122+ refetchType : 'all' ,
123+ } ) ;
124+ // Also explicitly refetch to guarantee fresh data
125+ console . log ( '[Delete] Explicitly refetching modelStatus query' ) ;
126+ await queryClient . refetchQueries ( { queryKey : [ 'modelStatus' ] } ) ;
127+ console . log ( '[Delete] Query refetched' ) ;
84128 } ,
85129 onError : ( error : Error ) => {
130+ console . log ( '[Delete] onError:' , error ) ;
86131 toast ( {
87132 title : 'Delete failed' ,
88133 description : error . message ,
@@ -124,7 +169,7 @@ export function ModelManagement() {
124169 < ModelItem
125170 key = { model . model_name }
126171 model = { model }
127- onDownload = { ( ) => downloadMutation . mutate ( model . model_name ) }
172+ onDownload = { ( ) => handleDownload ( model . model_name ) }
128173 onDelete = { ( ) => {
129174 setModelToDelete ( {
130175 name : model . model_name ,
@@ -152,7 +197,7 @@ export function ModelManagement() {
152197 < ModelItem
153198 key = { model . model_name }
154199 model = { model }
155- onDownload = { ( ) => downloadMutation . mutate ( model . model_name ) }
200+ onDownload = { ( ) => handleDownload ( model . model_name ) }
156201 onDelete = { ( ) => {
157202 setModelToDelete ( {
158203 name : model . model_name ,
@@ -168,21 +213,6 @@ export function ModelManagement() {
168213 </ div >
169214 </ div >
170215
171- { /* Progress indicators */ }
172- < div className = "pt-4 border-t" >
173- < h3 className = "text-sm font-semibold mb-3 text-muted-foreground" >
174- Download Progress
175- </ h3 >
176- < div className = "space-y-2" >
177- { modelStatus . models . map ( ( model ) => (
178- < ModelProgress
179- key = { model . model_name }
180- modelName = { model . model_name }
181- displayName = { model . display_name }
182- />
183- ) ) }
184- </ div >
185- </ div >
186216 </ div >
187217 ) : null }
188218 </ CardContent >
@@ -235,16 +265,20 @@ interface ModelItemProps {
235265 model_name : string ;
236266 display_name : string ;
237267 downloaded : boolean ;
268+ downloading ?: boolean ; // From server - true if download in progress
238269 size_mb ?: number ;
239270 loaded : boolean ;
240271 } ;
241272 onDownload : ( ) => void ;
242273 onDelete : ( ) => void ;
243- isDownloading : boolean ;
274+ isDownloading : boolean ; // Local state - true if user just clicked download
244275 formatSize : ( sizeMb ?: number ) => string ;
245276}
246277
247278function ModelItem ( { model, onDownload, onDelete, isDownloading, formatSize } : ModelItemProps ) {
279+ // Use server's downloading state OR local state (for immediate feedback before server updates)
280+ const showDownloading = model . downloading || isDownloading ;
281+
248282 return (
249283 < div className = "flex items-center justify-between p-3 border rounded-lg" >
250284 < div className = "flex-1" >
@@ -255,20 +289,21 @@ function ModelItem({ model, onDownload, onDelete, isDownloading, formatSize }: M
255289 Loaded
256290 </ Badge >
257291 ) }
258- { model . downloaded && ! model . loaded && (
292+ { /* Only show Downloaded if actually downloaded AND not downloading */ }
293+ { model . downloaded && ! model . loaded && ! showDownloading && (
259294 < Badge variant = "secondary" className = "text-xs" >
260295 Downloaded
261296 </ Badge >
262297 ) }
263298 </ div >
264- { model . downloaded && model . size_mb && (
299+ { model . downloaded && model . size_mb && ! showDownloading && (
265300 < div className = "text-xs text-muted-foreground mt-1" >
266301 Size: { formatSize ( model . size_mb ) }
267302 </ div >
268303 ) }
269304 </ div >
270305 < div className = "flex items-center gap-2" >
271- { model . downloaded ? (
306+ { model . downloaded && ! showDownloading ? (
272307 < div className = "flex items-center gap-2" >
273308 < div className = "flex items-center gap-1 text-sm text-muted-foreground" >
274309 < span > Ready</ span >
@@ -283,19 +318,15 @@ function ModelItem({ model, onDownload, onDelete, isDownloading, formatSize }: M
283318 < Trash2 className = "h-4 w-4" />
284319 </ Button >
285320 </ div >
321+ ) : showDownloading ? (
322+ < Button size = "sm" variant = "outline" disabled >
323+ < Loader2 className = "h-4 w-4 mr-2 animate-spin" />
324+ Downloading...
325+ </ Button >
286326 ) : (
287- < Button size = "sm" onClick = { onDownload } disabled = { isDownloading } variant = "outline" >
288- { isDownloading ? (
289- < >
290- < Loader2 className = "h-4 w-4 mr-2 animate-spin" />
291- Downloading...
292- </ >
293- ) : (
294- < >
295- < Download className = "h-4 w-4 mr-2" />
296- Download
297- </ >
298- ) }
327+ < Button size = "sm" onClick = { onDownload } variant = "outline" >
328+ < Download className = "h-4 w-4 mr-2" />
329+ Download
299330 </ Button >
300331 ) }
301332 </ div >
0 commit comments