diff --git a/core/download/api.py b/core/download/api.py index 1ddbffb..980d53e 100644 --- a/core/download/api.py +++ b/core/download/api.py @@ -212,6 +212,9 @@ context._delete_python_partial_download_file = _bind( ) context._delete_xet_partial_file = _bind(aria2_backend.delete_xet_partial_file) context._ensure_aria2_daemon = _bind(aria2_backend.ensure_aria2_daemon) +context._recover_aria2_missing_control_file = _bind( + aria2_backend.recover_aria2_missing_control_file +) context._force_remove_aria2_transfer = _bind( aria2_backend.force_remove_aria2_transfer ) @@ -226,6 +229,7 @@ context._resolve_download_url_for_aria2 = _bind( context._run_aria2_desired_state_worker = _bind( aria2_backend.run_aria2_desired_state_worker ) +context._aria2_idle_stop_worker = _bind(aria2_backend.aria2_idle_stop_worker) context._schedule_aria2_idle_stop = _bind(aria2_backend.schedule_aria2_idle_stop) context._set_download_progress_status = _bind( aria2_backend.set_download_progress_status @@ -238,6 +242,9 @@ context._queue_aria2_desired_state = _bind( aria2_backend.queue_aria2_desired_state ) context._aria2_action_error_is_ok = aria2_backend.aria2_action_error_is_ok +context._is_aria2_missing_control_file_error = ( + aria2_backend.is_aria2_missing_control_file_error +) context.Aria2Error = aria2_backend.Aria2Error context._find_free_port = aria2_backend.find_free_port context._parse_aria2_int = aria2_backend.parse_aria2_int diff --git a/core/download/aria2_backend.py b/core/download/aria2_backend.py index c13d194..fb1e96d 100644 --- a/core/download/aria2_backend.py +++ b/core/download/aria2_backend.py @@ -92,6 +92,16 @@ def parse_aria2_int(value: Any) -> int: return 0 +def is_aria2_missing_control_file_error(error: Any) -> bool: + """Identify aria2's stale-output error caused by a missing .aria2 file.""" + message = str(error or "").lower() + return ( + "control file" in message + and "does not exist" in message + and "allow-overwrite" in message + ) + + def resolve_aria2_completed_path(status: Dict[str, Any], default_path: str) -> str: """Use aria2's completed file path when it reports one.""" files = status.get("files") @@ -605,6 +615,58 @@ def ensure_aria2_daemon( ) +def recover_aria2_missing_control_file( + download_id: str, + gid: str, + settings: Optional[Dict[str, Any]] = None, + *, + dependencies: Any = None, +) -> bool: + """Restart a stale resolver daemon when a download loses its control file.""" + facade = _require_dependencies(dependencies) + with facade.aria2_lock: + other_transfers = [ + active_id + for active_id in facade.aria2_transfers + if active_id != download_id + ] + if other_transfers: + log.warning( + "Skipping automatic aria2 recovery because other downloads are active: " + f"{', '.join(other_transfers)}" + ) + return False + facade.aria2_transfers.pop(download_id, None) + + try: + if gid: + try: + facade._aria2_rpc( + "aria2.removeDownloadResult", + [gid], + ) + except Exception as exc: + log.debug(f"Could not remove stale aria2 result {gid}: {exc}") + + stopped = facade.stop_aria2_daemon(reason="stale-control-file") + if not stopped.get("success"): + log.warning( + "Could not restart aria2 for stale control-file recovery: " + f"{stopped.get('error') or stopped.get('message') or 'unknown error'}" + ) + return False + + facade._ensure_aria2_daemon(settings) + log.info( + "Restarted aria2 after detecting a missing control file; " + f"retrying download {download_id}" + ) + return True + except Exception as exc: + log.warning(f"Automatic aria2 recovery failed for {download_id}: {exc}") + return False + + def aria2_tell_status(gid: str, *, dependencies: Any = None) -> Dict[str, Any]: """Read an aria2 transfer status with retries for transient resets.""" facade = _require_dependencies(dependencies) @@ -614,6 +676,7 @@ def aria2_tell_status(gid: str, *, dependencies: Any = None) -> Dict[str, Any]: "totalLength", "completedLength", "downloadSpeed", + "connections", "errorMessage", "files", ] @@ -721,21 +784,59 @@ def download_file_with_aria2( if header_values: options["header"] = header_values - gid = facade._aria2_rpc( - "aria2.addUri", - [[aria2_url], options], - ) - if not isinstance(gid, str) or not gid: - raise Aria2Error("aria2 did not return a download gid") + def add_aria2_transfer() -> str: + new_gid = facade._aria2_rpc( + "aria2.addUri", + [[aria2_url], options], + ) + if not isinstance(new_gid, str) or not new_gid: + raise Aria2Error("aria2 did not return a download gid") - with facade.aria2_lock: - facade.aria2_transfers[download_id] = { - "gid": gid, - "path": dest_path, - } - with facade.download_lock: - facade.download_progress[download_id]["aria2_gid"] = gid - facade.download_progress[download_id]["status"] = "downloading" + with facade.aria2_lock: + facade.aria2_transfers[download_id] = { + "gid": new_gid, + "path": dest_path, + } + with facade.download_lock: + facade.download_progress[download_id].update( + { + "aria2_gid": new_gid, + "status": "downloading", + "error": None, + "progress": 0, + "downloaded": 0, + "total_size": 0, + "speed": 0, + } + ) + return new_gid + + recovery_attempted = False + try: + gid = add_aria2_transfer() + except Exception as exc: + if ( + not recovery_attempted + and download_id not in facade.cancelled_downloads + and facade._is_aria2_missing_control_file_error(exc) + ): + recovery_attempted = True + facade._set_download_progress_status( + download_id, + "starting", + error=None, + speed=0, + ) + if facade._recover_aria2_missing_control_file( + download_id, + "", + settings, + ): + gid = add_aria2_transfer() + else: + raise + else: + raise log.info(f"Starting aria2 download: {filename}") last_cli_log = start_time @@ -767,6 +868,7 @@ def download_file_with_aria2( total_size = facade._parse_aria2_int(status.get("totalLength")) downloaded = facade._parse_aria2_int(status.get("completedLength")) speed = facade._parse_aria2_int(status.get("downloadSpeed")) + connections = facade._parse_aria2_int(status.get("connections")) progress = int((downloaded / total_size) * 100) if total_size > 0 else 0 mapped_status = { "active": "downloading", @@ -790,6 +892,7 @@ def download_file_with_aria2( if mapped_status in {"paused", "completed"} else speed ), + "connections": connections, "download_backend": "aria2", "aria2_gid": gid, } @@ -862,6 +965,26 @@ def download_file_with_aria2( if state == "error": error_msg = status.get("errorMessage") or "aria2 download failed" + if ( + not recovery_attempted + and download_id not in facade.cancelled_downloads + and facade._is_aria2_missing_control_file_error(error_msg) + ): + recovery_attempted = True + facade._set_download_progress_status( + download_id, + "starting", + error=None, + speed=0, + ) + recovered = facade._recover_aria2_missing_control_file( + download_id, + gid, + settings, + ) + if recovered: + gid = add_aria2_transfer() + continue with facade.download_lock: facade.download_progress[download_id]["status"] = "error" facade.download_progress[download_id]["error"] = error_msg diff --git a/core/download/orchestrator.py b/core/download/orchestrator.py index 71819e1..5713918 100644 --- a/core/download/orchestrator.py +++ b/core/download/orchestrator.py @@ -3,6 +3,36 @@ from typing import Any, Callable, Dict, Optional +def find_active_download_for_path( + download_progress: Dict[str, Dict[str, Any]], + target_path: str, + os_module: Any, + *, + exclude_id: str = "", +) -> Optional[str]: + """Return the active download using the same normalized destination path.""" + if not target_path: + return None + normalized_target = os_module.path.normcase( + os_module.path.abspath(target_path) + ) + terminal_statuses = {"completed", "error", "cancelled"} + for download_id, progress in download_progress.items(): + if download_id == exclude_id or not isinstance(progress, dict): + continue + if str(progress.get("status") or "").lower() in terminal_statuses: + continue + existing_path = progress.get("path") + if not existing_path: + continue + normalized_existing = os_module.path.normcase( + os_module.path.abspath(str(existing_path)) + ) + if normalized_existing == normalized_target: + return download_id + return None + + def download_file( url: str, dest_path: str, @@ -378,11 +408,36 @@ def download_model( "error": "Download target is outside the selected model directory", } + with download_lock: + active_download_id = find_active_download_for_path( + download_progress, + dest_path, + os, + exclude_id=download_id, + ) + if active_download_id: + message = f"A download is already active for this file: {dest_path}" + log.info(f"{message} (download_id={active_download_id})") + return { + "success": False, + "download_id": download_id, + "active_download_id": active_download_id, + "error": message, + "path": dest_path, + } + resume_aria2_partial = bool( facade._download_backend_from_settings() == "aria2" and os.path.isfile(dest_path) and os.path.isfile(f"{dest_path}.aria2") ) + if ( + facade._download_backend_from_settings() == "aria2" + and os.path.isfile(f"{dest_path}.aria2") + and not os.path.isfile(dest_path) + ): + log.warning(f"Removing orphaned aria2 control file: {dest_path}.aria2") + facade._delete_partial_download_files(dest_path) if resume_aria2_partial: log.info(f"Resuming partial aria2 download: {dest_path}") @@ -519,6 +574,7 @@ def start_background_download( time = facade.time download_progress = facade.download_progress download_lock = facade.download_lock + log = facade.log download_id = facade.generate_download_id() filename = facade.sanitize_download_filename(filename) @@ -529,6 +585,17 @@ def start_background_download( initial_path = os.path.join(initial_directory, filename) if initial_directory and filename else "" with download_lock: + active_download_id = find_active_download_for_path( + download_progress, + initial_path, + os, + ) + if active_download_id: + log.info( + f"Reusing active download for {initial_path} " + f"(download_id={active_download_id})" + ) + return active_download_id download_progress[download_id] = { "status": "starting", "progress": 0, diff --git a/tests/test_downloader_boundaries.py b/tests/test_downloader_boundaries.py index dba9903..2ffa8f7 100644 --- a/tests/test_downloader_boundaries.py +++ b/tests/test_downloader_boundaries.py @@ -1,10 +1,13 @@ import os import tempfile +import threading import unittest from types import SimpleNamespace from unittest.mock import MagicMock, patch from core.download.api import context as downloader +from core.download.aria2_backend import recover_aria2_missing_control_file +from core.download.orchestrator import find_active_download_for_path from core.path_utils import normalize_metadata_file_path @@ -132,6 +135,76 @@ class DownloaderBoundaryTests(unittest.TestCase): self.assertFalse(downloader._aria2_action_error_is_ok("paused", "permission denied")) self.assertFalse(downloader._aria2_action_error_is_ok("completed", "already paused")) + def test_aria2_missing_control_file_error_is_detected(self): + self.assertTrue( + downloader._is_aria2_missing_control_file_error( + "File exists, but a control file(*.aria2) does not exist. " + "Download was canceled; add --allow-overwrite=true." + ) + ) + self.assertFalse( + downloader._is_aria2_missing_control_file_error( + "TLS handshake failed" + ) + ) + + def test_aria2_missing_control_file_recovery_restarts_only_when_alone(self): + facade = SimpleNamespace( + aria2_lock=threading.Lock(), + aria2_transfers={"download-1": {"gid": "gid-1"}}, + _aria2_rpc=MagicMock(), + stop_aria2_daemon=MagicMock(return_value={"success": True}), + _ensure_aria2_daemon=MagicMock(), + ) + + self.assertTrue( + recover_aria2_missing_control_file( + "download-1", + "gid-1", + {"download_backend": "aria2"}, + dependencies=facade, + ) + ) + facade.stop_aria2_daemon.assert_called_once_with( + reason="stale-control-file" + ) + facade._ensure_aria2_daemon.assert_called_once() + self.assertNotIn("download-1", facade.aria2_transfers) + + facade.aria2_transfers.update( + { + "download-1": {"gid": "gid-1"}, + "download-2": {"gid": "gid-2"}, + } + ) + facade.stop_aria2_daemon.reset_mock() + self.assertFalse( + recover_aria2_missing_control_file( + "download-1", + "gid-1", + dependencies=facade, + ) + ) + facade.stop_aria2_daemon.assert_not_called() + + def test_active_download_path_is_reused_case_insensitively(self): + active_id = find_active_download_for_path( + { + "active": { + "status": "downloading", + "path": r"C:\Models\Qwen.safetensors", + }, + "finished": { + "status": "completed", + "path": r"C:\Models\Other.safetensors", + }, + }, + r"c:\models\qwen.safetensors", + os, + ) + + self.assertEqual("active", active_id) + def test_progress_snapshot_is_copied_and_completed_entries_can_be_cleared(self): with patch.dict( downloader.download_progress, diff --git a/web/css/resolver-shell.css b/web/css/resolver-shell.css index 13bbee4..7752412 100644 --- a/web/css/resolver-shell.css +++ b/web/css/resolver-shell.css @@ -738,8 +738,8 @@ #model-resolver-modal .mr-download-queue-actions { display: inline-flex; align-items: center; - gap: 3px; - padding: 3px 5px; + gap: 2px; + padding: 2px 4px; background: var(--mr-surface-glass-02); border: 1px solid var(--mr-border-strong); border-radius: var(--mr-missing-chip-radius); @@ -749,8 +749,8 @@ display: inline-flex; align-items: center; justify-content: center; - width: 32px; - height: 32px; + width: 28px; + height: 28px; padding: 0; cursor: pointer; color: var(--mr-text-muted); @@ -776,8 +776,8 @@ } .mr-download-queue-action-icon { - width: 20px; - height: 20px; + width: 18px; + height: 18px; } #model-resolver-modal .mr-download-queue-pause { diff --git a/web/resolver/actions/queue_methods.js b/web/resolver/actions/queue_methods.js index 23ca434..b3648cb 100644 --- a/web/resolver/actions/queue_methods.js +++ b/web/resolver/actions/queue_methods.js @@ -1204,14 +1204,16 @@ export const queueMethods = { const isAria2 = backend === 'aria2'; const isPaused = status === 'paused'; const statusLabel = status === 'downloading' - ? (presentation.isFinalizing ? 'Finalizing' : `${percentLabel}%`) + ? (presentation.isFinalizing ? 'Finalizing' : '') : (status === 'starting' ? 'Starting' : (isPaused ? 'Paused' : status.replace(/_/g, ' '))); - const sizeText = presentation.isFinalizing && logicalTotal + const sizeTextBase = presentation.isFinalizing && logicalTotal ? `Finalizing file: ${logicalDownloaded} / ${logicalTotal}` : (presentation.totalSize ? `${downloaded} / ${total}` : downloaded); + const sizeText = presentation.totalSize ? `${sizeTextBase} (${percentLabel}%)` : sizeTextBase; const speed = Number(progress.speed) || 0; const speedText = speed > 0 ? `${this.formatBytes(speed)}/s` : ''; const etaText = this.getDownloadEtaText?.(progress)?.replace(/^ETA\s+/i, 'Remaining ') || ''; + const connections = Number(progress.connections) || 0; const targetPath = progress.directory || info.downloadDirectory || info.downloadPath || ''; const workflowLabel = this.getDownloadWorkflowLabel?.(info) || 'Unknown workflow'; const contextModel = this.getDownloadQueueContext?.(progress, info, workflowLabel, downloadId); @@ -1236,7 +1238,9 @@ export const queueMethods = { html += `
`; html += `
`; html += `${this.escapeHtml(filename)}`; - html += `${this.escapeHtml(statusLabel)}`; + if (statusLabel) { + html += `${this.escapeHtml(statusLabel)}`; + } html += `
`; html += `
`; html += actionButton( @@ -1262,6 +1266,9 @@ export const queueMethods = { if (etaText) { html += `${this.escapeHtml(etaText)}`; } + if (connections > 0) { + html += `${renderIcon('network', 'mr-download-queue-stat-icon')}${connections}`; + } if (!speedText && !etaText) { html += `${this.escapeHtml(statusLabel)}`; } diff --git a/web/utils/icon_utils.js b/web/utils/icon_utils.js index 6a3b87a..14de397 100644 --- a/web/utils/icon_utils.js +++ b/web/utils/icon_utils.js @@ -25,6 +25,7 @@ const SVG_ICONS = { pause: ``, play: ``, info: ``, + network: ``, cloudDownload: ``, skull: ``, shield: ``,