From 7a36659a200dbe3c914d19eb321a4858af392411 Mon Sep 17 00:00:00 2001 From: Will Miao Date: Sat, 29 Aug 2026 11:31:16 +0800 Subject: [PATCH] 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 | 49 +++++- py/services/download_manager.py | 75 +++++++-- tests/services/test_aria2_downloader.py | 146 ++++++++++++++++++ tests/services/test_download_manager_error.py | 137 ++++++++++++++++ 4 files changed, 390 insertions(+), 17 deletions(-) diff --git a/py/services/aria2_downloader.py b/py/services/aria2_downloader.py index 9acb7aab..d4f6a0ec 100644 --- a/py/services/aria2_downloader.py +++ b/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" diff --git a/py/services/download_manager.py b/py/services/download_manager.py index 14e63207..0cdd760d 100644 --- a/py/services/download_manager.py +++ b/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) diff --git a/tests/services/test_aria2_downloader.py b/tests/services/test_aria2_downloader.py index 67665a55..453ff611 100644 --- a/tests/services/test_aria2_downloader.py +++ b/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 == {} diff --git a/tests/services/test_download_manager_error.py b/tests/services/test_download_manager_error.py index 8cd2d7f3..25872dc0 100644 --- a/tests/services/test_download_manager_error.py +++ b/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