mcp/utils/rate_limiter.py

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]