Issue #42 : [decorators] @cache never evicts; @rate_limit is process-global and raises bare ValueError

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",

nitrostack/core/additional_decorators.py

@@ -1,33 +1,43 @@
 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."""
     for attr in dir(src):
         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)

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

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"))