fix(downloads): preserve aria2 partial pair and refresh expired CivitAI signed URLs

A failed aria2 transfer deleted the partial payload while keeping its
.aria2 control file, and "No URI available" (expired CivitAI signed URL)
was treated as a permanent failure, wasting nearly-complete downloads.

- Re-schedule the transfer with a freshly resolved signed URL and
  continue=true when aria2 reports "No URI available", bounded by
  MAX_TRANSFER_RECOVERY_ATTEMPTS
- Keep payload and .aria2 control file together as a resumable pair
  after a failed transfer instead of deleting the payload
- Report and remove orphaned .aria2 control files that have no payload,
  both after failures and when restoring persisted downloads

Fixes #1088

py/services/aria2_downloader.py

@@ -82,6 +82,17 @@ CIVITAI_DOWNLOAD_URL_PREFIXES = (
 )
 
 
+def _is_no_uri_available_error(message: str) -> bool:
+    """Return True for aria2's "No URI available" transfer failure.
+
+    aria2 reports this when every URI for the transfer has become unusable.
+    For CivitAI downloads this typically means the temporary signed URL
+    expired mid-download; the transfer can be recovered by resolving a fresh
+    signed URL and re-scheduling with ``continue=true``.
+    """
+    return "no uri available" in message.lower()
+
+
 class Aria2Error(RuntimeError):
     """Raised when aria2 integration fails."""
 
@@ -145,8 +156,11 @@ class Aria2Downloader:
         disappears (e.g. another download restarted the daemon and
         ``close()`` cleared ``_transfers``) or the RPC becomes unreachable,
         the transfer is re-scheduled with ``continue=true`` so the download
