feat(recipes): skip rate-limited batch-import items and register download 429s (#1085)

Phase 2 of docs/plans/issue-1085-rate-limit-design.md:

- Batch import: items that fail due to vendor rate limiting are now
  SKIPPED with a "re-run the import later" hint instead of FAILED, so a
  transient 429 no longer pollutes failure accounting; the progress
  broadcast carries a rate_limited flag.
- Batch import UI: show a one-time "rate limited — slowing down" toast
  and swap the running status text while rate_limited; i18n keys synced
  to all locales.
- Downloader: download_file / download_to_memory / get_response_headers
  register 429 cooldowns with the RateLimitCoordinator, so subsequent
  API calls queue behind a download-triggered rate-limit window.
This commit is contained in:
Will Miao
2026-08-27 10:08:32 +08:00
parent c2a2048c8b
commit df34efafbc
16 changed files with 273 additions and 9 deletions
+100
View File
@@ -655,3 +655,103 @@ class TestInputValidation:
assert service._validate_local_path("../etc/passwd") is False
assert service._validate_local_path("relative/path.png") is False
assert service._validate_local_path("") is False
class TestRateLimitSkipMapping:
"""#1085: vendor rate limiting must mark items SKIPPED, not FAILED."""
@pytest.fixture
def mock_services(self):
ws_manager = MockWebSocketManager()
persistence_service = MockPersistenceService()
logger = logging.getLogger("test")
return ws_manager, persistence_service, logger
def test_is_rate_limit_error_matching(self):
assert BatchImportService._is_rate_limit_error("Rate limited") is True
assert BatchImportService._is_rate_limit_error(
"Rate limit wait for 'civarchive.com' exceeds the 300s cap"
) is True
assert BatchImportService._is_rate_limit_error("Request rate limited") is True
assert BatchImportService._is_rate_limit_error("No metadata found") is False
assert BatchImportService._is_rate_limit_error(None) is False
assert BatchImportService._is_rate_limit_error("") is False
@pytest.mark.asyncio
async def test_rate_limited_item_becomes_skipped_and_sets_flag(self, mock_services):
ws_manager, persistence_service, logger = mock_services
analysis_service = MockAnalysisService(
{
"https://example.com/limited.png": MockAnalysisResult(
{"error": "Rate limited"}
),
}
)
service = BatchImportService(
analysis_service=analysis_service, # pyright: ignore[reportArgumentType]
persistence_service=persistence_service,
ws_manager=ws_manager,
logger=logger,
)
operation_id = await service.start_batch_import(
recipe_scanner_getter=lambda: SimpleNamespace(),
civitai_client_getter=lambda: SimpleNamespace(),
items=[{"source": "https://example.com/limited.png"}],
)
await asyncio.sleep(0.5)
# The operation may already be cleaned up; inspect the broadcasts.
final = next(
(
b
for b in reversed(ws_manager.broadcasts)
if b.get("type") == "batch_import_progress"
),
None,
)
assert final is not None
assert final["rate_limited"] is True
assert final["skipped"] == 1
assert final["failed"] == 0
item = final["items"][0]
assert item["status"] == "skipped"
assert "re-run the import later" in item["error_message"]
assert service.get_progress(operation_id) is None or True
@pytest.mark.asyncio
async def test_non_rate_limit_error_stays_failed(self, mock_services):
ws_manager, persistence_service, logger = mock_services
analysis_service = MockAnalysisService(
{
"https://example.com/broken.png": MockAnalysisResult(
{"error": "No metadata found"}
),
}
)
service = BatchImportService(
analysis_service=analysis_service, # pyright: ignore[reportArgumentType]
persistence_service=persistence_service,
ws_manager=ws_manager,
logger=logger,
)
await service.start_batch_import(
recipe_scanner_getter=lambda: SimpleNamespace(),
civitai_client_getter=lambda: SimpleNamespace(),
items=[{"source": "https://example.com/broken.png"}],
)
await asyncio.sleep(0.5)
final = next(
(
b
for b in reversed(ws_manager.broadcasts)
if b.get("type") == "batch_import_progress"
),
None,
)
assert final is not None
assert final["rate_limited"] is False
assert final["failed"] == 1
assert final["skipped"] == 0
+64 -1
View File
@@ -188,6 +188,12 @@ class _FakeSession:
assert self._responses, "unexpected extra request"
return self._responses.pop(0)
def get(self, url, headers=None, **kwargs):
return self.request("GET", url, headers=headers, **kwargs)
def head(self, url, headers=None, **kwargs):
return self.request("HEAD", url, headers=headers, **kwargs)
async def close(self):
return None
@@ -310,7 +316,13 @@ def _stub_provider(*, result=None, error=None, exc: Exception | None = None):
return SimpleNamespace(get_model_by_hash=call)
async def test_fallback_does_not_fail_over_to_network_provider_on_429():
async def test_fallback_does_not_fail_over_to_network_provider_on_429(monkeypatch):
# The stub error is not gate_handled, so the retry helper would sleep
# retry_after between attempts; patch it out (the helper's own behavior
# is covered by the double-wait tests below).
monkeypatch.setattr(
"py.services.model_metadata_provider.asyncio.sleep", AsyncMock()
)
civitai = _stub_provider(exc=RateLimitError("limited", retry_after=30))
civarchive = _stub_provider(result={"id": 1}, error=None)
sqlite = _stub_provider(result=None, error="not in archive")
@@ -398,3 +410,54 @@ async def test_retry_helper_keeps_legacy_retry_for_ungated_errors():
await helper.run("civitai_api", failing)
assert calls == 2 # legacy retry behavior unchanged
# ----------------------------------------------------------------------
# Download-path 429 registration (Phase 2)
class _FakeDownloadResponse(_FakeResponse):
async def read(self):
return b"data"
async def test_download_to_memory_429_registers_cooldown(monkeypatch):
_patch_gate_settings(
monkeypatch,
rate_limit_gate_enabled=True,
rate_limit_min_interval_seconds=0.0,
)
downloader = _build_downloader(
[_FakeDownloadResponse(429, headers={"Retry-After": "120"})]
)
success, error, _ = await downloader.download_to_memory(
"https://api.example.com/preview.png"
)
assert success is False
assert "Rate limited" in error
coordinator = await RateLimitCoordinator.get_instance()
remaining = coordinator.remaining_seconds("api.example.com")
assert 110.0 < remaining <= 120.0
async def test_get_response_headers_429_registers_cooldown(monkeypatch):
_patch_gate_settings(
monkeypatch,
rate_limit_gate_enabled=True,
rate_limit_min_interval_seconds=0.0,
)
downloader = _build_downloader(
[_FakeResponse(429, headers={"Retry-After": "60"})]
)
success, error = await downloader.get_response_headers(
"https://api.example.com/model/file.safetensors"
)
assert success is False
assert "rate limited" in error.lower()
coordinator = await RateLimitCoordinator.get_instance()
remaining = coordinator.remaining_seconds("api.example.com")
assert 50.0 < remaining <= 60.0