ceos/fastmcp
1
0
Fork 0
mirror of https://github.com/PrefectHQ/fastmcp.git synced 2026-08-23 05:54:19 +02:00
Code Issues Projects Releases Packages Wiki Activity

Fix: async rate limiting middleware get_client_id callbacks (#4319)

Browse source 
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
Tomasz Czochański 2026-06-24 16:27:30 +02:00 • committed by GitHub
commit 4ce8e2a5d0
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 75 additions and 17 deletions
Download patch file Download diff file Expand all files Collapse all files

2
docs/servers/middleware.mdx
View file

@ -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:

37
fastmcp_slim/fastmcp/server/middleware/rate_limiting.py
View file

@ -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()

53
tests/server/middleware/test_rate_limiting.py
View file

@ -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."""