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
This commit is contained in:
Will Miao
2026-08-29 11:31:16 +08:00
parent cb18281b14
commit 7a36659a20
4 changed files with 390 additions and 17 deletions
+146
View File
@@ -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 == {}
@@ -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