diff --git a/docs/servers/middleware.mdx b/docs/servers/middleware.mdx index 37713a8d7..e141d4f9e 100644 --- a/docs/servers/middleware.mdx +++ b/docs/servers/middleware.mdx @@ -473,7 +473,7 @@ mcp.add_middleware(RateLimitingMiddleware( |-----------|------|---------|-------------| | `max_requests_per_second` | `float` | `10.0` | Sustained request rate | | `burst_capacity` | `int` | `20` | Maximum burst size | -| `client_id_func` | `Callable` | `None` | Custom client identification | +| `get_client_id` | `Callable` | `None` | Custom client identification | For sliding window rate limiting: diff --git a/fastmcp_slim/fastmcp/server/middleware/rate_limiting.py b/fastmcp_slim/fastmcp/server/middleware/rate_limiting.py index 703fd393a..e34202521 100644 --- a/fastmcp_slim/fastmcp/server/middleware/rate_limiting.py +++ b/fastmcp_slim/fastmcp/server/middleware/rate_limiting.py @@ -1,9 +1,10 @@ """Rate limiting middleware for protecting FastMCP servers from abuse.""" +import inspect import time from collections import defaultdict, deque -from collections.abc import Callable -from typing import Any +from collections.abc import Awaitable, Callable +from typing import Any, cast import anyio from mcp import McpError @@ -114,7 +115,9 @@ class RateLimitingMiddleware(Middleware): self, max_requests_per_second: float = 10.0, burst_capacity: int | None = None, - get_client_id: Callable[[MiddlewareContext], str] | None = None, + get_client_id: Callable[[MiddlewareContext], str] + | Callable[[MiddlewareContext], Awaitable[str]] + | None = None, global_limit: bool = False, ): """Initialize rate limiting middleware. @@ -122,7 +125,8 @@ class RateLimitingMiddleware(Middleware): Args: max_requests_per_second: Sustained requests per second allowed burst_capacity: Maximum burst capacity. If None, defaults to 2x max_requests_per_second - get_client_id: Function to extract client ID from context. If None, uses global limiting + get_client_id: Function to extract client ID from context. Can be sync or async. + If None, uses global limiting global_limit: If True, apply limit globally; if False, per-client """ self.max_requests_per_second = max_requests_per_second @@ -143,10 +147,13 @@ class RateLimitingMiddleware(Middleware): self.burst_capacity, self.max_requests_per_second ) - def _get_client_identifier(self, context: MiddlewareContext) -> str: + async def _get_client_identifier(self, context: MiddlewareContext) -> str: """Get client identifier for rate limiting.""" if self.get_client_id: - return self.get_client_id(context) + client_id = self.get_client_id(context) + if inspect.isawaitable(client_id): + return cast(str, await client_id) + return client_id return "global" async def on_request(self, context: MiddlewareContext, call_next: CallNext) -> Any: @@ -158,7 +165,7 @@ class RateLimitingMiddleware(Middleware): raise RateLimitError("Global rate limit exceeded") else: # Per-client rate limiting - client_id = self._get_client_identifier(context) + client_id = await self._get_client_identifier(context) limiter = self.limiters[client_id] allowed = await limiter.consume() if not allowed: @@ -192,14 +199,17 @@ class SlidingWindowRateLimitingMiddleware(Middleware): self, max_requests: int, window_minutes: int = 1, - get_client_id: Callable[[MiddlewareContext], str] | None = None, + get_client_id: Callable[[MiddlewareContext], str] + | Callable[[MiddlewareContext], Awaitable[str]] + | None = None, ): """Initialize sliding window rate limiting middleware. Args: max_requests: Maximum requests allowed in the time window window_minutes: Time window in minutes - get_client_id: Function to extract client ID from context + get_client_id: Function to extract client ID from context. Can be sync or async. + If None, uses global limiting """ self.max_requests = max_requests self.window_seconds = window_minutes * 60 @@ -210,15 +220,18 @@ class SlidingWindowRateLimitingMiddleware(Middleware): lambda: SlidingWindowRateLimiter(self.max_requests, self.window_seconds) ) - def _get_client_identifier(self, context: MiddlewareContext) -> str: + async def _get_client_identifier(self, context: MiddlewareContext) -> str: """Get client identifier for rate limiting.""" if self.get_client_id: - return self.get_client_id(context) + client_id = self.get_client_id(context) + if inspect.isawaitable(client_id): + return cast(str, await client_id) + return client_id return "global" async def on_request(self, context: MiddlewareContext, call_next: CallNext) -> Any: """Apply sliding window rate limiting to requests.""" - client_id = self._get_client_identifier(context) + client_id = await self._get_client_identifier(context) limiter = self.limiters[client_id] allowed = await limiter.is_allowed() diff --git a/tests/server/middleware/test_rate_limiting.py b/tests/server/middleware/test_rate_limiting.py index 59fb074f1..d2d81ef98 100644 --- a/tests/server/middleware/test_rate_limiting.py +++ b/tests/server/middleware/test_rate_limiting.py @@ -176,19 +176,28 @@ class TestRateLimitingMiddleware: assert middleware.get_client_id is get_client_id assert middleware.global_limit is True - def test_get_client_identifier_default(self, mock_context): + async def test_get_client_identifier_default(self, mock_context): """Test default client identifier.""" middleware = RateLimitingMiddleware() - assert middleware._get_client_identifier(mock_context) == "global" + assert await middleware._get_client_identifier(mock_context) == "global" - def test_get_client_identifier_custom(self, mock_context): + async def test_get_client_identifier_custom(self, mock_context): """Test custom client identifier.""" def get_client_id(ctx): return "custom_client" middleware = RateLimitingMiddleware(get_client_id=get_client_id) - assert middleware._get_client_identifier(mock_context) == "custom_client" + assert await middleware._get_client_identifier(mock_context) == "custom_client" + + async def test_get_client_identifier_async_custom(self, mock_context): + """Test custom async client identifier.""" + + async def get_client_id(ctx): + return "async_client" + + middleware = RateLimitingMiddleware(get_client_id=get_client_id) + assert await middleware._get_client_identifier(mock_context) == "async_client" async def test_on_request_success(self, mock_context, mock_call_next): """Test successful request within rate limit.""" @@ -225,6 +234,25 @@ class TestRateLimitingMiddleware: with pytest.raises(RateLimitError, match="Global rate limit exceeded"): await middleware.on_request(mock_context, mock_call_next) + async def test_on_request_async_get_client_id(self, mock_context, mock_call_next): + """Test per-client rate limiting with an async get_client_id.""" + + async def get_client_id(ctx): + return "async_client" + + middleware = RateLimitingMiddleware( + max_requests_per_second=1.0, burst_capacity=1, get_client_id=get_client_id + ) + + # First request should succeed + await middleware.on_request(mock_context, mock_call_next) + + # Second request should be rate limited for the async-resolved client + with pytest.raises( + RateLimitError, match="Rate limit exceeded for client: async_client" + ): + await middleware.on_request(mock_context, mock_call_next) + class TestSlidingWindowRateLimitingMiddleware: """Test sliding window rate limiting middleware.""" @@ -269,6 +297,23 @@ class TestSlidingWindowRateLimitingMiddleware: with pytest.raises(RateLimitError, match="Rate limit exceeded"): await middleware.on_request(mock_context, mock_call_next) + async def test_on_request_async_get_client_id(self, mock_context, mock_call_next): + """Test sliding window rate limiting with an async get_client_id.""" + + async def get_client_id(ctx): + return "async_client" + + middleware = SlidingWindowRateLimitingMiddleware( + max_requests=1, get_client_id=get_client_id + ) + + # First request should succeed + await middleware.on_request(mock_context, mock_call_next) + + # Second request should be rate limited for the async-resolved client + with pytest.raises(RateLimitError, match="client: async_client"): + await middleware.on_request(mock_context, mock_call_next) + class TestRateLimitError: """Test rate limit error."""