"""Cost-abuse guardrails: per-user rate limiting + USD spend caps. The owner's #1 worry is a runaway paid bill (Gemini OCR / Anthropic / embeddings / rerank). Every endpoint that can fan out to a paid provider calls :func:`enforce_paid_request` *before* it does any work. Two independent gates: 1. **Rate limit** — an in-process per-user token bucket so a single login can't fire paid requests in a tight loop. Sustained ``api.rate_limit_per_min`` with an ``api.rate_limit_burst`` allowance. Checked first because it's in-memory and sheds a flood without touching the DB. 2. **Spend caps** — checked against the cost ledger: * global : ``SUM(llm_calls.cost_usd)`` for the current UTC day >= ``api.daily_spend_cap_usd`` * per-user: ``SUM(query_history.cost_usd)`` for the user, today >= ``api.per_user_daily_spend_cap_usd`` (non-admins only) Either tripping raises :class:`SpendCapExceeded` (429) until the UTC day rolls over (or the operator raises the cap in ``config.yaml > api``). Single-process assumption: the HF Space runs one replica, so an in-memory bucket is correct and lock-free contention is negligible. If the deploy is ever scaled out, swap the bucket for a shared store (Turso/Redis) — the call sites (``enforce_paid_request``) don't change. """ from __future__ import annotations import threading import time from src.api.errors import RateLimited, SpendCapExceeded from src.config import load_config from src.lib.auth.models import User # Defaults used when config.yaml has no `api:` section (fresh checkout). The # numbers are safety ceilings, not budgets — set generous enough not to block # legitimate solo-researcher use, low enough that a leaked login can't run up # a catastrophic bill. Tune in config.yaml > api. _DEFAULTS = { "daily_spend_cap_usd": 10.0, "per_user_daily_spend_cap_usd": 5.0, "rate_limit_per_min": 12.0, "rate_limit_burst": 6.0, "max_queued_jobs": 12.0, # backpressure: reject paid work past this many active jobs } def _api_cfg() -> dict: try: return dict(load_config().section("api") or {}) except Exception: # noqa: BLE001 — never let config IO break a request gate return {} def _f(cfg: dict, key: str) -> float: try: return float(cfg.get(key, _DEFAULTS[key])) except (TypeError, ValueError): return _DEFAULTS[key] # ---------- Per-user token bucket ----------------------------------------- class _Bucket: __slots__ = ("tokens", "updated") def __init__(self, tokens: float, updated: float) -> None: self.tokens = tokens self.updated = updated _buckets: dict[str, _Bucket] = {} _buckets_lock = threading.Lock() def _check_rate(username: str, per_min: float, burst: float) -> None: if per_min <= 0: return # limiter disabled refill = per_min / 60.0 # tokens per second cap = max(1.0, burst) now = time.monotonic() with _buckets_lock: b = _buckets.get(username) if b is None: # First request: start full, spend one. _buckets[username] = _Bucket(tokens=cap - 1.0, updated=now) return b.tokens = min(cap, b.tokens + (now - b.updated) * refill) b.updated = now if b.tokens < 1.0: deficit = 1.0 - b.tokens retry = max(1, int(deficit / refill + 0.999)) if refill > 0 else 60 raise RateLimited(retry_after_seconds=retry) b.tokens -= 1.0 def _reset_buckets_for_test() -> None: """Clear all buckets — test-only hook so a flood test starts clean.""" with _buckets_lock: _buckets.clear() # ---------- Spend caps ---------------------------------------------------- def _check_spend(user: User) -> None: cfg = _api_cfg() # Imported here (not at module top) so importing this module never drags # in the cost repo / DB layer until a paid request actually gates. from src.lib.costs import repo as costs global_cap = _f(cfg, "daily_spend_cap_usd") if global_cap > 0: spent = costs.cost_today() if spent >= global_cap: raise SpendCapExceeded( f"The daily spend ceiling of ${global_cap:.2f} has been reached " f"(${spent:.2f} spent today). Paid actions are paused until " f"tomorrow (UTC).", scope="global", cap_usd=global_cap, spent_usd=spent, ) # Per-user query cap applies to the public-facing risk (non-admins). user_cap = _f(cfg, "per_user_daily_spend_cap_usd") if user_cap > 0 and not user.is_admin: uspent = costs.user_query_spend_today_usd(user.username) if uspent >= user_cap: raise SpendCapExceeded( f"Your daily query budget of ${user_cap:.2f} has been used up " f"(${uspent:.2f} today). It resets tomorrow (UTC).", scope="user", cap_usd=user_cap, spent_usd=uspent, ) # ---------- Login throttle (brute-force protection) ----------------------- # Passwords on this deploy are human-set and weak; without a throttle the # public login is brute-forceable. Tight bucket per (username / client IP). _LOGIN_PER_MIN = 8.0 _LOGIN_BURST = 5.0 def enforce_login_attempt(identifier: str) -> None: """Throttle login attempts keyed by ``identifier`` (a username or IP). Raises :class:`RateLimited` (429) once the per-key bucket is empty. Call once per axis (username and client IP) so neither a single-account brute force nor a username-spray from one host can run unbounded. """ _check_rate(f"login:{identifier}", per_min=_LOGIN_PER_MIN, burst=_LOGIN_BURST) def _check_job_capacity() -> None: """Backpressure: refuse new paid work when too many jobs are already queued/running, so accepted-but-waiting jobs can't grow unbounded (each spawns a daemon thread that blocks on the global semaphore). """ cap = int(_f(_api_cfg(), "max_queued_jobs")) if cap <= 0: return from src.api.jobs import store active = 0 for st in ("queued", "running"): try: active += len(store.list_jobs(status=st, limit=cap + 1)) except Exception: # noqa: BLE001 — don't block on a transient store error return if active >= cap: raise RateLimited( "The server is at job capacity right now; please retry shortly.", retry_after_seconds=10, ) # ---------- Public gate --------------------------------------------------- def enforce_paid_request(user: User) -> None: """Gate a paid endpoint. Raises before any spend happens. Order: the in-memory rate bucket first (cheap, sheds floods without a DB hit), then job-capacity backpressure, then the ledger-backed spend caps (the hard money ceiling). Call at the very top of every endpoint that fans out to a paid provider. """ cfg = _api_cfg() _check_rate( user.username, per_min=_f(cfg, "rate_limit_per_min"), burst=_f(cfg, "rate_limit_burst"), ) _check_job_capacity() _check_spend(user)