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:
Douwe Osinga
2026-03-17 10:30:39 -04:00
committed by GitHub
parent a762fe1000
commit aa3090a519
2 changed files with 67 additions and 34 deletions
+26 -8
View File
@@ -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);
}