"""
Rate Limiting & Throttling for MCP Server (Phase 7.3)
This module implements Token Bucket-based rate limiting to prevent API abuse
and ensure fair resource usage across all MCP clients.
Features:
- Multi-level rate limits (per minute, hour, day)
- Per-client tracking with token bucket algorithm
- Configurable limits per plugin type
- Statistics and monitoring capabilities
- Integration with audit logging
Author: Phase 7.3 Implementation
Date: 2025-01-11
"""
import logging
import os
import time
from dataclasses import dataclass, field
from datetime import UTC, datetime
from typing import Any
logger = logging.getLogger(__name__)
@dataclass
class RateLimitConfig:
"""Configuration for rate limits at different time intervals."""
per_minute: int = 60
per_hour: int = 1000
per_day: int = 10000
@classmethod
def from_env(cls, prefix: str = "") -> "RateLimitConfig":
"""Create config from environment variables."""
env_prefix = f"{prefix}_" if prefix else ""
return cls(
per_minute=int(os.getenv(f"{env_prefix}RATE_LIMIT_PER_MINUTE", "60")),
per_hour=int(os.getenv(f"{env_prefix}RATE_LIMIT_PER_HOUR", "1000")),
per_day=int(os.getenv(f"{env_prefix}RATE_LIMIT_PER_DAY", "10000")),
)
@dataclass
class TokenBucket:
"""
Token Bucket implementation for rate limiting.
The token bucket algorithm allows for burst traffic while maintaining
an average rate limit over time.
"""
capacity: int
refill_rate: float # tokens per second
tokens: float = field(default=0.0)
last_refill: float = field(default_factory=time.time)
def __post_init__(self):
"""Initialize bucket with full capacity."""
if self.tokens == 0.0:
self.tokens = float(self.capacity)
def refill(self) -> None:
"""Refill tokens based on elapsed time."""
now = time.time()
elapsed = now - self.last_refill
# Add tokens based on elapsed time
self.tokens = min(self.capacity, self.tokens + (elapsed * self.refill_rate))
self.last_refill = now
def consume(self, tokens: int = 1) -> bool:
"""
Attempt to consume tokens from the bucket.
Args:
tokens: Number of tokens to consume
Returns:
True if tokens were available and consumed, False otherwise
"""
self.refill()
if self.tokens >= tokens:
self.tokens -= tokens
return True
return False
def get_available_tokens(self) -> int:
"""Get current number of available tokens."""
self.refill()
return int(self.tokens)
def get_wait_time(self, tokens: int = 1) -> float:
"""
Calculate wait time in seconds until enough tokens are available.
Args:
tokens: Number of tokens needed
Returns:
Wait time in seconds (0 if tokens already available)
"""
self.refill()
if self.tokens >= tokens:
return 0.0
tokens_needed = tokens - self.tokens
return tokens_needed / self.refill_rate
@dataclass
class ClientRateLimitState:
"""Track rate limit state for a single client."""
client_id: str
minute_bucket: TokenBucket
hour_bucket: TokenBucket
day_bucket: TokenBucket
total_requests: int = 0
rejected_requests: int = 0
last_request_time: float = field(default_factory=time.time)
first_request_time: float = field(default_factory=time.time)
def check_and_consume(self) -> tuple[bool, str, float]:
"""
Check if request is allowed and consume tokens.
Returns:
Tuple of (allowed, reason, retry_after_seconds)
"""
# Check each time window (most restrictive first)
if not self.minute_bucket.consume():
wait_time = self.minute_bucket.get_wait_time()
self.rejected_requests += 1
return False, "Rate limit exceeded: too many requests per minute", wait_time
if not self.hour_bucket.consume():
wait_time = self.hour_bucket.get_wait_time()
# Refund the minute token since we're rejecting
self.minute_bucket.tokens = min(
self.minute_bucket.capacity, self.minute_bucket.tokens + 1
)
self.rejected_requests += 1
return False, "Rate limit exceeded: too many requests per hour", wait_time
if not self.day_bucket.consume():
wait_time = self.day_bucket.get_wait_time()
# Refund tokens since we're rejecting
self.minute_bucket.tokens = min(
self.minute_bucket.capacity, self.minute_bucket.tokens + 1
)
self.hour_bucket.tokens = min(self.hour_bucket.capacity, self.hour_bucket.tokens + 1)
self.rejected_requests += 1
return False, "Rate limit exceeded: daily limit reached", wait_time
# All checks passed
self.total_requests += 1
self.last_request_time = time.time()
return True, "", 0.0
def get_stats(self) -> dict[str, Any]:
"""Get statistics for this client."""
now = time.time()
uptime = now - self.first_request_time
return {
"client_id": self.client_id,
"total_requests": self.total_requests,
"rejected_requests": self.rejected_requests,
"success_rate": (
(self.total_requests - self.rejected_requests) / self.total_requests
if self.total_requests > 0
else 1.0
),
"available_tokens": {
"per_minute": self.minute_bucket.get_available_tokens(),
"per_hour": self.hour_bucket.get_available_tokens(),
"per_day": self.day_bucket.get_available_tokens(),
},
"limits": {
"per_minute": self.minute_bucket.capacity,
"per_hour": self.hour_bucket.capacity,
"per_day": self.day_bucket.capacity,
},
"last_request": datetime.fromtimestamp(self.last_request_time, tz=UTC).isoformat(),
"uptime_seconds": uptime,
}
class RateLimiter:
"""
Rate limiter using Token Bucket algorithm.
Provides multi-level rate limiting (per minute, hour, day) with
per-client tracking and configurable limits.
"""
def __init__(self):
"""Initialize rate limiter with default configuration."""
self.clients: dict[str, ClientRateLimitState] = {}
self.global_stats = {"total_requests": 0, "total_rejected": 0, "start_time": time.time()}
# Load default configuration from environment
self.default_config = RateLimitConfig.from_env()
# Plugin-specific configurations
self.plugin_configs: dict[str, RateLimitConfig] = {
"wordpress": RateLimitConfig.from_env("WORDPRESS"),
"woocommerce": RateLimitConfig.from_env("WOOCOMMERCE"),
}
logger.info(
"Rate limiter initialized with default limits: "
f"{self.default_config.per_minute}/min, "
f"{self.default_config.per_hour}/hour, "
f"{self.default_config.per_day}/day"
)
def _get_or_create_client_state(
self, client_id: str, plugin_type: str | None = None
) -> ClientRateLimitState:
"""Get or create rate limit state for a client."""
if client_id not in self.clients:
# Determine which config to use
config = self.plugin_configs.get(plugin_type, self.default_config)
# Create token buckets for each time window
minute_bucket = TokenBucket(
capacity=config.per_minute,
refill_rate=config.per_minute / 60.0, # tokens per second
)
hour_bucket = TokenBucket(
capacity=config.per_hour, refill_rate=config.per_hour / 3600.0
)
day_bucket = TokenBucket(capacity=config.per_day, refill_rate=config.per_day / 86400.0)
self.clients[client_id] = ClientRateLimitState(
client_id=client_id,
minute_bucket=minute_bucket,
hour_bucket=hour_bucket,
day_bucket=day_bucket,
)
logger.debug(f"Created rate limit state for client: {client_id}")
return self.clients[client_id]
def check_rate_limit(
self, client_id: str, tool_name: str | None = None, plugin_type: str | None = None
) -> tuple[bool, str, float]:
"""
Check if request should be allowed based on rate limits.
Args:
client_id: Identifier for the client (e.g., auth token hash)
tool_name: Name of the tool being called (for logging)
plugin_type: Type of plugin (wordpress, woocommerce, etc.)
Returns:
Tuple of (allowed, message, retry_after_seconds)
"""
# Get or create client state
client_state = self._get_or_create_client_state(client_id, plugin_type)
# Update global stats
self.global_stats["total_requests"] += 1
# Check and consume tokens
allowed, message, retry_after = client_state.check_and_consume()
if not allowed:
# Track rejection
client_state.rejected_requests += 1
self.global_stats["total_rejected"] += 1
logger.warning(
f"Rate limit exceeded for client {client_id[:8]}... "
f"(tool: {tool_name}, reason: {message}, "
f"retry_after: {retry_after:.1f}s)"
)
else:
logger.debug(
f"Rate limit check passed for client {client_id[:8]}... " f"(tool: {tool_name})"
)
return allowed, message, retry_after
def get_client_stats(self, client_id: str) -> dict[str, Any] | None:
"""
Get statistics for a specific client.
Args:
client_id: Client identifier
Returns:
Client statistics or None if client not found
"""
if client_id not in self.clients:
return None
return self.clients[client_id].get_stats()
def get_all_stats(self) -> dict[str, Any]:
"""Get global rate limiter statistics."""
now = time.time()
uptime = now - self.global_stats["start_time"]
# Calculate per-client stats
client_stats = []
for _client_id, client_state in self.clients.items():
client_stats.append(client_state.get_stats())
return {
"global": {
"total_requests": self.global_stats["total_requests"],
"total_rejected": self.global_stats["total_rejected"],
"rejection_rate": (
self.global_stats["total_rejected"] / self.global_stats["total_requests"]
if self.global_stats["total_requests"] > 0
else 0.0
),
"active_clients": len(self.clients),
"uptime_seconds": uptime,
"start_time": datetime.fromtimestamp(
self.global_stats["start_time"], tz=UTC
).isoformat(),
},
"default_limits": {
"per_minute": self.default_config.per_minute,
"per_hour": self.default_config.per_hour,
"per_day": self.default_config.per_day,
},
"plugin_limits": {
plugin: {
"per_minute": config.per_minute,
"per_hour": config.per_hour,
"per_day": config.per_day,
}
for plugin, config in self.plugin_configs.items()
},
"clients": client_stats,
}
def reset_client(self, client_id: str) -> bool:
"""
Reset rate limit state for a specific client.
Args:
client_id: Client identifier
Returns:
True if client was reset, False if client not found
"""
if client_id in self.clients:
del self.clients[client_id]
logger.info(f"Reset rate limit state for client: {client_id}")
return True
return False
def reset_all(self) -> int:
"""
Reset all client rate limit states.
Returns:
Number of clients reset
"""
count = len(self.clients)
self.clients.clear()
self.global_stats = {"total_requests": 0, "total_rejected": 0, "start_time": time.time()}
logger.info(f"Reset rate limit state for {count} clients")
return count
def configure_limits(
self,
plugin_type: str,
per_minute: int | None = None,
per_hour: int | None = None,
per_day: int | None = None,
) -> None:
"""
Configure rate limits for a specific plugin type.
Args:
plugin_type: Plugin type identifier
per_minute: Requests per minute limit
per_hour: Requests per hour limit
per_day: Requests per day limit
"""
if plugin_type not in self.plugin_configs:
self.plugin_configs[plugin_type] = RateLimitConfig()
config = self.plugin_configs[plugin_type]
if per_minute is not None:
config.per_minute = per_minute
if per_hour is not None:
config.per_hour = per_hour
if per_day is not None:
config.per_day = per_day
logger.info(
f"Updated rate limits for {plugin_type}: "
f"{config.per_minute}/min, {config.per_hour}/hour, {config.per_day}/day"
)
# Singleton instance
_rate_limiter: RateLimiter | None = None
def get_rate_limiter() -> RateLimiter:
"""Get or create the global rate limiter instance."""
global _rate_limiter
if _rate_limiter is None:
_rate_limiter = RateLimiter()
return _rate_limiter
