diff --git a/crates/goose/src/download_manager.rs b/crates/goose/src/download_manager.rs index 6fc4718cbc..4d599472b8 100644 --- a/crates/goose/src/download_manager.rs +++ b/crates/goose/src/download_manager.rs @@ -46,6 +46,9 @@ pub struct DownloadProgress { pub eta_seconds: Option, /// Error message if failed pub error: Option, + /// Whether the background download task has exited + #[serde(skip)] + pub task_exited: bool, } #[derive(Debug, Clone, Serialize, Deserialize, ToSchema, PartialEq)] @@ -108,8 +111,15 @@ impl DownloadManager { .lock() .map_err(|_| anyhow::anyhow!("Failed to acquire lock"))?; - if downloads.contains_key(&model_id) { - anyhow::bail!("Download already in progress"); + if let Some(existing) = downloads.get(&model_id) { + if existing.status == DownloadStatus::Downloading { + anyhow::bail!("Download already in progress"); + } + if existing.status == DownloadStatus::Cancelled && !existing.task_exited { + anyhow::bail!( + "Download is being cancelled; wait for it to finish before restarting" + ); + } } downloads.insert( @@ -123,6 +133,7 @@ impl DownloadManager { speed_bps: None, eta_seconds: None, error: None, + task_exited: false, }, ); } @@ -148,6 +159,7 @@ impl DownloadManager { if let Some(progress) = downloads.get_mut(&model_id_clone) { progress.status = DownloadStatus::Completed; progress.progress_percent = 100.0; + progress.task_exited = true; } } @@ -162,8 +174,11 @@ impl DownloadManager { if let Ok(mut downloads) = downloads.lock() { if let Some(progress) = downloads.get_mut(&model_id_clone) { - progress.status = DownloadStatus::Failed; + if progress.status != DownloadStatus::Cancelled { + progress.status = DownloadStatus::Failed; + } progress.error = Some(e.to_string()); + progress.task_exited = true; } } } @@ -179,7 +194,10 @@ impl DownloadManager { downloads: &DownloadMap, model_id: &str, ) -> Result<(), anyhow::Error> { - let client = reqwest::Client::new(); + let client = reqwest::Client::builder() + .connect_timeout(std::time::Duration::from_secs(30)) + .read_timeout(std::time::Duration::from_secs(60)) + .build()?; let mut response = client.get(url).send().await?; if !response.status().is_success() { @@ -217,7 +235,7 @@ impl DownloadManager { if should_cancel { let _ = tokio::fs::remove_file(&partial_path).await; - return Ok(()); + anyhow::bail!("Download cancelled"); } file.write_all(&chunk).await?; @@ -264,10 +282,10 @@ impl DownloadManager { pub fn clear_completed(&self, model_id: &str) { if let Ok(mut downloads) = self.downloads.lock() { if let Some(progress) = downloads.get(model_id) { - if progress.status == DownloadStatus::Completed + let is_terminal = progress.status == DownloadStatus::Completed || progress.status == DownloadStatus::Failed - || progress.status == DownloadStatus::Cancelled - { + || progress.status == DownloadStatus::Cancelled; + if is_terminal && progress.task_exited { downloads.remove(model_id); } } diff --git a/ui/desktop/src/components/settings/localInference/LocalInferenceSettings.tsx b/ui/desktop/src/components/settings/localInference/LocalInferenceSettings.tsx index b2191d32de..76015e6ee0 100644 --- a/ui/desktop/src/components/settings/localInference/LocalInferenceSettings.tsx +++ b/ui/desktop/src/components/settings/localInference/LocalInferenceSettings.tsx @@ -30,41 +30,33 @@ export const LocalInferenceSettings = () => { const [settingsOpenFor, setSettingsOpenFor] = useState(null); const { currentModel, currentProvider, refreshCurrentModelAndProvider } = useModelAndProvider(); const downloadSectionRef = useRef(null); + const activePolls = useRef(new Set()); const selectedModelId = currentProvider === 'local' ? currentModel : null; - const loadModels = useCallback(async () => { + const loadModels = useCallback(async (): Promise => { try { const response = await listLocalModels(); if (response.data) { setModels(response.data); + response.data.forEach((model) => { + if (model.status.state === 'Downloading') { + pollDownloadProgress(model.id); + } + }); + + return response.data; } } catch (error) { console.error('Failed to load models:', error); } - }, []); - - // Check for any in-progress downloads when models list changes - const detectActiveDownloads = useCallback(async () => { - for (const model of models) { - if (downloads.has(model.id)) continue; - // Check models that the API reports as downloading - if (model.status.state === 'Downloading') { - pollDownloadProgress(model.id); - } - } + return undefined; // eslint-disable-next-line react-hooks/exhaustive-deps - }, [models, downloads]); + }, []); useEffect(() => { loadModels(); - }, [loadModels]); - - useEffect(() => { - if (models.length > 0) { - detectActiveDownloads(); - } // eslint-disable-next-line react-hooks/exhaustive-deps - }, [models]); + }, []); const selectModel = async (modelId: string) => { try { @@ -97,6 +89,14 @@ export const LocalInferenceSettings = () => { }, []); const pollDownloadProgress = (modelId: string) => { + if (activePolls.current.has(modelId)) return; + activePolls.current.add(modelId); + + const stopPolling = (interval: ReturnType) => { + clearInterval(interval); + activePolls.current.delete(modelId); + }; + const interval = setInterval(async () => { try { const response = await getLocalModelDownloadProgress({ path: { model_id: modelId } }); @@ -105,7 +105,7 @@ export const LocalInferenceSettings = () => { setDownloads((prev) => new Map(prev).set(modelId, progress)); if (progress.status === 'completed') { - clearInterval(interval); + stopPolling(interval); setDownloads((prev) => { const next = new Map(prev); next.delete(modelId); @@ -113,15 +113,20 @@ export const LocalInferenceSettings = () => { }); await loadModels(); await selectModel(modelId); - } else if (progress.status === 'failed') { - clearInterval(interval); + } else if (progress.status === 'failed' || progress.status === 'cancelled') { + stopPolling(interval); + setDownloads((prev) => { + const next = new Map(prev); + next.delete(modelId); + return next; + }); await loadModels(); } } else { - clearInterval(interval); + stopPolling(interval); } } catch { - clearInterval(interval); + stopPolling(interval); } }, 1000); }; @@ -134,6 +139,7 @@ export const LocalInferenceSettings = () => { next.delete(modelId); return next; }); + await loadModels(); } catch (error) { console.error('Failed to cancel download:', error); } @@ -143,7 +149,16 @@ export const LocalInferenceSettings = () => { if (!window.confirm('Delete this model? You can re-download it later.')) return; try { await deleteLocalModel({ path: { model_id: modelId } }); - await loadModels(); + const updatedModels = await loadModels(); + + if (selectedModelId === modelId && updatedModels) { + const remainingDownloaded = updatedModels.filter( + (m) => m.id !== modelId && m.status.state === 'Downloaded' + ); + if (remainingDownloaded.length > 0) { + selectModel(remainingDownloaded[0].id); + } + } } catch (error) { console.error('Failed to delete model:', error); }