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