""" Rate limiting, TTL caching, and retry with backoff for Super Search MCP server. Provides: - TokenBucket: async-aware token bucket rate limiter (per provider) - TTLCache: dict-based cache with per-key expiration - rate_limited_search: decorator that applies rate limiting + caching + retry """ import asyncio import time import logging from functools import wraps from typing import Any, Callable, Optional logger = logging.getLogger("super-search.ratelimit") # ── Rate limits (tokens per interval in seconds) ──────────────────────────────── RATE_LIMITS: dict[str, dict[str, float]] = { "searxng": {"tokens": 30, "interval": 1.0}, # local, high limit "exa": {"tokens": 10, "interval": 1.0}, # paid tier "firecrawl": {"tokens": 5, "interval": 60.0}, # 1k/mo, be careful "opencorporates": {"tokens": 1, "interval": 1.0}, # ~1/sec free tier "courtlistener": {"tokens": 10, "interval": 6.0}, # ~100/min free tier "duckduckgo": {"tokens": 5, "interval": 1.0}, # free, be respectful "wikipedia": {"tokens": 30, "interval": 1.0}, # public API, generous } # ── Cache TTLs (seconds) ──────────────────────────────────────────────────────── CACHE_TTL: dict[str, int] = { "searxng": 120, # 2 min — web results change "exa": 300, # 5 min "opencorporates": 3600, # 1 hour — company records don't change often "courtlistener": 3600, # 1 hour — court cases are infrequent "duckduckgo": 300, # 5 min — instant answers can change "wikipedia": 3600, # 1 hour — encyclopedia articles are stable } # ── Exceptions considered "retryable" ─────────────────────────────────────────── RETRYABLE_STATUSES: set[int] = {429, 503, 502, 504} # ══════════════════════════════════════════════════════════════════════════════════ # Token Bucket Rate Limiter # ══════════════════════════════════════════════════════════════════════════════════ class TokenBucket: """Async-aware token bucket rate limiter. Tokens refill at a constant rate (tokens per interval). Each acquire() consumes one token; if none are available the caller must wait. """ def __init__(self, tokens: float, interval: float) -> None: self._capacity = float(tokens) self._tokens = float(tokens) self._interval = float(interval) self._refill_rate = self._capacity / self._interval self._last_refill = time.monotonic() self._lock = asyncio.Lock() def _refill(self) -> None: now = time.monotonic() elapsed = now - self._last_refill self._tokens = min(self._capacity, self._tokens + elapsed * self._refill_rate) self._last_refill = now async def acquire(self) -> None: """Acquire one token, waiting if necessary until one is available.""" async with self._lock: self._refill() if self._tokens >= 1.0: self._tokens -= 1.0 return # Calculate wait time until next token wait = (1.0 - self._tokens) / self._refill_rate logger.debug("TokenBucket: waiting %.2fs for token", wait) # Sleep outside the lock so other consumers can also wait await asyncio.sleep(wait) # Retry after sleep async with self._lock: self._refill() self._tokens -= 1.0 # ── Global bucket registry ────────────────────────────────────────────────────── _buckets: dict[str, TokenBucket] = {} def get_bucket(provider: str) -> TokenBucket: """Get or create a token bucket for a provider.""" if provider not in _buckets: limits = RATE_LIMITS.get(provider, {"tokens": 5, "interval": 1.0}) _buckets[provider] = TokenBucket( tokens=limits["tokens"], interval=limits["interval"], ) return _buckets[provider] # ══════════════════════════════════════════════════════════════════════════════════ # TTL Cache # ══════════════════════════════════════════════════════════════════════════════════ class TTLCache: """Simple dict-based TTL cache with per-key expiration.""" def __init__(self) -> None: self._cache: dict[str, tuple[Any, float]] = {} # key → (value, expires_at) self._lock = asyncio.Lock() async def get(self, key: str) -> Optional[Any]: """Return cached value if not expired, else None.""" async with self._lock: entry = self._cache.get(key) if entry is None: return None value, expires_at = entry if time.monotonic() >= expires_at: del self._cache[key] return None return value async def set(self, key: str, value: Any, ttl: int) -> None: """Store a value with the given TTL in seconds.""" async with self._lock: self._cache[key] = (value, time.monotonic() + ttl) async def invalidate(self, key: str) -> None: """Remove an entry from the cache.""" async with self._lock: self._cache.pop(key, None) # ── Global cache instance ─────────────────────────────────────────────────────── _cache = TTLCache() def cache_key(provider: str, query: str, limit: int) -> str: """Build a deterministic cache key for a search.""" return f"{provider}:{query.strip().lower()}:{limit}" # ══════════════════════════════════════════════════════════════════════════════════ # Retry with exponential backoff # ══════════════════════════════════════════════════════════════════════════════════ import httpx # for HTTPStatusError type hint async def _retry_with_backoff( coro_factory: Callable[[], Any], provider: str, max_retries: int = 3, base_delay: float = 1.0, ) -> Any: """Execute a coroutine with retry on 429/5xx, using exponential backoff. Respects Retry-After header when present. """ for attempt in range(max_retries + 1): try: result = await coro_factory() return result except httpx.HTTPStatusError as e: status = e.response.status_code if status not in RETRYABLE_STATUSES or attempt == max_retries: raise # Check for Retry-After header retry_after = e.response.headers.get("Retry-After", "") if retry_after and retry_after.isdigit(): delay = float(retry_after) else: delay = base_delay * (2 ** attempt) logger.warning( "Provider %s returned %d (attempt %d/%d). Retrying in %.1fs…", provider, status, attempt + 1, max_retries, delay, ) await asyncio.sleep(delay) # ══════════════════════════════════════════════════════════════════════════════════ # Rate-limited + cached search wrapper # ══════════════════════════════════════════════════════════════════════════════════ def rate_limited_search(provider: str): """Decorator: wraps a search coroutine with rate limiting + TTL caching. Usage: @rate_limited_search("opencorporates") async def _opencorporates_search(query, limit=5) -> list[dict]: ... """ def decorator(func): @wraps(func) async def wrapper(query: str, limit: int = 5, **kwargs): # 1. Check cache key = cache_key(provider, query, limit) cached = await _cache.get(key) if cached is not None: logger.debug("Cache HIT for %s:%s", provider, key) return cached # 2. Acquire rate-limit token bucket = get_bucket(provider) await bucket.acquire() # 3. Execute with retry ttl = CACHE_TTL.get(provider, 300) async def _call(): return await func(query, limit=limit, **kwargs) try: result = await _retry_with_backoff(_call, provider) except Exception: # Don't cache failures raise # 4. Cache result if result: await _cache.set(key, result, ttl) logger.debug("Cache SET for %s:%s (TTL=%ds)", provider, key, ttl) return result return wrapper return decorator