mirror of
https://github.com/PrefectHQ/fastmcp.git
synced 2026-08-23 22:14:18 +02:00
fix: validate max_clients>=1 and add LRU eviction test coverage
🤖 Generated with Claude Code
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
parent
d280a50bfd
commit
88c5982a75
2 changed files with 85 additions and 5 deletions
|
|
@ -125,9 +125,12 @@ class RateLimitingMiddleware(Middleware):
|
|||
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
|
||||
global_limit: If True, apply limit globally; if False, per-client
|
||||
max_clients: Maximum number of per-client limiters to track. When reached,
|
||||
the least recently used client is evicted to make room for a new client.
|
||||
max_clients: Maximum number of per-client limiters to track. Must be >= 1.
|
||||
When reached, the least recently used client is evicted to make
|
||||
room for a new client.
|
||||
"""
|
||||
if max_clients < 1:
|
||||
raise ValueError(f"max_clients must be >= 1, got {max_clients}")
|
||||
self.max_requests_per_second = max_requests_per_second
|
||||
self.burst_capacity = burst_capacity or int(max_requests_per_second * 2)
|
||||
self.get_client_id = get_client_id
|
||||
|
|
@ -135,7 +138,6 @@ class RateLimitingMiddleware(Middleware):
|
|||
self._max_clients = max_clients
|
||||
|
||||
# Per-client limiters stored in LRU order (oldest at front, newest at back)
|
||||
|
||||
self._client_limiters: OrderedDict[str, TokenBucketRateLimiter] = OrderedDict()
|
||||
|
||||
# Global rate limiter
|
||||
|
|
@ -222,9 +224,12 @@ class SlidingWindowRateLimitingMiddleware(Middleware):
|
|||
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
|
||||
max_clients: Maximum number of per-client limiters to track. When reached,
|
||||
the least recently used client is evicted to make room for a new client.
|
||||
max_clients: Maximum number of per-client limiters to track. Must be >= 1.
|
||||
When reached, the least recently used client is evicted to make
|
||||
room for a new client.
|
||||
"""
|
||||
if max_clients < 1:
|
||||
raise ValueError(f"max_clients must be >= 1, got {max_clients}")
|
||||
self.max_requests = max_requests
|
||||
self.window_seconds = window_minutes * 60
|
||||
self.get_client_id = get_client_id
|
||||
|
|
|
|||
|
|
@ -270,6 +270,81 @@ class TestSlidingWindowRateLimitingMiddleware:
|
|||
await middleware.on_request(mock_context, mock_call_next)
|
||||
|
||||
|
||||
class TestRateLimiterLRUEviction:
|
||||
"""LRU bounding of per-client limiters (issue #4053)."""
|
||||
|
||||
def test_token_bucket_evicts_lru_when_full(self):
|
||||
mw = RateLimitingMiddleware(
|
||||
get_client_id=lambda ctx: ctx, max_clients=2
|
||||
)
|
||||
mw._get_limiter("a")
|
||||
mw._get_limiter("b")
|
||||
# Cache full at 2; adding "c" evicts the LRU ("a"), not the newcomer.
|
||||
mw._get_limiter("c")
|
||||
assert list(mw._client_limiters) == ["b", "c"]
|
||||
|
||||
def test_token_bucket_move_to_end_on_access_protects_active_client(self):
|
||||
mw = RateLimitingMiddleware(
|
||||
get_client_id=lambda ctx: ctx, max_clients=2
|
||||
)
|
||||
mw._get_limiter("a")
|
||||
mw._get_limiter("b")
|
||||
# Touch "a" so it is most-recently-used; "b" is now LRU.
|
||||
mw._get_limiter("a")
|
||||
mw._get_limiter("c")
|
||||
assert list(mw._client_limiters) == ["a", "c"]
|
||||
|
||||
def test_token_bucket_reuses_same_limiter_instance(self):
|
||||
mw = RateLimitingMiddleware(get_client_id=lambda ctx: ctx)
|
||||
first = mw._get_limiter("a")
|
||||
assert mw._get_limiter("a") is first
|
||||
|
||||
def test_token_bucket_bound_never_exceeds_max_clients(self):
|
||||
mw = RateLimitingMiddleware(
|
||||
get_client_id=lambda ctx: ctx, max_clients=5
|
||||
)
|
||||
for i in range(100):
|
||||
mw._get_limiter(f"client-{i}")
|
||||
assert len(mw._client_limiters) == 5
|
||||
# Only the 5 most-recent survive.
|
||||
assert list(mw._client_limiters) == [f"client-{i}" for i in range(95, 100)]
|
||||
|
||||
def test_sliding_window_evicts_lru_when_full(self):
|
||||
mw = SlidingWindowRateLimitingMiddleware(
|
||||
max_requests=10, get_client_id=lambda ctx: ctx, max_clients=2
|
||||
)
|
||||
mw._get_limiter("a")
|
||||
mw._get_limiter("b")
|
||||
mw._get_limiter("c")
|
||||
assert list(mw._client_limiters) == ["b", "c"]
|
||||
|
||||
def test_sliding_window_bound_never_exceeds_max_clients(self):
|
||||
mw = SlidingWindowRateLimitingMiddleware(
|
||||
max_requests=10, get_client_id=lambda ctx: ctx, max_clients=3
|
||||
)
|
||||
for i in range(50):
|
||||
mw._get_limiter(f"c{i}")
|
||||
assert len(mw._client_limiters) == 3
|
||||
|
||||
@pytest.mark.parametrize("bad", [0, -1, -100])
|
||||
def test_token_bucket_rejects_non_positive_max_clients(self, bad):
|
||||
with pytest.raises(ValueError, match="max_clients must be >= 1"):
|
||||
RateLimitingMiddleware(max_clients=bad)
|
||||
|
||||
@pytest.mark.parametrize("bad", [0, -1])
|
||||
def test_sliding_window_rejects_non_positive_max_clients(self, bad):
|
||||
with pytest.raises(ValueError, match="max_clients must be >= 1"):
|
||||
SlidingWindowRateLimitingMiddleware(max_requests=10, max_clients=bad)
|
||||
|
||||
def test_max_clients_one_keeps_only_newest(self):
|
||||
mw = RateLimitingMiddleware(
|
||||
get_client_id=lambda ctx: ctx, max_clients=1
|
||||
)
|
||||
mw._get_limiter("a")
|
||||
mw._get_limiter("b")
|
||||
assert list(mw._client_limiters) == ["b"]
|
||||
|
||||
|
||||
class TestRateLimitError:
|
||||
"""Test rate limit error."""
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue