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):
        self.max_requests = int(max_requests)
        self.window_seconds = int(window_seconds)
        self._requests = defaultdict(list)
        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:
            # 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 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]