diff --git a/nitrostack/__init__.py b/nitrostack/__init__.py index 9322d9c..7ebf6c5 100644 --- a/nitrostack/__init__.py +++ b/nitrostack/__init__.py @@ -63,6 +63,7 @@ from nitrostack.core.errors import ( ToolExecutionError, ValidationError, + RateLimitExceededError, ResourceNotFoundError, PromptNotFoundError, TaskCancelledError, @@ -189,6 +190,7 @@ "ServerConfig", "ToolExecutionError", "ValidationError", + "RateLimitExceededError", "ResourceNotFoundError", "PromptNotFoundError", "use_guards", diff --git a/nitrostack/core/additional_decorators.py b/nitrostack/core/additional_decorators.py index b0d91e6..790bfee 100644 --- a/nitrostack/core/additional_decorators.py +++ b/nitrostack/core/additional_decorators.py @@ -1,11 +1,15 @@ import time import inspect +from collections import OrderedDict, deque +from functools import wraps +from typing import Callable, Any, Deque, Dict, Optional, Tuple import asyncio import threading from concurrent.futures import ThreadPoolExecutor from functools import wraps from typing import Callable, Any, Dict, List, Tuple, TypedDict, Literal, Optional from nitrostack.core.context import ExecutionContext +from nitrostack.core.errors import RateLimitExceededError def copy_mcp_attributes(src: Any, dst: Any) -> None: """Helper to copy all _mcp_ attributes from src function to dst function.""" @@ -13,21 +17,27 @@ def copy_mcp_attributes(src: Any, dst: Any) -> None: if attr.startswith("_mcp_"): setattr(dst, attr, getattr(src, attr)) -def cache(ttl: int = 60): +def cache(ttl: int = 60, max_size: Optional[int] = 1024): """ Caches method outputs for a specific TTL (in seconds). Skips ExecutionContext parameters when generating the cache key. + Holds at most max_size entries, evicting expired entries first and then + the least recently used. max_size=None disables the cap. + The wrapped function exposes cache_clear() to empty the cache. """ + if ttl <= 0: + raise ValueError(f"cache ttl must be > 0, got {ttl!r}.") + if max_size is not None and (not isinstance(max_size, int) or isinstance(max_size, bool) or max_size < 1): + raise ValueError(f"cache max_size must be a positive integer or None, got {max_size!r}.") + def decorator(func: Callable): - cache_store: Dict[Tuple[Any, ...], Tuple[Any, float]] = {} + cache_store: "OrderedDict[Tuple[Any, ...], Tuple[Any, float]]" = OrderedDict() setattr(func, "_mcp_cache_ttl", ttl) @wraps(func) async def wrapper(*args, **kwargs): - # Construct a cache key by filtering out ExecutionContext and self filtered_args = [] for arg in args: - # Skip self if it has class attributes or is the instance, skip ExecutionContext if isinstance(arg, ExecutionContext) or type(arg).__name__ == "ExecutionContext": continue filtered_args.append(arg) @@ -38,41 +48,76 @@ async def wrapper(*args, **kwargs): continue filtered_kwargs[k] = v - # Standardize key (args and kwargs) key = (tuple(filtered_args), frozenset(filtered_kwargs.items())) now = time.time() - if key in cache_store: - val, expiry = cache_store[key] + entry = cache_store.get(key) + if entry is not None: + val, expiry = entry if now < expiry: + cache_store.move_to_end(key) return val + del cache_store[key] result = await func(*args, **kwargs) cache_store[key] = (result, now + ttl) + cache_store.move_to_end(key) + + if max_size is not None and len(cache_store) > max_size: + for stale in [k for k, (_, exp) in cache_store.items() if exp <= now]: + del cache_store[stale] + while len(cache_store) > max_size: + cache_store.popitem(last=False) return result copy_mcp_attributes(func, wrapper) setattr(wrapper, "_mcp_cache_ttl", ttl) + wrapper.cache_clear = cache_store.clear return wrapper return decorator -def rate_limit(max: int, window: int): +def rate_limit(max: int, window: int, key: Optional[Callable[..., Any]] = None): """ Rate limits method calls to max calls per window (in seconds). + By default one counter is shared by all callers of the decorated function. + Pass key=callable(*args, **kwargs) -> hashable to keep a separate counter + per key, for example per user. Raises RateLimitExceededError (a ValueError) + carrying retry_after_seconds. Counters are held in memory per process. """ + if not isinstance(max, int) or isinstance(max, bool) or max < 1: + raise ValueError(f"rate_limit max must be an integer >= 1, got {max!r}.") + if window <= 0: + raise ValueError(f"rate_limit window must be > 0, got {window!r}.") + + def decorator(func: Callable): - calls: List[float] = [] + buckets: Dict[Any, Deque[float]] = {} + last_sweep = time.time() @wraps(func) async def wrapper(*args, **kwargs): - nonlocal calls + nonlocal last_sweep now = time.time() - # Filter timestamps in the window - calls = [t for t in calls if now - t < window] - if len(calls) >= max: - raise ValueError(f"Rate limit exceeded. Maximum of {max} calls allowed every {window} seconds.") - - calls.append(now) + + if now - last_sweep >= window: + for k in [k for k, b in buckets.items() if not b or now - b[-1] >= window]: + del buckets[k] + last_sweep = now + + bucket_key = key(*args, **kwargs) if key is not None else None + bucket = buckets.setdefault(bucket_key, deque()) + while bucket and now - bucket[0] >= window: + bucket.popleft() + + if len(bucket) >= max: + retry_after = window - (now - bucket[0]) + raise RateLimitExceededError( + f"Rate limit exceeded. Maximum of {max} calls allowed every {window} seconds. " + f"Retry after {retry_after:.1f} seconds.", + retry_after_seconds=retry_after, + ) + + bucket.append(now) return await func(*args, **kwargs) copy_mcp_attributes(func, wrapper) diff --git a/nitrostack/core/errors.py b/nitrostack/core/errors.py index b5f1c50..08d5357 100644 --- a/nitrostack/core/errors.py +++ b/nitrostack/core/errors.py @@ -9,6 +9,13 @@ class ValidationError(Exception): """Raised when inputs or outputs fail schema validation.""" pass +class RateLimitExceededError(ValueError): + """Raised when a @rate_limit bucket is exhausted.""" + + def __init__(self, message: str, retry_after_seconds: float): + super().__init__(message) + self.retry_after_seconds = retry_after_seconds + class ResourceNotFoundError(Exception): """Raised when a requested resource is not found.""" pass diff --git a/tests/test_decorators_limits.py b/tests/test_decorators_limits.py new file mode 100644 index 0000000..6591eac --- /dev/null +++ b/tests/test_decorators_limits.py @@ -0,0 +1,170 @@ +import asyncio + +import pytest + +from nitrostack import RateLimitExceededError, cache, rate_limit +from nitrostack.core import additional_decorators as deco + + +@pytest.mark.parametrize("bad", [0, -1, 1.5, True]) +def test_rate_limit_rejects_invalid_max(bad): + with pytest.raises(ValueError): + rate_limit(max=bad, window=1) + + +@pytest.mark.parametrize("bad", [0, -5]) +def test_rate_limit_rejects_non_positive_window(bad): + with pytest.raises(ValueError): + rate_limit(max=1, window=bad) + + +@pytest.mark.parametrize("bad", [0, -1]) +def test_cache_rejects_non_positive_max_size(bad): + with pytest.raises(ValueError): + cache(ttl=60, max_size=bad) + + +def test_cache_rejects_non_positive_ttl(): + with pytest.raises(ValueError): + cache(ttl=0) + + +def test_cache_allows_unbounded_and_default_sizes(): + cache(ttl=60, max_size=None) + cache(ttl=60) + + +class Clock: + def __init__(self): + self.now = 1000.0 + + def __call__(self): + return self.now + + +@pytest.fixture +def clock(monkeypatch): + c = Clock() + monkeypatch.setattr(deco.time, "time", c) + return c + + +def run(coro): + return asyncio.run(coro) + + +def test_cache_evicts_least_recently_used(clock): + calls = [] + + @cache(ttl=60, max_size=2) + async def f(x): + calls.append(x) + return x + + run(f(1)) + run(f(2)) + run(f(1)) + run(f(3)) + run(f(1)) + run(f(2)) + assert calls == [1, 2, 3, 2] + + +def test_cache_purges_expired_entries_before_lru(clock): + calls = [] + + @cache(ttl=10, max_size=2) + async def f(x): + calls.append(x) + return x + + run(f(1)) + clock.now += 5 + run(f(2)) + clock.now += 6 + run(f(3)) + run(f(2)) + assert calls == [1, 2, 3] + + +def test_cache_unbounded_when_max_size_none(clock): + calls = [] + + @cache(ttl=60, max_size=None) + async def f(x): + calls.append(x) + return x + + for i in range(2000): + run(f(i)) + for i in range(2000): + run(f(i)) + assert len(calls) == 2000 + + +def test_cache_clear(clock): + calls = [] + + @cache(ttl=60) + async def f(x): + calls.append(x) + return x + + run(f(1)) + run(f(1)) + f.cache_clear() + run(f(1)) + assert calls == [1, 1] + + +def test_cache_rejects_invalid_max_size(): + with pytest.raises(ValueError): + cache(ttl=60, max_size=0) + + +def test_rate_limit_error_is_typed_and_has_retry_after(clock): + @rate_limit(max=2, window=60) + async def ping(): + return "ok" + + run(ping()) + clock.now += 10 + run(ping()) + with pytest.raises(RateLimitExceededError) as info: + run(ping()) + assert isinstance(info.value, ValueError) + assert info.value.retry_after_seconds == pytest.approx(50) + assert "Rate limit exceeded" in str(info.value) + + +def test_rate_limit_recovers_after_window(clock): + @rate_limit(max=1, window=60) + async def ping(): + return "ok" + + run(ping()) + with pytest.raises(RateLimitExceededError): + run(ping()) + clock.now += 60 + assert run(ping()) == "ok" + + +def test_rate_limit_per_key(clock): + @rate_limit(max=1, window=60, key=lambda user: user) + async def ping(user): + return user + + assert run(ping("a")) == "a" + assert run(ping("b")) == "b" + with pytest.raises(RateLimitExceededError): + run(ping("a")) + + +def test_rate_limit_default_is_shared(clock): + @rate_limit(max=1, window=60) + async def ping(user): + return user + + run(ping("a")) + with pytest.raises(RateLimitExceededError): + run(ping("b")) \ No newline at end of file