-        resumes from the on-disk ``.aria2`` control file. Recovery is bounded
-        by ``MAX_TRANSFER_RECOVERY_ATTEMPTS``.
+        resumes from the on-disk ``.aria2`` control file.  The same
+        re-scheduling happens when aria2 fails with "No URI available"
+        (typically an expired CivitAI signed URL): a fresh URL is resolved
+        and the partial download continues.  Recovery is bounded by
+        ``MAX_TRANSFER_RECOVERY_ATTEMPTS``.
         """
 
         await self._ensure_process()
@@ -201,7 +215,36 @@ class Aria2Downloader:
                     completed_path = self._resolve_completed_path(status, save_path)
                     return True, completed_path
                 if state == "error":
-                    return False, status.get("errorMessage") or "aria2 download failed"
+                    error_message = status.get("errorMessage") or "aria2 download failed"
+                    if (
+                        _is_no_uri_available_error(error_message)
+                        and recovery_attempts < MAX_TRANSFER_RECOVERY_ATTEMPTS
+                    ):
+                        # The signed URL (e.g. CivitAI's) expired before the
+                        # transfer finished.  Re-registering resolves a fresh
+                        # URL and resumes from the on-disk partial payload and
+                        # .aria2 control file via ``continue=true``.
+                        recovery_attempts += 1
+                        logger.warning(
+                            "aria2 transfer %s failed with %r; refreshing the "
+                            "URL and resuming the partial download "
+                            "(attempt %d/%d)",
+                            download_id,
+                            error_message,
+                            recovery_attempts,
+                            MAX_TRANSFER_RECOVERY_ATTEMPTS,
+                        )
+                        await asyncio.sleep(1.0)
+                        await self._ensure_process()
+                        async with self._register_lock:
+                            transfer = await self._register_transfer(
+                                url,
+                                save_path,
+                                download_id=download_id,
+                                headers=headers,
+                            )
+                        continue
+                    return False, error_message
                 if state == "removed":
                     return False, "Download was cancelled"
 

py/services/download_manager.py

@@ -717,6 +717,47 @@ class DownloadManager:
                 await asyncio.sleep(delay)
         return False
 
+    @staticmethod
+    def _reconcile_failed_aria2_partial(save_path: str) -> None:
+        """Reconcile on-disk partial state after a failed aria2 transfer.
+
+        The payload and its ``.aria2`` control file form a resumable pair and
+        are preserved together so a retry (with a refreshed URL when needed)
+        can resume via aria2's ``continue=true``.  A control file without its
+        payload cannot resume anything, so the orphan is reported and removed.
+        """
+        control_path = f"{save_path}.aria2"
+        payload_exists = os.path.exists(save_path)
+        control_exists = os.path.exists(control_path)
+
+        if payload_exists and not control_exists:
+            # If the .aria2 control file is missing, aria2 considers the
+            # download complete.  A transient RPC failure may have made us
+            # think the download failed even though the file is fully on disk.
+            # Keep the file so a retry can find it already complete.
+            logger.warning(
+                "aria2 download reported failure but .aria2 file is absent "
+                "for %s — the file is likely complete.  Preserving it for retry.",
+                save_path,
+            )
+        elif payload_exists and control_exists:
+            logger.info(
+                "Preserving aria2 partial download for resume: %s", save_path
+            )
+        elif control_exists:
+            logger.warning(
+                "Orphaned aria2 control file without payload: %s — removing it",
+                control_path,
+            )
+            try:
+                os.remove(control_path)
+            except OSError as exc:
+                logger.warning(
+                    "Failed to remove orphaned aria2 control file %s: %s",
+                    control_path,
+                    exc,
+                )
+
     async def _cleanup_cancelled_download_files(
         self,
         download_id: str,
@@ -1226,6 +1267,24 @@ class DownloadManager:
                     )
                     continue
 
+                if not os.path.exists(save_path) and os.path.exists(control_path):
+                    # A control file without its payload cannot resume
+                    # anything; report it and clean up the orphan.
+                    logger.warning(
+                        "Orphaned aria2 control file without payload for %s: "
+                        "%s — removing it",
+                        download_id,
+                        control_path,
+                    )
+                    try:
+                        os.remove(control_path)
+                    except OSError as exc:
+                        logger.warning(
+                            "Failed to remove orphaned aria2 control file %s: %s",
+                            control_path,
+                            exc,
+                        )
+
                 await self._aria2_state_store.remove(download_id)
 
             self._restored_persisted_downloads = True
@@ -2423,20 +2482,8 @@ class DownloadManager:
                     break
 
                 last_error = result
-                # For aria2: if the .aria2 control file is missing, aria2 considers
-                # the download complete.  A transient RPC failure may have made us
-                # think the download failed even though the file is fully on disk.
-                # Keep the file so a retry can find it already complete.
-                if (
-                    transfer_backend == "aria2"
-                    and os.path.exists(save_path)
-                    and not os.path.exists(f"{save_path}.aria2")
-                ):
-                    logger.warning(
-                        "aria2 download reported failure but .aria2 file is absent "
-                        "for %s — the file is likely complete.  Preserving it for retry.",
-                        save_path,
-                    )
+                if transfer_backend == "aria2":
+                    self._reconcile_failed_aria2_partial(save_path)
                 elif os.path.exists(save_path):
                     try:
                         os.remove(save_path)

tests/services/test_aria2_downloader.py

@@ -926,3 +926,149 @@ async def test_rpc_call_suppresses_error_log_when_log_errors_false(
     debug_records = [r for r in caplog.records if r.levelno == logging.DEBUG]
     assert error_records == []
     assert any("GID x is not found" in r.message for r in debug_records)
+
+
+@pytest.mark.asyncio
+async def test_download_file_refreshes_signed_url_on_no_uri_available(
+    tmp_path, monkeypatch
+):
+    """An expired CivitAI signed URL ("No URI available") is recovered by
+    resolving a fresh URL and re-scheduling with continue=true."""
+    downloader = Aria2Downloader()
+    downloader._rpc_url = "http://127.0.0.1/jsonrpc"
+    downloader._rpc_secret = "secret"
+
+    save_path = tmp_path / "downloads" / "model.safetensors"
+    add_uri_urls = []
+    gids = iter(["gid-1", "gid-2"])
+
+    async def fake_rpc_call(method, params, **_kwargs):
+        if method == "aria2.addUri":
+            add_uri_urls.append(params[0][0])
+            return next(gids)
+        if method == "aria2.tellStatus":
+            gid = params[0]
+            if gid == "gid-1":
+                return {
+                    "gid": "gid-1",
+                    "status": "error",
+                    "errorMessage": "No URI available.",
+                }
+            return {
+                "gid": "gid-2",
+                "status": "complete",
+                "completedLength": "10",
+                "totalLength": "10",
+                "downloadSpeed": "0",
+                "files": [{"path": str(save_path)}],
+            }
+        raise AssertionError(f"Unexpected RPC method: {method}")
+
+    resolve = AsyncMock(
+        side_effect=[
+            "https://signed.example.com/model.safetensors?token=old",
+            "https://signed.example.com/model.safetensors?token=fresh",
+        ]
+    )
+    monkeypatch.setattr(downloader, "_ensure_process", AsyncMock())
+    monkeypatch.setattr(downloader, "_resolve_authenticated_redirect_url", resolve)
+    monkeypatch.setattr(downloader, "_rpc_call", fake_rpc_call)
+    monkeypatch.setattr("py.services.aria2_downloader.asyncio.sleep", AsyncMock())
+
+    success, result = await downloader.download_file(
+        "https://civitai.com/api/download/models/123",
+        str(save_path),
+        download_id="download-1",
+        headers={"Authorization": "Bearer token"},
+    )
+
+    assert success is True
+    assert result == str(save_path)
+    # The second addUri used a freshly resolved signed URL, and both
+    # re-schedules keep continue=true so the partial payload is resumed.
+    assert add_uri_urls == [
+        "https://signed.example.com/model.safetensors?token=old",
+        "https://signed.example.com/model.safetensors?token=fresh",
+    ]
+    assert resolve.await_count == 2
+    assert downloader._transfers == {}
+
+
+@pytest.mark.asyncio
+async def test_download_file_fails_when_no_uri_available_exceeds_recovery_limit(
+    tmp_path, monkeypatch
+):
+    """URL refresh is bounded by MAX_TRANSFER_RECOVERY_ATTEMPTS."""
+    downloader = Aria2Downloader()
+    downloader._rpc_url = "http://127.0.0.1/jsonrpc"
+    downloader._rpc_secret = "secret"
+
+    save_path = tmp_path / "downloads" / "model.safetensors"
+    add_uri_count = {"n": 0}
+
+    async def fake_rpc_call(method, params, **_kwargs):
+        if method == "aria2.addUri":
+            add_uri_count["n"] += 1
+            return f"gid-{add_uri_count['n']}"
+        if method == "aria2.tellStatus":
+            return {
+                "gid": params[0],
+                "status": "error",
+                "errorMessage": "No URI available.",
+            }
+        raise AssertionError(f"Unexpected RPC method: {method}")
+
+    monkeypatch.setattr(downloader, "_ensure_process", AsyncMock())
+    monkeypatch.setattr(downloader, "_rpc_call", fake_rpc_call)
+    monkeypatch.setattr("py.services.aria2_downloader.asyncio.sleep", AsyncMock())
+
+    success, result = await downloader.download_file(
+        "https://example.com/model.safetensors",
+        str(save_path),
+        download_id="download-1",
+    )
+
+    assert success is False
+    assert result == "No URI available."
+    assert add_uri_count["n"] == 1 + MAX_TRANSFER_RECOVERY_ATTEMPTS
+    assert downloader._transfers == {}
+
+
+@pytest.mark.asyncio
+async def test_download_file_does_not_refresh_url_for_other_errors(
+    tmp_path, monkeypatch
+):
+    """Errors other than "No URI available" remain permanent failures."""
+    downloader = Aria2Downloader()
+    downloader._rpc_url = "http://127.0.0.1/jsonrpc"
+    downloader._rpc_secret = "secret"
+
+    save_path = tmp_path / "downloads" / "model.safetensors"
+    add_uri_count = {"n": 0}
+
+    async def fake_rpc_call(method, params, **_kwargs):
+        if method == "aria2.addUri":
+            add_uri_count["n"] += 1
+            return "gid-1"
+        if method == "aria2.tellStatus":
+            return {
+                "gid": "gid-1",
+                "status": "error",
+                "errorMessage": "Download aborted. URI=https://example.com/model.safetensors",
+            }
+        raise AssertionError(f"Unexpected RPC method: {method}")
+
+    monkeypatch.setattr(downloader, "_ensure_process", AsyncMock())
+    monkeypatch.setattr(downloader, "_rpc_call", fake_rpc_call)
+    monkeypatch.setattr("py.services.aria2_downloader.asyncio.sleep", AsyncMock())
+
+    success, result = await downloader.download_file(
+        "https://example.com/model.safetensors",
+        str(save_path),
+        download_id="download-1",
+    )
+
+    assert success is False
+    assert "Download aborted" in result
+    assert add_uri_count["n"] == 1
+    assert downloader._transfers == {}

tests/services/test_download_manager_error.py

@@ -1790,3 +1790,140 @@ async def test_concurrent_downloads_with_same_target_path_do_not_destroy_each_ot
         f"Task B reused task A's exact target path instead of a unique one: "
         f"{downloader_paths}"
     )
+
+
+def test_reconcile_failed_aria2_partial_preserves_resumable_pair(tmp_path):
+    """Payload and .aria2 control file are a resumable pair: keep both."""
+    save_path = tmp_path / "model.safetensors"
+    save_path.write_text("partial")
+    control_path = tmp_path / "model.safetensors.aria2"
+    control_path.write_text("control")
+
+    DownloadManager._reconcile_failed_aria2_partial(str(save_path))
+
+    assert save_path.exists()
+    assert control_path.exists()
+
+
+def test_reconcile_failed_aria2_partial_preserves_likely_complete_file(tmp_path):
+    """Payload without a control file means aria2 finished it: keep it."""
+    save_path = tmp_path / "model.safetensors"
+    save_path.write_text("complete")
+
+    DownloadManager._reconcile_failed_aria2_partial(str(save_path))
+
+    assert save_path.exists()
+    assert not (tmp_path / "model.safetensors.aria2").exists()
+
+
+def test_reconcile_failed_aria2_partial_removes_orphaned_control_file(tmp_path):
+    """A control file without its payload cannot resume: remove the orphan."""
+    save_path = tmp_path / "model.safetensors"
+    control_path = tmp_path / "model.safetensors.aria2"
+    control_path.write_text("control")
+
+    DownloadManager._reconcile_failed_aria2_partial(str(save_path))
+
+    assert not control_path.exists()
+
+
+@pytest.mark.asyncio
+async def test_execute_download_preserves_aria2_partial_pair_on_failure(
+    monkeypatch, tmp_path
+):
+    """A failed aria2 transfer keeps the payload and control file for resume."""
+    manager = DownloadManager()
+    settings = get_settings_manager()
+    settings.settings["download_backend"] = "aria2"
+
+    save_dir = tmp_path / "downloads"
+    save_dir.mkdir()
+    target_path = save_dir / "file.safetensors"
+
+    class DummyMetadata:
+        def __init__(self, path: Path):
+            self.file_path = str(path)
+            self.sha256 = "sha256"
+            self.file_name = path.stem
+            self.preview_url = None
+            self.autov3: Optional[str] = None
+
+        def generate_unique_filename(self, *_args, **_kwargs):
+            return os.path.basename(self.file_path)
+
+        def update_file_info(self, _path):
+            return None
+
+        def to_dict(self):
+            return {"file_path": self.file_path}
+
+    class FailingAria2Downloader:
+        async def download_file(
+            self,
+            url,
+            save_path,
+            *,
+            download_id,
+            progress_callback=None,
+            headers=None,
+        ):
+            Path(save_path).write_text("partial")
+            Path(f"{save_path}.aria2").write_text("control")
+            return False, "No URI available."
+
+    monkeypatch.setattr(
+        download_manager,
+        "get_aria2_downloader",
+        AsyncMock(return_value=FailingAria2Downloader()),
+    )
+
+    result = await manager._execute_download(
+        download_urls=["https://civitai.com/api/download/models/1"],
+        save_dir=str(save_dir),
+        metadata=DummyMetadata(target_path),
+        version_info={"images": []},
+        relative_path="",
+        progress_callback=None,
+        model_type="lora",
+        download_id="download-fail",
+    )
+
+    assert result["success"] is False
+    assert result["error"] == "No URI available."
+    assert target_path.exists()
+    assert (save_dir / "file.safetensors.aria2").exists()
+
+
+@pytest.mark.asyncio
+async def test_restore_persisted_downloads_removes_orphaned_aria2_control_file(
+    monkeypatch, tmp_path
+):
+    """A persisted record whose .aria2 control file lost its payload is
+    dropped, and the orphaned control file is removed."""
+    manager = DownloadManager()
+    download_id = "download-orphan"
+    save_path = tmp_path / "file.safetensors"
+    control_path = tmp_path / "file.safetensors.aria2"
+    control_path.write_text("control")
+
+    await manager._aria2_state_store.upsert(
+        download_id,
+        {
+            "download_id": download_id,
+            "transfer_backend": "aria2",
+            "status": "failed",
+            "save_path": str(save_path),
+            "file_path": str(save_path),
+        },
+    )
+
+    monkeypatch.setattr(
+        download_manager,
+        "get_aria2_downloader",
+        AsyncMock(return_value=SimpleNamespace()),
+    )
+
+    await manager._restore_persisted_downloads()
+
+    assert not control_path.exists()
+    assert await manager._aria2_state_store.get(download_id) is None