import time from collections import defaultdict from threading import Lock class SimpleRateLimiter: """ Simple in-memory rate limiter using sliding window. For production, consider using Redis-based rate limiting. """ def __init__(self, max_requests=60, window_seconds=60, max_in_flight=2): self.max_requests = int(max_requests) self.window_seconds = int(window_seconds) self.max_in_flight = int(max_in_flight) self._requests = defaultdict(list) self._in_flight = defaultdict(int) self._lock = Lock() def is_allowed(self, key): """ Check if the key is allowed to make a request. Args: key: Identifier (IP address, session_id, etc.) Returns: bool: True if allowed, False if rate limit exceeded """ now = time.time() window_start = now - self.window_seconds with self._lock: if self.max_requests <= 0: return True # Clean old requests requests = self._requests[key] self._requests[key] = [ts for ts in requests if ts > window_start] # Check limit if len(self._requests[key]) >= self.max_requests: return False # Record this request self._requests[key].append(now) return True def try_acquire(self, key): """Acquire one reusable in-flight slot for a session/tool key.""" with self._lock: if self.max_in_flight <= 0: return True if self._in_flight[key] >= self.max_in_flight: return False self._in_flight[key] += 1 return True def release(self, key): """Release a previously acquired slot; extra releases are harmless.""" with self._lock: if self.max_in_flight <= 0 or key not in self._in_flight: return self._in_flight[key] -= 1 if self._in_flight[key] <= 0: del self._in_flight[key] def cleanup(self, max_age_seconds=3600): """ Remove old entries to prevent memory leak. Call this periodically in a background thread. """ now = time.time() cutoff = now - max_age_seconds with self._lock: keys_to_delete = [] for key, requests in self._requests.items(): if not requests or requests[-1] < cutoff: keys_to_delete.append(key) for key in keys_to_delete: del self._requests[key]