src/shared/ratelimit.py

"""Abuse controls for AI-powered Spaces.

Layers (see docs/rate-limiting.md):
1. Per-IP sliding window (server-side, survives incognito/new sessions).
2. Global daily cap — cannot be bypassed by identity rotation; bounds
   worst-case daily spend to DAILY_CAP x cost-per-generation.
3. Queue concurrency (configured on the Blocks) bounds burn rate.
4. api_name=False on generate events removes the scripted-API surface.

State is in-memory and resets on Space restart; that is acceptable for
abuse control (not accounting).
"""

import os
import threading
import time
from collections import defaultdict


class RateLimiter:
    def __init__(
        self,
        per_ip_limit: int = int(os.environ.get("LF_IP_LIMIT", 8)),
        window_seconds: int = int(os.environ.get("LF_IP_WINDOW_SECONDS", 3600)),
        daily_cap: int = int(os.environ.get("LF_DAILY_CAP", 500)),
    ):
        self.per_ip_limit = per_ip_limit
        self.window = window_seconds
        self.daily_cap = daily_cap
        self._hits: dict[str, list[float]] = defaultdict(list)
        self._day: str = ""
        self._day_count = 0
        self._lock = threading.Lock()

    @staticmethod
    def client_ip(request) -> str:
        """Real client IP. On HF Spaces the app sits behind a proxy, so the
        client is the first entry of x-forwarded-for, not request.client."""
        if request is None:
            return "unknown"
        fwd = dict(request.headers).get("x-forwarded-for", "")
        if fwd:
            return fwd.split(",")[0].strip()
        try:
            return request.client.host
        except Exception:
            return "unknown"

    def check(self, request) -> tuple[bool, str]:
        """Returns (allowed, user_safe_message)."""
        ip = self.client_ip(request)
        now = time.time()
        today = time.strftime("%Y-%m-%d", time.gmtime(now))
        with self._lock:
            if today != self._day:
                self._day, self._day_count = today, 0
            if self._day_count >= self.daily_cap:
                return False, (
                    "The forge is cooling down for today — daily generation "
                    "limit reached. Come back tomorrow!"
                )
            window_start = now - self.window
            self._hits[ip] = [t for t in self._hits[ip] if t > window_start]
            if len(self._hits[ip]) >= self.per_ip_limit:
                minutes = max(1, int((self._hits[ip][0] + self.window - now) / 60))
                return False, (
                    f"You've hit the hourly generation limit. Try again in about "
                    f"{minutes} minute{'s' if minutes != 1 else ''}."
                )
            self._hits[ip].append(now)
            self._day_count += 1
            # Opportunistic cleanup so the dict doesn't grow unbounded.
            if len(self._hits) > 10000:
                stale = [k for k, v in self._hits.items() if not v or v[-1] < window_start]
                for k in stale:
                    del self._hits[k]
        return True, ""