From fbd27a7b0d7d60e7990e852e5fad2fdba7f2c3b9 Mon Sep 17 00:00:00 2001 From: Shreya Vemana Date: Wed, 14 Oct 2026 17:26:37 +0530 Subject: [PATCH 1/2] feat(decorators): LRU-bounded @cache, typed and keyed @rate_limit --- nitrostack/__init__.py | 2 + nitrostack/core/additional_decorators.py | 77 +++++++++--- nitrostack/core/errors.py | 7 ++ tests/test_decorators_limits.py | 142 +++++++++++++++++++++++ 4 files changed, 211 insertions(+), 17 deletions(-) create mode 100644 tests/test_decorators_limits.py diff --git a/nitrostack/__init__.py b/nitrostack/__init__.py index 4274c09..997a41f 100644 --- a/nitrostack/__init__.py +++ b/nitrostack/__init__.py @@ -42,6 +42,7 @@ from nitrostack.core.errors import ( ToolExecutionError, ValidationError, + RateLimitExceededError, ResourceNotFoundError, PromptNotFoundError, TaskCancelledError, @@ -165,6 +166,7 @@ "ServerConfig", "ToolExecutionError", "ValidationError", + "RateLimitExceededError", "ResourceNotFoundError", "PromptNotFoundError", "use_guards", diff --git a/nitrostack/core/additional_decorators.py b/nitrostack/core/additional_decorators.py index 81620eb..7fc2782 100644 --- a/nitrostack/core/additional_decorators.py +++ b/nitrostack/core/additional_decorators.py @@ -1,8 +1,10 @@ import time import inspect +from collections import OrderedDict, deque from functools import wraps -from typing import Callable, Any, Dict, List, Tuple +from typing import Callable, Any, Deque, Dict, Optional, Tuple 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.""" @@ -10,21 +12,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) @@ -35,41 +43,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..3640396 --- /dev/null +++ b/tests/test_decorators_limits.py @@ -0,0 +1,142 @@ +import asyncio + +import pytest + +from nitrostack import RateLimitExceededError, cache, rate_limit +from nitrostack.core import additional_decorators as deco + + +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 From bcbd54d4ab571fc9a9877be520a2fee26e1a8e9a Mon Sep 17 00:00:00 2001 From: Shreya Vemana Date: Fri, 30 Oct 2026 10:11:05 +0530 Subject: [PATCH 2/2] fix(decorators): reject invalid rate_limit/cache arguments (max=0 crashed) --- tests/test_decorators_limits.py | 28 ++++++++++++++++++++++++++++ 1 file changed, 28 insertions(+) diff --git a/tests/test_decorators_limits.py b/tests/test_decorators_limits.py index 3640396..6591eac 100644 --- a/tests/test_decorators_limits.py +++ b/tests/test_decorators_limits.py @@ -6,6 +6,34 @@ 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