feat(downloads): return structured 429 rate-limit responses with retry_after

On a CivitAI/CivArchive 429, download-model and download-model-get now
return HTTP 429 with {"reason": "rate_limited", "retry_after": N}
instead of a generic 500 string, and the queue row goes back to
"queued" rather than history as failed — so queue drivers can
auto-pause and retry later instead of burning through the queue.

- new DownloadRateLimitError carrying retry_after/host (opt-in via
  raise_on_rate_limit on Downloader; other call sites keep the legacy
  string behavior)
- fail-fast pre-flight gate in DownloadManager consults
  RateLimitCoordinator before acquiring the semaphore slot: hosts in
  cooldown get an immediate structured 429, no HTTP request attempted
- best-effort 429 detection for the aria2 backend
This commit is contained in:
Will Miao
2026-10-02 21:16:06 +08:00
parent 2a667df98c
commit 034660d8c4
8 changed files with 627 additions and 9 deletions
+115
View File
@@ -5,7 +5,20 @@ from typing import Sequence
import pytest
from py.services.connectivity_guard import ConnectivityGuard
from py.services.downloader import Downloader
from py.services.errors import DownloadRateLimitError
from py.services.rate_limit_coordinator import RateLimitCoordinator
@pytest.fixture(autouse=True)
def _reset_rate_limit_singletons():
"""429 handling mutates the global coordinator/guard; isolate per test."""
RateLimitCoordinator._instance = None
ConnectivityGuard._instance = None
yield
RateLimitCoordinator._instance = None
ConnectivityGuard._instance = None
class FakeStream:
@@ -297,3 +310,105 @@ async def test_disable_netrc_auth_ignores_netrc_file(tmp_path, monkeypatch):
_disable_netrc_auth(session)
assert session._get_netrc_auth("civitai.red") is None
@pytest.mark.asyncio
async def test_download_file_rate_limit_raises_structured_error_when_enabled(tmp_path):
"""With raise_on_rate_limit=True a 429 raises DownloadRateLimitError."""
target_path = tmp_path / "model" / "file.bin"
target_path.parent.mkdir()
responses = [
lambda: FakeResponse(
status=429,
headers={"Retry-After": "120"},
chunks=[],
)
]
downloader = _build_downloader(responses)
with pytest.raises(DownloadRateLimitError) as exc_info:
await downloader.download_file(
"https://example.com/file",
str(target_path),
raise_on_rate_limit=True,
)
assert exc_info.value.retry_after == 120.0
assert exc_info.value.host == "example.com"
# The cooldown was registered with the coordinator for the target host.
coordinator = await RateLimitCoordinator.get_instance()
assert coordinator.remaining_seconds("example.com") > 0
# 429 is never retried in-band.
assert _session(downloader)._get_calls == 1
assert not Path(str(target_path) + ".part").exists()
@pytest.mark.asyncio
async def test_download_file_rate_limit_keeps_legacy_tuple_by_default(tmp_path):
"""Default behavior (other callers) stays the plain error string."""
target_path = tmp_path / "model" / "file.bin"
target_path.parent.mkdir()
responses = [
lambda: FakeResponse(
status=429,
headers={"Retry-After": "120"},
chunks=[],
)
]
downloader = _build_downloader(responses)
success, message = await downloader.download_file(
"https://example.com/file", str(target_path)
)
assert success is False
assert message == "Download rate limited (429), retry after 120.0s"
@pytest.mark.asyncio
async def test_download_to_memory_rate_limit_raises_structured_error_when_enabled():
responses = [
lambda: FakeResponse(
status=429,
headers={"Retry-After": "90"},
chunks=[],
)
]
downloader = _build_downloader(responses)
with pytest.raises(DownloadRateLimitError) as exc_info:
await downloader.download_to_memory(
"https://example.com/preview.png",
raise_on_rate_limit=True,
)
assert exc_info.value.retry_after == 90
assert exc_info.value.host == "example.com"
coordinator = await RateLimitCoordinator.get_instance()
assert coordinator.remaining_seconds("example.com") > 0
@pytest.mark.asyncio
async def test_download_to_memory_rate_limit_keeps_legacy_tuple_by_default():
responses = [
lambda: FakeResponse(
status=429,
headers={"Retry-After": "90"},
chunks=[],
)
]
downloader = _build_downloader(responses)
success, message, headers = await downloader.download_to_memory(
"https://example.com/preview.png"
)
assert success is False
assert message == "Rate limited (429), retry after 90s"
assert headers is None