server/ratelimit.py

"""How often one viewer may ask.

This does **not** protect a wallet. Inference is billed to whoever signed in,
so a busy visitor spends their own credits and not the owner's — see
``credentials.py``. What it protects is the Space itself: one free CPU serving
a public URL, where a single script can make it unresponsive for everyone else
without spending a penny of anyone's money.

So the ceiling is deliberately generous. A rate limiter tight enough to be felt
by a real person is worse than none at all on a Space nobody is paying for.

Configuration, all optional:

===========================  =======  ====================================
``ASE_RATE_LIMIT``           ``30``   requests per window; **0 disables it**
``ASE_RATE_WINDOW``          ``60``   seconds in a window
===========================  =======  ====================================

In memory, per process. A restart forgives everyone and a second replica keeps
its own count — both acceptable for one small Space, and both reasons not to
reach for this as a security control.
"""

from __future__ import annotations

import os
import time
from collections import deque
from dataclasses import dataclass

DEFAULT_LIMIT = 30
DEFAULT_WINDOW_S = 60
# One entry per caller seen this window. Bounded so a flood of distinct callers
# cannot grow the table without limit.
MAX_TRACKED = 4096


def limit() -> int:
    try:
        return max(0, int(os.environ.get("ASE_RATE_LIMIT", DEFAULT_LIMIT)))
    except ValueError:
        return DEFAULT_LIMIT


def window_s() -> float:
    try:
        return max(1.0, float(os.environ.get("ASE_RATE_WINDOW", DEFAULT_WINDOW_S)))
    except ValueError:
        return float(DEFAULT_WINDOW_S)


def enabled() -> bool:
    return limit() > 0


@dataclass(frozen=True)
class Verdict:
    ok: bool
    remaining: int
    retry_after_s: float = 0.0

    @property
    def detail(self) -> str:
        if self.ok:
            return ""
        return (
            f"That is more than {limit()} requests in {int(window_s())} seconds. "
            f"Try again in {max(1, int(self.retry_after_s))}s — this Space runs on one "
            f"free CPU, so the limit is there to keep it responsive for everyone."
        )


class RateLimiter:
    """A sliding window per caller.

    A sliding window rather than a fixed one: a fixed window lets twice the
    limit through across a boundary, which on a small box is exactly the burst
    it exists to stop.
    """

    def __init__(self) -> None:
        self._seen: dict[str, deque[float]] = {}

    def check(self, key: str, now: float | None = None) -> Verdict:
        cap = limit()
        if cap <= 0:
            return Verdict(True, -1)

        now = time.monotonic() if now is None else now
        span = window_s()
        hits = self._seen.get(key)
        if hits is None:
            self._prune(now, span)
            hits = self._seen.setdefault(key, deque())

        while hits and now - hits[0] >= span:
            hits.popleft()

        if len(hits) >= cap:
            return Verdict(False, 0, span - (now - hits[0]))

        hits.append(now)
        return Verdict(True, cap - len(hits))

    def _prune(self, now: float, span: float) -> None:
        """Drop callers who have gone quiet, before adding another."""
        if len(self._seen) < MAX_TRACKED:
            return
        stale = [k for k, v in self._seen.items() if not v or now - v[-1] >= span]
        for k in stale:
            del self._seen[k]
        # Still full of active callers: this is a flood, and forgetting the
        # oldest is better than growing without bound.
        if len(self._seen) >= MAX_TRACKED:
            for k in list(self._seen)[: MAX_TRACKED // 4]:
                del self._seen[k]


def caller_key(user: str | None, client_host: str | None) -> str:
    """Who is being counted.

    A signed-in viewer is counted as themselves, so sharing an office does not
    mean sharing a quota. Everyone else is counted by address, which is the
    only handle an anonymous caller has.
    """
    if user:
        return f"user:{user}"
    return f"ip:{client_host or 'unknown'}"