mirror of
https://github.com/block/goose.git
synced 2026-07-17 12:56:20 +02:00
Fix download manager (#7933)
Co-authored-by: Douwe Osinga <douwe@squareup.com> Co-authored-by: jh-block <jhugo@block.xyz> Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
@@ -46,6 +46,9 @@ pub struct DownloadProgress {
|
||||
pub eta_seconds: Option<u64>,
|
||||
/// Error message if failed
|
||||
pub error: Option<String>,
|
||||
/// 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);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -30,41 +30,33 @@ export const LocalInferenceSettings = () => {
|
||||
const [settingsOpenFor, setSettingsOpenFor] = useState<string | null>(null);
|
||||
const { currentModel, currentProvider, refreshCurrentModelAndProvider } = useModelAndProvider();
|
||||
const downloadSectionRef = useRef<HTMLDivElement>(null);
|
||||
const activePolls = useRef(new Set<string>());
|
||||
const selectedModelId = currentProvider === 'local' ? currentModel : null;
|
||||
|
||||
const loadModels = useCallback(async () => {
|
||||
const loadModels = useCallback(async (): Promise<LocalModelResponse[] | undefined> => {
|
||||
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<typeof setInterval>) => {
|
||||
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);
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user