feat: improve aria2 download resilience and queue presentation

2026-08-03 15:55:08 +02:00
parent 88b6c6761e
commit 415bcb93f4
7 changed files with 301 additions and 23 deletions
@@ -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
@@ -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
@@ -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,
@@ -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,
@@ -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 {
@@ -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 += `<div class="mr-download-queue-header">`;
            html += `<div class="mr-download-queue-heading">`;
            html += `<span class="mr-download-queue-filename" data-tooltip="${this.escapeHtml(filename)}">${this.escapeHtml(filename)}</span>`;
-            html += `<span class="mr-download-queue-status">${this.escapeHtml(statusLabel)}</span>`;
+            if (statusLabel) {
+                html += `<span class="mr-download-queue-status">${this.escapeHtml(statusLabel)}</span>`;
+            }
            html += `</div>`;
            html += `<div class="mr-download-queue-actions" role="group" aria-label="Download actions">`;
            html += actionButton(
@@ -1262,6 +1266,9 @@ export const queueMethods = {
            if (etaText) {
                html += `<span class="mr-download-queue-stat"><span>${this.escapeHtml(etaText)}</span></span>`;
            }
+            if (connections > 0) {
+                html += `<span class="mr-download-queue-stat" data-tooltip="Active download connections">${renderIcon('network', 'mr-download-queue-stat-icon')}<span>${connections}</span></span>`;
+            }
            if (!speedText && !etaText) {
                html += `<span class="mr-download-queue-stat mr-download-queue-stat-muted">${this.escapeHtml(statusLabel)}</span>`;
            }
@@ -25,6 +25,7 @@ const SVG_ICONS = {
    pause: `<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24" fill="none" stroke="{color}" stroke-width="1.9" stroke-linecap="round" stroke-linejoin="round" aria-hidden="true"><rect x="6.5" y="4.5" width="4" height="15" rx="1"></rect><rect x="13.5" y="4.5" width="4" height="15" rx="1"></rect></svg>`,
    play: `<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24" fill="none" stroke="{color}" stroke-width="2.1" stroke-linecap="round" stroke-linejoin="round" aria-hidden="true"><path d="m8 5 11 7-11 7V5Z"></path></svg>`,
    info: `<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24" fill="none" stroke="{color}" stroke-width="2" stroke-linecap="round" stroke-linejoin="round" aria-hidden="true"><circle cx="12" cy="12" r="9"></circle><path d="M12 11v5"></path><path d="M12 8h.01"></path></svg>`,
+    network: `<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24" fill="none" stroke="{color}" stroke-width="2" stroke-linecap="round" stroke-linejoin="round" aria-hidden="true"><circle cx="6" cy="12" r="2.5"></circle><circle cx="18" cy="6" r="2.5"></circle><circle cx="18" cy="18" r="2.5"></circle><path d="m8.2 10.9 7.5-3.7"></path><path d="m8.2 13.1 7.5 3.7"></path></svg>`,
    cloudDownload: `<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24" fill="none" stroke="{color}" stroke-width="2" stroke-linecap="round" stroke-linejoin="round" aria-hidden="true"><path d="M12 13v8l-4-4"></path><path d="m12 21 4-4"></path><path d="M4.393 15.269A7 7 0 1 1 15.71 8h1.79a4.5 4.5 0 0 1 2.436 8.284"></path></svg>`,
    skull: `<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24" fill="none" stroke="{color}" stroke-width="2" stroke-linecap="round" stroke-linejoin="round" aria-hidden="true"><path d="m12.5 17-.5-1-.5 1h1z"></path><path d="M15 22a1 1 0 0 0 1-1v-1a2 2 0 0 0 1.56-3.25 8 8 0 1 0-11.12 0A2 2 0 0 0 8 20v1a1 1 0 0 0 1 1z"></path><circle cx="15" cy="12" r="1"></circle><circle cx="9" cy="12" r="1"></circle></svg>`,
    shield: `<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24" fill="none" stroke="{color}" stroke-width="2" stroke-linecap="round" stroke-linejoin="round" aria-hidden="true"><path d="M20 13c0 5-3.5 7.5-7.66 8.95a1 1 0 0 1-.67-.01C7.5 20.5 4 18 4 13V6a1 1 0 0 1 1-1c2 0 4.5-1.2 6.24-2.72a1.17 1.17 0 0 1 1.52 0C14.51 3.81 17 5 19 5a1 1 0 0 1 1 1z"></path></svg>`,