From c2a2048c8b4538b836a81c1db14af7caf1e60ca4 Mon Sep 17 00:00:00 2001 From: Will Miao Date: Thu, 27 Aug 2026 09:53:07 +0800 Subject: [PATCH] feat(services): add per-destination rate-limit gate for API traffic (#1085) Implement Phase 1 of docs/plans/issue-1085-rate-limit-design.md: - New RateLimitCoordinator: per-host shared Retry-After gate with exponential backoff (30s base, 1800s cap), minimum inter-request pacing (default 0.75s), herd-free waiter serialization via per-destination locks, and a bounded wait (default 300s) that raises instead of parking. - Downloader.make_request: connectivity-guard fail-fast first, then gate pacing; on 429 register the cooldown and wait-and-resend (bounded); errors that passed through the gate are marked gate_handled. - FallbackMetadataProvider / MetadataSyncService: a network provider 429 no longer fails over to other network providers (stops the CivArchive flood); sqlite stays as local last resort. Rate-limited lookups now report "Rate limited" instead of "Model not found", so transient 429s no longer mark models civitai_deleted. - _RateLimitRetryHelper skips its own sleep for gate_handled errors, removing the double wait. - New settings: rate_limit_gate_enabled, rate_limit_max_wait_seconds, rate_limit_min_interval_seconds. --- docs/plans/issue-1085-rate-limit-design.md | 93 ++-- py/services/downloader.py | 154 ++++--- py/services/metadata_sync_service.py | 9 +- py/services/model_metadata_provider.py | 65 ++- py/services/rate_limit_coordinator.py | 213 ++++++++++ py/services/settings_manager.py | 3 + tests/services/test_metadata_sync_service.py | 53 +++ .../services/test_model_metadata_provider.py | 23 +- tests/services/test_rate_limit_coordinator.py | 400 ++++++++++++++++++ 9 files changed, 909 insertions(+), 104 deletions(-) create mode 100644 py/services/rate_limit_coordinator.py create mode 100644 tests/services/test_rate_limit_coordinator.py diff --git a/docs/plans/issue-1085-rate-limit-design.md b/docs/plans/issue-1085-rate-limit-design.md index 738097aa..46393d87 100644 --- a/docs/plans/issue-1085-rate-limit-design.md +++ b/docs/plans/issue-1085-rate-limit-design.md @@ -1,7 +1,11 @@ # Plan: Global Rate-Limit Abidance for Recipe Ingest & Metadata Fetching **Issue:** [#1085 — Large Recipe Ingest Appears to not abide by vendor rate limits, possibly a few other errors?](https://github.com/willmiao/ComfyUI-Lora-Manager/issues/1085) -**Status:** v1 — awaiting review +**Status:** v2 — reviewed; decisions recorded in §10. **Phase 1 implemented** +(2026-08-27): coordinator + downloader gate + Fix C failover semantics + +helper double-wait fix + settings; full regression 2385 passed. Changes vs v1: +Fix C moved to Phase 1, helper double-wait resolved in Phase 1, gate/guard +ordering specified, WebSocket slowdown hint confirmed in scope (Phase 2). **Scope:** HTTP API traffic to CivitAI (`civitai.red`) and CivArchive (`civarchive.com`) from metadata fetching (bulk refresh, metadata sync, recipe analysis/enrichment, usage-control lookups). Large binary downloads (model files / preview images via `download_file`) are out of scope for *pacing* (they are already single-connection transfers) but their 429 responses should still be *registered*. > Context: a first batch of fixes for this issue was already committed as @@ -187,12 +191,15 @@ API: Enforcement points: -1. **`Downloader.make_request`** (`downloader.py:1102-1132`): before - `session.request`, `await coordinator.wait_for_slot(destination)`. On 429: - `coordinator.register_rate_limit(...)`, then *wait for the gate and - re-send* (loop, bounded by `rate_limit_max_wait_seconds`, default 300; - `retry_after ≥ cap` ⇒ fail immediately). After the loop, return the - `RateLimitError` to the caller (unchanged contract). 200 path calls +1. **`Downloader.make_request`** (`downloader.py:1102-1132`): ordering inside + the method is **connectivity-guard fail-fast first** (offline short-circuit + costs nothing to check), **then** `await coordinator.wait_for_slot(destination)` + before `session.request`. On 429: `coordinator.register_rate_limit(...)`, + then *wait for the gate and re-send* (loop, bounded by + `rate_limit_max_wait_seconds`, default 300; `retry_after ≥ cap` ⇒ fail + immediately). After the loop, return the `RateLimitError` to the caller + (unchanged contract) **with `exc.gate_handled = True` set** so downstream + retry helpers know the wait already happened. 200 path calls `register_success`. 2. **`Downloader.download_to_memory` / `get_response_headers`** (phase 2): register 429s (so API calls queue); waiting only in `make_request` @@ -211,10 +218,12 @@ Enforcement points: `"rate limited (retry_after=…s); re-run the import later"` instead of `FAILED`, and surface a `rate_limited` flag in the WebSocket progress broadcast. -6. **`_RateLimitRetryHelper` retries** (`model_metadata_provider.py`): with the - gate at the downloader, the helper's `retry_after` sleeps become mostly - redundant; demote/simplify in phase 2 (keep the wiring so a - `RateLimitError` still propagates cleanly). +6. **`_RateLimitRetryHelper` retries** (`model_metadata_provider.py`): + **Phase 1** — when the raised `RateLimitError` carries `gate_handled = True` + (set by the downloader after honoring the gate), the helper skips its own + `retry_after` sleep and re-raises immediately, eliminating the double wait. + The wiring stays so a `RateLimitError` still propagates cleanly; full + demotion/removal can follow once the gate proves out. Settings (`settings.json`, schema extension in `SettingsManager`): @@ -230,7 +239,7 @@ Settings (`settings.json`, schema extension in `SettingsManager`): |---|---| | `py/services/rate_limit_coordinator.py` (new) | coordinator singleton + per-destination state + tests seam | | `py/services/downloader.py` | gate pre-check + 429 register/wait/retry loop + `register_success`; log the 429 notice at INFO once per cooldown, then DEBUG | -| `py/services/model_metadata_provider.py` | `FallbackMetadataProvider`: stop network failover on `RateLimitError`; helper simplification | +| `py/services/model_metadata_provider.py` | `FallbackMetadataProvider`: stop network failover on `RateLimitError`; helper skips its sleep when the error is marked `gate_handled` | | `py/services/metadata_sync_service.py` | `fetch_and_update_model`: same failover semantics; keep sqlite last resort | | `py/services/batch_import_service.py` | (phase 2) rate-limit failures → `SKIPPED` + `rate_limited` progress flag | | `py/services/settings_manager.py` | new settings keys + defaults | @@ -245,18 +254,24 @@ Settings (`settings.json`, schema extension in `SettingsManager`): up to the wait cap — UI actions that call the API (e.g. a model-details fetch) may take longer during cooldowns. Mitigation: bounded cap + INFO log + the existing async request handling already tolerates slow responses. - Open question: should interactive (non-batch) requests skip the wait and - fail fast instead? (recommend: same wait — simpler, and cooldowns are short) + **Decided (§10): interactive requests take the same bounded wait** — one + behavior, no call-source plumbing; cooldowns are usually short. +- **Gate waits occupy batch slots**: with the 1–5 batch semaphore, all slots + can park on a gate simultaneously, freezing visible progress for up to one + wait cap per wave. Bounded and acceptable; the phase-2 `SKIPPED` mapping + + WebSocket `rate_limited` flag (both confirmed in scope, §10) make the stall + visible and recoverable. - **Rate limit reality check**: CivitAI anonymous vs keyed limits, and whether `civitai.red` differs, is unverified. Default pacing `0.75 s/req` is a conservative guess (R6). Open question for maintainer: preferred default and whether an API-keyed ceiling should be higher. - **Long CivArchive windows**: `Retry-After ~1500 s` observed in code - comments. A 300 s default cap means such lookups fail rather than wait. - Open question: raise the default cap, or accept failure+skip semantics? -- **Double waiting**: `_RateLimitRetryHelper` + gate could stack waits; the - phase-2 simplification removes the helper's own sleeps for requests that go - through the downloader. + comments. **Decided (§10): keep the 300 s default cap** — such lookups + fail/skip rather than park a request path for 25 minutes; batch import maps + them to `SKIPPED` (phase 2) so the user can re-run later. +- **Double waiting**: `_RateLimitRetryHelper` + gate could stack waits. + **Resolved in Phase 1**: the downloader marks gate-honored errors with + `gate_handled = True` and the helper skips its own sleep for those. - **Downloads**: `download_file` 429s return an error to download managers unchanged (already handled); only *registration* is proposed, so future API calls queue behind a large `Retry-After` from a download burst. @@ -289,22 +304,30 @@ Settings (`settings.json`, schema extension in `SettingsManager`): ## 9. Implementation Phases -- **Phase 1 (this plan, after review):** `RateLimitCoordinator` + - `Downloader.make_request` integration (pre-check pacing + 429 - register/wait/retry loop + cap) + settings + coordinator/downloader tests. -- **Phase 2:** failover semantics (`FallbackMetadataProvider`, - `fetch_and_update_model`), helper simplification, batch-import - `SKIPPED`-on-rate-limit + progress flag, `download_to_memory`/HEAD 429 - registration, provider/sync/batch tests. +- **Phase 1 (this plan, reviewed):** `RateLimitCoordinator` + + `Downloader.make_request` integration (guard fail-fast → gate pre-check + pacing → 429 register/wait/retry loop with cap → `gate_handled` marking) + + settings + **Fix C failover semantics** (`FallbackMetadataProvider`, + `fetch_and_update_model` — moved up from phase 2: smallest diff, kills the + CivArchive flood immediately, independent of coordinator correctness) + + `_RateLimitRetryHelper` double-wait fix + coordinator/downloader/provider/ + sync tests. +- **Phase 2:** batch-import `SKIPPED`-on-rate-limit + `rate_limited` WebSocket + progress flag + slowdown hint (confirmed, §10), + `download_to_memory`/HEAD 429 registration, batch tests. - **Phase 3:** full regression + docs + commit referencing `(#1085)`. -## 10. Review Checklist +## 10. Review Checklist — Decisions (2026-08-27) -- [ ] Default pacing interval acceptable (`0.75 s`)? Prefer higher/lower? -- [ ] Wait cap `300 s` acceptable, or should long-window CivArchive lookups - wait longer? -- [ ] OK that interactive API calls also wait (bounded) instead of failing - fast? -- [ ] Keep sqlite as last resort behind a network rate limit? -- [ ] Add a UI hint ("rate limited — slowing down") surfaced via WebSocket, - or is INFO logging enough? \ No newline at end of file +- [x] Default pacing interval `0.75 s` — **accepted** as conservative default; + tunable via `rate_limit_min_interval_seconds`. Revisit if CivitAI + publishes keyed/anonymous ceilings. +- [x] Wait cap `300 s` — **accepted**; long-window CivArchive lookups fail → + batch import marks them `SKIPPED` with a rate-limit reason (phase 2). +- [x] Interactive API calls also wait (bounded) — **yes**, same behavior for + all callers. +- [x] Keep sqlite as last resort behind a network rate limit — **yes** + (local-only, no vendor cost). +- [x] UI hint — **yes**: WebSocket `rate_limited` flag + "rate limited — + slowing down" hint in batch-import progress (phase 2); INFO logging + regardless. \ No newline at end of file diff --git a/py/services/downloader.py b/py/services/downloader.py index 2697ba26..16211942 100644 --- a/py/services/downloader.py +++ b/py/services/downloader.py @@ -32,6 +32,7 @@ from .connectivity_guard import ( ConnectivityGuard, ) from .errors import RateLimitError +from .rate_limit_coordinator import RateLimitCoordinator logger = logging.getLogger(__name__) @@ -1074,74 +1075,113 @@ class Downloader: Returns: Tuple[bool, Union[Dict, str]]: (success, response data or error message) + + When the rate-limit gate is enabled (``rate_limit_gate_enabled``), + requests are paced per destination and 429 responses are honored by + waiting out the ``Retry-After`` window (bounded by + ``rate_limit_max_wait_seconds``) before re-sending. A ``RateLimitError`` + returned after gate involvement is marked with ``gate_handled = True`` + so downstream retry helpers do not wait a second time. """ guard = await ConnectivityGuard.get_instance() destination = self._guard_destination(url) + # Fail fast on transport-level outages before pacing: there is no + # point waiting out a vendor cooldown while the network is down. if guard.should_block_request(destination): return False, OFFLINE_COOLDOWN_ERROR - try: - session = await self.session - # Debug log for proxy mode at request time - if self.proxy_url: - logger.debug(f"[make_request] Using app-level proxy: {self.proxy_url}") - else: - logger.debug( - "[make_request] Using system-level proxy (trust_env) if configured." - ) + coordinator = await RateLimitCoordinator.get_instance() + gate_enabled = coordinator.enabled + # Safety bound on the wait-and-resend loop; each 429 normally exits + # via the wait cap in wait_for_slot, this covers pathological 429s + # with tiny Retry-After values. + max_resend_attempts = 5 + attempt = 0 - # Prepare headers - headers = self._get_auth_headers(use_auth) - if custom_headers: - headers.update(custom_headers) + while True: + if gate_enabled: + try: + await coordinator.wait_for_slot(destination) + except RateLimitError as exc: + exc.gate_handled = True + return False, exc - # Add proxy to kwargs if not already present - if "proxy" not in kwargs: - kwargs["proxy"] = self.proxy_url - - async with session.request( - method, url, headers=headers, **kwargs - ) as response: - if response.status == 200: - guard.register_success(destination) - # Try to parse as JSON, fall back to text - try: - data = await response.json() - return True, data - except: - text = await response.text() - return True, text - elif response.status == 401: - return False, "Unauthorized access - invalid or missing API key" - elif response.status == 403: - return False, "Access forbidden" - elif response.status == 404: - return False, "Resource not found" - elif response.status == 429: - retry_after = self._extract_retry_after(response.headers) - error_msg = "Request rate limited" - logger.warning( - "Rate limit encountered for %s %s; retry_after=%s", - method, - url, - retry_after, - ) - return False, RateLimitError( - error_msg, - retry_after=retry_after, - ) + try: + session = await self.session + # Debug log for proxy mode at request time + if self.proxy_url: + logger.debug(f"[make_request] Using app-level proxy: {self.proxy_url}") else: - return False, f"Request failed with status {response.status}" + logger.debug( + "[make_request] Using system-level proxy (trust_env) if configured." + ) - except Exception as e: - if guard.is_network_unreachable_error(e): - guard.register_network_failure(e, destination) - if guard.should_block_request(destination): - return False, OFFLINE_COOLDOWN_ERROR - logger.debug("Network unavailable for %s %s: %s", method, url, e) + # Prepare headers + headers = self._get_auth_headers(use_auth) + if custom_headers: + headers.update(custom_headers) + + # Add proxy to kwargs if not already present + if "proxy" not in kwargs: + kwargs["proxy"] = self.proxy_url + + async with session.request( + method, url, headers=headers, **kwargs + ) as response: + if response.status == 200: + guard.register_success(destination) + if gate_enabled: + coordinator.register_success(destination) + # Try to parse as JSON, fall back to text + try: + data = await response.json() + return True, data + except: + text = await response.text() + return True, text + elif response.status == 401: + return False, "Unauthorized access - invalid or missing API key" + elif response.status == 403: + return False, "Access forbidden" + elif response.status == 404: + return False, "Resource not found" + elif response.status == 429: + retry_after = self._extract_retry_after(response.headers) + error_msg = "Request rate limited" + if not gate_enabled: + logger.warning( + "Rate limit encountered for %s %s; retry_after=%s", + method, + url, + retry_after, + ) + return False, RateLimitError( + error_msg, + retry_after=retry_after, + ) + # The coordinator logs the cooldown notice (INFO once + # per window, DEBUG on extension). + coordinator.register_rate_limit(destination, retry_after) + attempt += 1 + if attempt >= max_resend_attempts: + error = RateLimitError(error_msg, retry_after=retry_after) + error.gate_handled = True + return False, error + # Loop back: wait_for_slot blocks until the cooldown + # elapses (or raises once the wait exceeds the cap). + continue + else: + return False, f"Request failed with status {response.status}" + + except Exception as e: + if guard.is_network_unreachable_error(e): + guard.register_network_failure(e, destination) + if guard.should_block_request(destination): + return False, OFFLINE_COOLDOWN_ERROR + logger.debug("Network unavailable for %s %s: %s", method, url, e) + return False, str(e) + logger.error(f"Error making {method} request to {url}: {e}") return False, str(e) - logger.error(f"Error making {method} request to {url}: {e}") - return False, str(e) async def close(self): """Close the HTTP session""" diff --git a/py/services/metadata_sync_service.py b/py/services/metadata_sync_service.py index e441a8cc..6c8088eb 100644 --- a/py/services/metadata_sync_service.py +++ b/py/services/metadata_sync_service.py @@ -245,16 +245,23 @@ class MetadataSyncService: civitai_api_not_found = False any_rate_limited = False + skip_network_providers = False for provider_name, provider in provider_attempts: + if skip_network_providers and provider_name != "sqlite": + # A network provider was already rate-limited; failing + # over to another network provider just spreads the flood + # (#1085). The local sqlite archive stays as last resort. + continue try: civitai_metadata_candidate, error = await provider.get_model_by_hash(sha256) except RateLimitError as exc: logger.warning( - "Provider %s is rate-limited (retry_after=%.0fs); skipping to next provider", + "Provider %s is rate-limited (retry_after=%.0fs); not failing over to other network providers", provider_name or provider.__class__.__name__, exc.retry_after or 0, ) any_rate_limited = True + skip_network_providers = True continue except Exception as exc: # pragma: no cover - defensive logging logger.error("Provider %s failed for hash %s: %s", provider_name, sha256, exc) diff --git a/py/services/model_metadata_provider.py b/py/services/model_metadata_provider.py index 9625f738..a6b6b528 100644 --- a/py/services/model_metadata_provider.py +++ b/py/services/model_metadata_provider.py @@ -66,6 +66,14 @@ class _RateLimitRetryHelper: except RateLimitError as exc: attempt += 1 + # The downloader's rate-limit gate already applied the wait + # policy for this request (waited out the vendor window or + # deliberately refused because it exceeds the cap). Sleeping + # again here would double the wait — just propagate. + if getattr(exc, "gate_handled", False): + exc.provider = exc.provider or label + raise + # Determine effective retry limit based on rate-limit magnitude effective_retry_limit = self._retry_limit # default: 3 if exc.retry_after is not None and exc.retry_after >= 120.0: @@ -101,6 +109,12 @@ class _RateLimitRetryHelper: return min(self._max_delay, max(0.0, base_delay)) + +# Labels of providers that are free to consult even while a network provider +# is rate-limited (local lookups, no vendor cost). +_LOCAL_PROVIDER_LABELS = frozenset({"sqlite"}) + + class ModelMetadataProvider(ABC): """Base abstract class for all model metadata providers""" @@ -451,7 +465,14 @@ class SQLiteModelMetadataProvider(ModelMetadataProvider): return None class FallbackMetadataProvider(ModelMetadataProvider): - """Try providers in order, return first successful result.""" + """Try providers in order, return first successful result. + + Rate-limit policy (#1085): once a *network* provider raises + ``RateLimitError``, the chain stops consulting further network providers — + failing over would just spread the flood to the next vendor. Local-only + providers (see ``_LOCAL_PROVIDER_LABELS``) are still allowed as a last + resort because they cost the vendor nothing. + """ def __init__( self, @@ -486,7 +507,10 @@ class FallbackMetadataProvider(ModelMetadataProvider): ) async def get_model_by_hash(self, model_hash: str) -> Tuple[Optional[Dict[str, Any]], Optional[str]]: + rate_limited = False for provider, label in self._iter_providers(): + if rate_limited and label not in _LOCAL_PROVIDER_LABELS: + continue try: result, error = await self._call_with_rate_limit( label, @@ -496,8 +520,9 @@ class FallbackMetadataProvider(ModelMetadataProvider): if result: return result, error except RateLimitError as exc: + rate_limited = True logger.warning( - "Provider %s is rate-limited (retry_after=%.0fs); skipping to next provider", + "Provider %s is rate-limited (retry_after=%.0fs); not failing over to other network providers", label, exc.retry_after or 0, ) @@ -505,11 +530,18 @@ class FallbackMetadataProvider(ModelMetadataProvider): except Exception as e: logger.debug("Provider %s failed for get_model_by_hash: %s", label, e) continue + if rate_limited: + # Distinct from "Model not found": callers must not mistake a + # rate-limited lookup for a confirmed deletion. + return None, "Rate limited" return None, "Model not found" async def get_model_versions(self, model_id: str) -> Optional[Dict[str, Any]]: not_found_confirmed = False + rate_limited = False for provider, label in self._iter_providers(): + if rate_limited and label not in _LOCAL_PROVIDER_LABELS: + continue try: result = await self._call_with_rate_limit( label, @@ -519,8 +551,9 @@ class FallbackMetadataProvider(ModelMetadataProvider): if result: return result except RateLimitError as exc: + rate_limited = True logger.warning( - "Provider %s is rate-limited (retry_after=%.0fs); skipping to next provider", + "Provider %s is rate-limited (retry_after=%.0fs); not failing over to other network providers", label, exc.retry_after or 0, ) @@ -539,7 +572,10 @@ class FallbackMetadataProvider(ModelMetadataProvider): return None async def get_model_version(self, model_id: Optional[int] = None, version_id: Optional[int] = None) -> Optional[Dict[str, Any]]: + rate_limited = False for provider, label in self._iter_providers(): + if rate_limited and label not in _LOCAL_PROVIDER_LABELS: + continue try: result = await self._call_with_rate_limit( label, @@ -550,8 +586,9 @@ class FallbackMetadataProvider(ModelMetadataProvider): if result: return result except RateLimitError as exc: + rate_limited = True logger.warning( - "Provider %s is rate-limited (retry_after=%.0fs); skipping to next provider", + "Provider %s is rate-limited (retry_after=%.0fs); not failing over to other network providers", label, exc.retry_after or 0, ) @@ -562,7 +599,10 @@ class FallbackMetadataProvider(ModelMetadataProvider): return None async def get_model_version_info(self, version_id: str) -> Tuple[Optional[Dict[str, Any]], Optional[str]]: + rate_limited = False for provider, label in self._iter_providers(): + if rate_limited and label not in _LOCAL_PROVIDER_LABELS: + continue try: result, error = await self._call_with_rate_limit( label, @@ -572,8 +612,9 @@ class FallbackMetadataProvider(ModelMetadataProvider): if result: return result, error except RateLimitError as exc: + rate_limited = True logger.warning( - "Provider %s is rate-limited (retry_after=%.0fs); skipping to next provider", + "Provider %s is rate-limited (retry_after=%.0fs); not failing over to other network providers", label, exc.retry_after or 0, ) @@ -581,12 +622,17 @@ class FallbackMetadataProvider(ModelMetadataProvider): except Exception as e: logger.debug("Provider %s failed for get_model_version_info: %s", label, e) continue + if rate_limited: + return None, "Rate limited" return None, "No provider could retrieve the data" async def get_model_versions_by_hashes( self, hashes: List[str] ) -> Optional[List[Dict[str, Any]]]: + rate_limited = False for provider, label in self._iter_providers(): + if rate_limited and label not in _LOCAL_PROVIDER_LABELS: + continue try: result = await self._call_with_rate_limit( label, @@ -598,8 +644,9 @@ class FallbackMetadataProvider(ModelMetadataProvider): except NotImplementedError: continue except RateLimitError as exc: + rate_limited = True logger.warning( - "Provider %s is rate-limited (retry_after=%.0fs); skipping to next provider", + "Provider %s is rate-limited (retry_after=%.0fs); not failing over to other network providers", label, exc.retry_after or 0, ) @@ -614,7 +661,10 @@ class FallbackMetadataProvider(ModelMetadataProvider): return None async def get_user_models(self, username: str, cursor: Optional[str] = None) -> Optional[Dict[str, Any]]: + rate_limited = False for provider, label in self._iter_providers(): + if rate_limited and label not in _LOCAL_PROVIDER_LABELS: + continue try: result = await self._call_with_rate_limit( label, @@ -625,8 +675,9 @@ class FallbackMetadataProvider(ModelMetadataProvider): if result is not None: return result except RateLimitError as exc: + rate_limited = True logger.warning( - "Provider %s is rate-limited (retry_after=%.0fs); skipping to next provider", + "Provider %s is rate-limited (retry_after=%.0fs); not failing over to other network providers", label, exc.retry_after or 0, ) diff --git a/py/services/rate_limit_coordinator.py b/py/services/rate_limit_coordinator.py new file mode 100644 index 00000000..70c25eac --- /dev/null +++ b/py/services/rate_limit_coordinator.py @@ -0,0 +1,213 @@ +"""Process-wide, per-destination rate-limit gate for outbound API traffic. + +Implements the pacing/gating layer designed in +``docs/plans/issue-1085-rate-limit-design.md``: + +- **Reactive gate**: a 429 response arms ``next_allowed_send`` from the + vendor's ``Retry-After`` (or exponential backoff when the header is + missing); subsequent requests to the same destination wait out the window. +- **Preemptive pacing**: a minimum inter-request interval per destination + spaces consecutive sends so bursts never form in the first place. +- **Herd-free**: waiters are serialized through a per-destination lock, so + each one claims a distinct send slot instead of thousands of coroutines + waking up together. +- **Bounded**: waits longer than ``rate_limit_max_wait_seconds`` are refused + by raising :class:`RateLimitError`, leaving the final decision to callers. +""" + +from __future__ import annotations + +import asyncio +import logging +import time +from dataclasses import dataclass, field +from typing import Dict, Optional + +from .errors import RateLimitError + +logger = logging.getLogger(__name__) + +DEFAULT_MIN_INTERVAL_SECONDS = 0.75 +DEFAULT_MAX_WAIT_SECONDS = 300.0 +BASE_BACKOFF_SECONDS = 30.0 +MAX_BACKOFF_SECONDS = 1800.0 + + +@dataclass +class _DestinationState: + """Rate-limit bookkeeping for one destination (hostname).""" + + next_allowed_send: float = 0.0 # time.monotonic() timestamp + consecutive_429: int = 0 + last_send_at: float = 0.0 # time.monotonic() timestamp + lock: asyncio.Lock = field(default_factory=asyncio.Lock) + + +class RateLimitCoordinator: + """Coordinates outbound request pacing per destination. + + Singleton mirroring :class:`ConnectivityGuard`'s pattern. All waits are + bounded by the ``rate_limit_max_wait_seconds`` setting; when the required + wait exceeds the cap, :meth:`wait_for_slot` raises :class:`RateLimitError` + instead of parking the caller. + """ + + _instance: "RateLimitCoordinator | None" = None + _instance_lock = asyncio.Lock() + + @classmethod + async def get_instance(cls) -> "RateLimitCoordinator": + async with cls._instance_lock: + if cls._instance is None: + cls._instance = cls() + return cls._instance + + def __init__(self) -> None: + if hasattr(self, "_initialized"): + return + self._initialized = True + self._states: Dict[str, _DestinationState] = {} + + # ------------------------------------------------------------------ + # Settings (read live so settings edits apply without a restart) + + @staticmethod + def _setting(key: str, default): + try: + from .settings_manager import get_settings_manager + + return get_settings_manager().get(key, default) + except Exception: # pragma: no cover - defensive: settings unavailable + return default + + @property + def enabled(self) -> bool: + return bool(self._setting("rate_limit_gate_enabled", True)) + + @property + def min_interval_seconds(self) -> float: + try: + return max(0.0, float(self._setting("rate_limit_min_interval_seconds", DEFAULT_MIN_INTERVAL_SECONDS))) + except (TypeError, ValueError): + return DEFAULT_MIN_INTERVAL_SECONDS + + @property + def max_wait_seconds(self) -> float: + try: + return max(0.0, float(self._setting("rate_limit_max_wait_seconds", DEFAULT_MAX_WAIT_SECONDS))) + except (TypeError, ValueError): + return DEFAULT_MAX_WAIT_SECONDS + + # ------------------------------------------------------------------ + # State helpers + + @staticmethod + def _normalize(destination: Optional[str]) -> str: + if destination is None or not destination.strip(): + return "__global__" + return destination.lower().strip() + + def _state_for(self, destination: Optional[str]) -> _DestinationState: + key = self._normalize(destination) + if key not in self._states: + self._states[key] = _DestinationState() + return self._states[key] + + def reset(self) -> None: + """Drop all per-destination state. Test seam.""" + self._states.clear() + + def in_cooldown(self, destination: Optional[str] = None) -> bool: + return self.remaining_seconds(destination) > 0 + + def remaining_seconds(self, destination: Optional[str] = None) -> float: + state = self._state_for(destination) + return max(0.0, state.next_allowed_send - time.monotonic()) + + # ------------------------------------------------------------------ + # Gate operations + + async def wait_for_slot(self, destination: Optional[str] = None) -> None: + """Block until this caller may send the next request to *destination*. + + Waits for both the rate-limit cooldown (``next_allowed_send``) and the + minimum inter-request interval (``last_send_at + min_interval``). + Waiters queue on the per-destination lock, so concurrent callers are + spaced out instead of stampeding when a cooldown expires. + + Raises: + RateLimitError: when the required wait exceeds + ``rate_limit_max_wait_seconds``. + """ + state = self._state_for(destination) + deadline = time.monotonic() + self.max_wait_seconds + async with state.lock: + now = time.monotonic() + wake_at = max( + state.next_allowed_send, + state.last_send_at + self.min_interval_seconds, + ) + if wake_at > deadline: + raise RateLimitError( + f"Rate limit wait for '{self._normalize(destination)}' " + f"exceeds the {self.max_wait_seconds:.0f}s cap", + retry_after=wake_at - now, + ) + delay = wake_at - now + if delay > 0: + logger.debug( + "Rate-limit gate: pacing request to '%s' by %.2fs", + self._normalize(destination), + delay, + ) + await asyncio.sleep(delay) + state.last_send_at = time.monotonic() + + def register_rate_limit( + self, + destination: Optional[str], + retry_after: Optional[float] = None, + ) -> float: + """Record a 429 for *destination* and arm the cooldown window. + + Honors the vendor's ``Retry-After`` when present; otherwise grows an + exponential backoff (30s base, doubling per consecutive 429, capped at + 1800s). Returns the cooldown duration in seconds. + """ + state = self._state_for(destination) + state.consecutive_429 += 1 + if retry_after is not None and retry_after > 0: + backoff = min(MAX_BACKOFF_SECONDS, float(retry_after)) + else: + backoff = min( + MAX_BACKOFF_SECONDS, + BASE_BACKOFF_SECONDS * (2 ** (state.consecutive_429 - 1)), + ) + now = time.monotonic() + already_cooling = state.next_allowed_send > now + state.next_allowed_send = max(state.next_allowed_send, now + backoff) + if already_cooling: + logger.debug( + "Rate-limit cooldown for '%s' extended by %.0fs (consecutive_429=%d)", + self._normalize(destination), + backoff, + state.consecutive_429, + ) + else: + logger.info( + "Rate limited by '%s'; pausing requests for %.0fs", + self._normalize(destination), + backoff, + ) + return backoff + + def register_success(self, destination: Optional[str]) -> None: + """Reset rate-limit state after a successful request. + + A 200 proves the vendor is accepting traffic again, so any armed + cooldown window is cleared alongside the backoff counter (mirrors + ``ConnectivityGuard.register_success`` semantics). + """ + state = self._state_for(destination) + state.consecutive_429 = 0 + state.next_allowed_send = 0.0 diff --git a/py/services/settings_manager.py b/py/services/settings_manager.py index 2b33fc63..8d46539c 100644 --- a/py/services/settings_manager.py +++ b/py/services/settings_manager.py @@ -70,6 +70,9 @@ DEFAULT_SETTINGS: Dict[str, Any] = { "enable_metadata_archive_db": False, "enable_civarchive_api": True, "metadata_provider_order": "civitai_archive_sqlite", + "rate_limit_gate_enabled": True, + "rate_limit_max_wait_seconds": 300, + "rate_limit_min_interval_seconds": 0.75, "proxy_enabled": False, "proxy_host": "", "proxy_port": "", diff --git a/tests/services/test_metadata_sync_service.py b/tests/services/test_metadata_sync_service.py index 813b5deb..84f45bda 100644 --- a/tests/services/test_metadata_sync_service.py +++ b/tests/services/test_metadata_sync_service.py @@ -812,3 +812,56 @@ async def test_fetch_and_update_model_does_not_overwrite_api_metadata_with_archi helpers.metadata_manager.save_metadata.assert_awaited() update_cache.assert_awaited() + + +@pytest.mark.asyncio +async def test_fetch_and_update_model_keeps_sqlite_last_resort_after_civarchive_rate_limit(tmp_path): + """A CivArchive 429 must not block the local sqlite last resort (#1085).""" + civarchive_provider = SimpleNamespace( + get_model_by_hash=AsyncMock( + side_effect=RateLimitError("limited", retry_after=30) + ), + get_model_version=AsyncMock(), + ) + sqlite_payload = { + "source": "archive_db", + "model": {"name": "Recovered", "description": "", "tags": []}, + "images": [], + "baseModel": "sdxl", + } + sqlite_provider = SimpleNamespace( + get_model_by_hash=AsyncMock(return_value=(sqlite_payload, None)), + get_model_version=AsyncMock(), + ) + + async def select_provider(name: str): + if name == "civarchive_api": + return civarchive_provider + if name == "sqlite": + return sqlite_provider + raise AssertionError(f"unexpected provider request: {name}") + + helpers = build_service( + settings_values={"enable_metadata_archive_db": True}, + provider_selector=AsyncMock(side_effect=select_provider), + ) + + model_path = tmp_path / "model.safetensors" + model_data = { + "civitai_deleted": True, + "db_checked": False, + "file_path": str(model_path), + } + update_cache = AsyncMock() + + ok, error = await helpers.service.fetch_and_update_model( + sha256="cafe", + file_path=str(model_path), + model_data=model_data, + update_cache_func=update_cache, + ) + + assert ok and error is None + civarchive_provider.get_model_by_hash.assert_awaited_once() + sqlite_provider.get_model_by_hash.assert_awaited_once() + assert model_data["metadata_source"] == "archive_db" diff --git a/tests/services/test_model_metadata_provider.py b/tests/services/test_model_metadata_provider.py index c8877ccf..bbde1f9e 100644 --- a/tests/services/test_model_metadata_provider.py +++ b/tests/services/test_model_metadata_provider.py @@ -101,7 +101,9 @@ async def test_fallback_retries_same_provider_on_rate_limit(monkeypatch): @pytest.mark.asyncio async def test_fallback_continues_to_next_provider_on_rate_limit(monkeypatch): - """After exhausting retries on primary, fallback should continue to secondary.""" + """#1085: a rate-limited network provider no longer fails over to another + network provider (that just spreads the flood); local providers such as + sqlite remain as a last resort.""" sleep_mock = AsyncMock() monkeypatch.setattr(provider_module.asyncio, "sleep", sleep_mock) monkeypatch.setattr(provider_module.random, "uniform", lambda *_: 0.0) @@ -114,13 +116,26 @@ async def test_fallback_continues_to_next_provider_on_rate_limit(monkeypatch): rate_limit_retry_limit=2, ) - # After Change A: no longer raises; falls through to secondary + result, error = await fallback.get_model_by_hash("abc") + + # Secondary is a network provider: it must NOT be consulted after the 429. + assert result is None + assert error == "Rate limited" + assert primary.calls == 2 # retry_limit exhausted on primary + assert secondary.calls == 0 # no network failover + + # A local sqlite provider behind the rate-limited one is still allowed. + sqlite = TrackingProvider() + fallback = FallbackMetadataProvider( + [("primary", AlwaysRateLimitedProvider()), ("sqlite", sqlite)], + rate_limit_retry_limit=2, + ) + result, error = await fallback.get_model_by_hash("abc") assert error is None assert result == {"id": "secondary"} - assert primary.calls == 2 # retry_limit exhausted on primary - assert secondary.calls == 1 # secondary IS called now + assert sqlite.calls == 1 @pytest.mark.asyncio diff --git a/tests/services/test_rate_limit_coordinator.py b/tests/services/test_rate_limit_coordinator.py new file mode 100644 index 00000000..a557c3d8 --- /dev/null +++ b/tests/services/test_rate_limit_coordinator.py @@ -0,0 +1,400 @@ +"""Tests for the per-destination rate-limit gate (#1085). + +Covers the RateLimitCoordinator itself, its integration into +``Downloader.make_request``, the failover semantics change in +``FallbackMetadataProvider``, and the ``_RateLimitRetryHelper`` double-wait +fix. +""" + +from __future__ import annotations + +import asyncio +import time +from datetime import datetime +from types import SimpleNamespace +from typing import Any, Dict, Optional +from unittest.mock import AsyncMock + +import pytest + +from py.services.connectivity_guard import ConnectivityGuard +from py.services.downloader import Downloader +from py.services.errors import RateLimitError +from py.services.model_metadata_provider import ( + FallbackMetadataProvider, + _RateLimitRetryHelper, +) +from py.services.rate_limit_coordinator import RateLimitCoordinator + + +@pytest.fixture(autouse=True) +def _reset_singletons(): + RateLimitCoordinator._instance = None + ConnectivityGuard._instance = None + yield + RateLimitCoordinator._instance = None + ConnectivityGuard._instance = None + + +def _patch_gate_settings(monkeypatch, **overrides): + """Override the coordinator's settings reads for the test.""" + monkeypatch.setattr( + RateLimitCoordinator, + "_setting", + staticmethod(lambda key, default: overrides.get(key, default)), + ) + + +async def _make_coordinator(monkeypatch, **overrides) -> RateLimitCoordinator: + _patch_gate_settings(monkeypatch, **overrides) + return await RateLimitCoordinator.get_instance() + + +# ---------------------------------------------------------------------- +# Coordinator unit tests + + +async def test_pacing_enforces_min_interval(monkeypatch): + coordinator = await _make_coordinator( + monkeypatch, rate_limit_min_interval_seconds=0.1 + ) + start = time.monotonic() + await coordinator.wait_for_slot("example.com") + await coordinator.wait_for_slot("example.com") + elapsed = time.monotonic() - start + assert elapsed >= 0.1 + + +async def test_pacing_is_per_destination(monkeypatch): + coordinator = await _make_coordinator( + monkeypatch, rate_limit_min_interval_seconds=0.2 + ) + await coordinator.wait_for_slot("a.example.com") + start = time.monotonic() + await coordinator.wait_for_slot("b.example.com") + elapsed = time.monotonic() - start + assert elapsed < 0.1 + + +async def test_register_rate_limit_arms_cooldown_and_waits(monkeypatch): + coordinator = await _make_coordinator( + monkeypatch, + rate_limit_min_interval_seconds=0.0, + rate_limit_max_wait_seconds=5.0, + ) + coordinator.register_rate_limit("example.com", retry_after=0.15) + assert coordinator.in_cooldown("example.com") + assert 0.1 < coordinator.remaining_seconds("example.com") <= 0.15 + + start = time.monotonic() + await coordinator.wait_for_slot("example.com") + elapsed = time.monotonic() - start + assert elapsed >= 0.14 + assert not coordinator.in_cooldown("example.com") + + +async def test_concurrent_waiters_share_one_cooldown_window(monkeypatch): + """Herd test: N waiters wake after ~one window, not N windows.""" + coordinator = await _make_coordinator( + monkeypatch, + rate_limit_min_interval_seconds=0.0, + rate_limit_max_wait_seconds=5.0, + ) + coordinator.register_rate_limit("example.com", retry_after=0.2) + + start = time.monotonic() + await asyncio.gather( + *(coordinator.wait_for_slot("example.com") for _ in range(4)) + ) + elapsed = time.monotonic() - start + # 4 independent windows would take ~0.8s; a shared window is ~0.2s. + assert 0.19 <= elapsed < 0.5 + + +async def test_backoff_grows_on_consecutive_429_and_resets_on_success( + monkeypatch, +): + coordinator = await _make_coordinator( + monkeypatch, rate_limit_min_interval_seconds=0.0 + ) + coordinator.register_rate_limit("example.com", retry_after=None) + first = coordinator.remaining_seconds("example.com") + assert 29.0 < first <= 30.0 + + coordinator.register_rate_limit("example.com", retry_after=None) + second = coordinator.remaining_seconds("example.com") + assert 59.0 < second <= 60.0 + + coordinator.register_success("example.com") + coordinator.register_rate_limit("example.com", retry_after=None) + third = coordinator.remaining_seconds("example.com") + assert 29.0 < third <= 30.0 + + +async def test_wait_beyond_cap_raises_rate_limit_error(monkeypatch): + coordinator = await _make_coordinator( + monkeypatch, + rate_limit_min_interval_seconds=0.0, + rate_limit_max_wait_seconds=0.05, + ) + coordinator.register_rate_limit("example.com", retry_after=30.0) + + start = time.monotonic() + with pytest.raises(RateLimitError) as excinfo: + await coordinator.wait_for_slot("example.com") + elapsed = time.monotonic() - start + assert elapsed < 1.0 # refused immediately instead of parking + assert excinfo.value.retry_after is not None + assert excinfo.value.retry_after > 1.0 + + +# ---------------------------------------------------------------------- +# Downloader integration tests + + +class _FakeResponse: + def __init__( + self, + status: int, + payload: Optional[Dict[str, Any]] = None, + headers: Optional[Dict[str, str]] = None, + ): + self.status = status + self._payload = payload + self.headers = headers or {} + + async def __aenter__(self): + return self + + async def __aexit__(self, exc_type, exc, tb): + return False + + async def json(self): + if self._payload is None: + raise ValueError("no json payload") + return self._payload + + async def text(self): + return "" + + +class _FakeSession: + def __init__(self, responses): + self._responses = list(responses) + self.requests = [] + + def request(self, method, url, headers=None, **kwargs): + self.requests.append({"method": method, "url": url}) + assert self._responses, "unexpected extra request" + return self._responses.pop(0) + + async def close(self): + return None + + +def _build_downloader(responses) -> Downloader: + downloader = Downloader() + fake_session = _FakeSession(responses) + downloader._session = fake_session # pyright: ignore[reportAttributeAccessIssue] + downloader._session_created_at = datetime.now() + downloader._proxy_url = None + + async def _noop_create_session(): + downloader._session = fake_session # pyright: ignore[reportAttributeAccessIssue] + downloader._session_created_at = datetime.now() + downloader._proxy_url = None + + downloader._create_session = _noop_create_session # type: ignore[assignment] + return downloader + + +async def test_make_request_waits_out_429_then_resends(monkeypatch): + _patch_gate_settings( + monkeypatch, + rate_limit_gate_enabled=True, + rate_limit_min_interval_seconds=0.0, + rate_limit_max_wait_seconds=5.0, + ) + downloader = _build_downloader( + [ + _FakeResponse(429, headers={"Retry-After": "1"}), + _FakeResponse(200, payload={"ok": True}), + ] + ) + + start = time.monotonic() + success, payload = await downloader.make_request( + "GET", "https://api.example.com/models/1" + ) + elapsed = time.monotonic() - start + + assert success is True + assert payload == {"ok": True} + assert len(downloader._session.requests) == 2 + assert elapsed >= 0.9 + + +async def test_make_request_paces_consecutive_calls(monkeypatch): + _patch_gate_settings( + monkeypatch, + rate_limit_gate_enabled=True, + rate_limit_min_interval_seconds=0.15, + rate_limit_max_wait_seconds=5.0, + ) + downloader = _build_downloader( + [_FakeResponse(200, payload={}), _FakeResponse(200, payload={})] + ) + + start = time.monotonic() + await downloader.make_request("GET", "https://api.example.com/a") + await downloader.make_request("GET", "https://api.example.com/b") + elapsed = time.monotonic() - start + assert elapsed >= 0.14 + + +async def test_make_request_gate_disabled_returns_429_immediately(monkeypatch): + _patch_gate_settings(monkeypatch, rate_limit_gate_enabled=False) + downloader = _build_downloader( + [_FakeResponse(429, headers={"Retry-After": "30"})] + ) + + start = time.monotonic() + success, payload = await downloader.make_request( + "GET", "https://api.example.com/models/1" + ) + elapsed = time.monotonic() - start + + assert success is False + assert isinstance(payload, RateLimitError) + assert payload.retry_after == 30.0 + # Gate was off: the error is NOT marked, so retry helpers keep their + # legacy behavior. + assert getattr(payload, "gate_handled", False) is False + assert len(downloader._session.requests) == 1 + assert elapsed < 1.0 + + +async def test_make_request_refuses_wait_beyond_cap(monkeypatch): + _patch_gate_settings( + monkeypatch, + rate_limit_gate_enabled=True, + rate_limit_min_interval_seconds=0.0, + rate_limit_max_wait_seconds=0.2, + ) + downloader = _build_downloader( + [_FakeResponse(429, headers={"Retry-After": "3600"})] + ) + + start = time.monotonic() + success, payload = await downloader.make_request( + "GET", "https://api.example.com/models/1" + ) + elapsed = time.monotonic() - start + + assert success is False + assert isinstance(payload, RateLimitError) + assert payload.gate_handled is True + assert len(downloader._session.requests) == 1 + assert elapsed < 1.0 + + +# ---------------------------------------------------------------------- +# FallbackMetadataProvider failover semantics (Fix C) + + +def _stub_provider(*, result=None, error=None, exc: Exception | None = None): + if exc is not None: + call = AsyncMock(side_effect=exc) + else: + call = AsyncMock(return_value=(result, error)) + return SimpleNamespace(get_model_by_hash=call) + + +async def test_fallback_does_not_fail_over_to_network_provider_on_429(): + 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") + + fallback = FallbackMetadataProvider( + [ + ("civitai_api", civitai), + ("civarchive_api", civarchive), + ("sqlite", sqlite), + ] + ) + + result, error = await fallback.get_model_by_hash("deadbeef") + + assert result is None + assert error == "Rate limited" + civarchive.get_model_by_hash.assert_not_called() # no network failover + sqlite.get_model_by_hash.assert_called_once() # local last resort kept + + +async def test_fallback_still_fails_over_on_not_found(): + civitai = _stub_provider(result=None, error="Model not found") + civarchive = _stub_provider(result={"id": 1}, error=None) + + fallback = FallbackMetadataProvider( + [("civitai_api", civitai), ("civarchive_api", civarchive)] + ) + + result, _ = await fallback.get_model_by_hash("deadbeef") + + assert result == {"id": 1} + civarchive.get_model_by_hash.assert_called_once() + + +async def test_fallback_404_failover_still_works_after_rate_limit_change(): + """A 404 from the first network provider still reaches the second.""" + civitai = _stub_provider(result=None, error="Resource not found") + civarchive = _stub_provider(result={"id": 2}, error=None) + + fallback = FallbackMetadataProvider( + [("civitai_api", civitai), ("civarchive_api", civarchive)] + ) + + result, _ = await fallback.get_model_by_hash("deadbeef") + assert result == {"id": 2} + + +# ---------------------------------------------------------------------- +# _RateLimitRetryHelper double-wait fix + + +async def test_retry_helper_does_not_sleep_for_gate_handled_errors(): + calls = 0 + + async def failing(): + nonlocal calls + calls += 1 + error = RateLimitError("limited", retry_after=30) + error.gate_handled = True + raise error + + helper = _RateLimitRetryHelper() + start = time.monotonic() + with pytest.raises(RateLimitError) as excinfo: + await helper.run("civitai_api", failing) + elapsed = time.monotonic() - start + + assert calls == 1 # propagated immediately, no retry loop + assert elapsed < 1.0 + assert excinfo.value.provider == "civitai_api" + + +async def test_retry_helper_keeps_legacy_retry_for_ungated_errors(): + calls = 0 + + async def failing(): + nonlocal calls + calls += 1 + raise RateLimitError("limited", retry_after=None) + + helper = _RateLimitRetryHelper( + retry_limit=2, base_delay=0.01, max_delay=0.05, jitter_ratio=0.0 + ) + with pytest.raises(RateLimitError): + await helper.run("civitai_api", failing) + + assert calls == 2 # legacy retry behavior unchanged