""" Lightweight in-process rate limiting. A sliding-window log limiter, keyed by client IP. No external dependencies and no Redis — fine for the single-process server this project runs. If you ever scale to multiple workers you'll want a shared store instead, since each worker keeps its own counters. limiter = RateLimiter(limit=30, window=60) # 30 requests / minute allowed, remaining, retry_after = limiter.hit("1.2.3.4", now=time.time()) """ from __future__ import annotations import threading import time from collections import defaultdict, deque from typing import Deque, Dict, Tuple from starlette.requests import Request from starlette.responses import JSONResponse class RateLimiter: """Sliding-window request counter shared across requests (thread-safe).""" def __init__(self, limit: int, window: float = 60.0): self.limit = limit self.window = window self._hits: Dict[str, Deque[float]] = defaultdict(deque) self._lock = threading.Lock() def hit(self, key: str, now: float) -> Tuple[bool, int, float]: """Record a request for `key`. Returns (allowed, remaining, retry_after). `retry_after` is seconds until the oldest in-window hit expires (0 when the request is allowed). """ cutoff = now - self.window with self._lock: q = self._hits[key] while q and q[0] <= cutoff: q.popleft() if len(q) >= self.limit: retry_after = q[0] + self.window - now return False, 0, max(0.0, retry_after) q.append(now) return True, self.limit - len(q), 0.0 def install_rate_limit(app, limiter: RateLimiter, *, protect_prefix: str = "/v1") -> None: """Attach `limiter` to a FastAPI/Starlette app as HTTP middleware. Only paths under `protect_prefix` are limited (so /healthz stays open). On every response we set the standard `X-RateLimit-*` headers; over-limit requests get a 429 with `Retry-After` and an OpenAI-shaped error body. """ @app.middleware("http") async def _rate_limit(request: Request, call_next): if not request.url.path.startswith(protect_prefix): return await call_next(request) key = _client_key(request) now = time.time() allowed, remaining, retry_after = limiter.hit(key, now) if not allowed: resp = JSONResponse( status_code=429, content={ "error": { "message": ( f"Rate limit exceeded: {limiter.limit} requests per " f"{int(limiter.window)}s. Retry in " f"{retry_after:.1f}s." ), "type": "rate_limit_error", } }, ) resp.headers["Retry-After"] = str(int(retry_after) + 1) else: resp = await call_next(request) resp.headers["X-RateLimit-Limit"] = str(limiter.limit) resp.headers["X-RateLimit-Remaining"] = str(remaining) resp.headers["X-RateLimit-Reset"] = str(int(now + limiter.window)) return resp def _client_key(request: Request) -> str: """Best-effort client identity: first X-Forwarded-For hop, else peer IP.""" fwd = request.headers.get("x-forwarded-for") if fwd: return fwd.split(",")[0].strip() return request.client.host if request.client else "unknown"