src/api/limits.py

"""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)