From a44623154f0e0a23dfa8e9ca62b43e13bedd2ef9 Mon Sep 17 00:00:00 2001 From: Aadarsh Mahesh K Date: Wed, 4 Nov 2026 16:35:45 +0530 Subject: [PATCH 1/2] fix(cache): handle unhashable args and add key= callable --- nitrostack/core/additional_decorators.py | 70 +++++-- tests/test_cache_unhashable.py | 242 +++++++++++++++++++++++ 2 files changed, 301 insertions(+), 11 deletions(-) create mode 100644 tests/test_cache_unhashable.py diff --git a/nitrostack/core/additional_decorators.py b/nitrostack/core/additional_decorators.py index 790bfee..4144813 100644 --- a/nitrostack/core/additional_decorators.py +++ b/nitrostack/core/additional_decorators.py @@ -1,13 +1,12 @@ +import json 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 collections import OrderedDict, deque from concurrent.futures import ThreadPoolExecutor from functools import wraps -from typing import Callable, Any, Dict, List, Tuple, TypedDict, Literal, Optional +from typing import Any, Callable, Deque, Dict, List, Literal, Optional, Tuple, TypedDict from nitrostack.core.context import ExecutionContext from nitrostack.core.errors import RateLimitExceededError @@ -17,13 +16,52 @@ def copy_mcp_attributes(src: Any, dst: Any) -> None: if attr.startswith("_mcp_"): setattr(dst, attr, getattr(src, attr)) -def cache(ttl: int = 60, max_size: Optional[int] = 1024): + +def _make_cache_key_part(value: Any, label: str) -> Any: + """Convert *value* into a hashable cache-key component. + + Pydantic BaseModel instances are serialized via model_dump(mode="json"). + Natively hashable values are used as-is. + Anything else raises a descriptive TypeError. + """ + try: + from pydantic import BaseModel as PydanticBaseModel # type: ignore + if isinstance(value, PydanticBaseModel): + dumped = value.model_dump(mode="json") + return ("__pydantic__", type(value).__name__, json.dumps(dumped, sort_keys=True)) + except ImportError: + pass + + try: + hash(value) + return value + except TypeError: + pass + + raise TypeError( + f"cache: argument {label!r} of type {type(value).__name__!r} is not hashable " + f"and has no built-in serialization strategy. " + f"Pass a key= callable to cache() to handle this type explicitly." + ) + + +def cache(ttl: int = 60, max_size: Optional[int] = 1024, key: Optional[Callable[..., Any]] = None): """ 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. + + Args: + ttl: Time-to-live for cached entries in seconds (must be > 0). + max_size: Maximum number of cache entries. None means unlimited. + key: Optional callable that receives the same *args* and **kwargs + as the decorated function (after ExecutionContext arguments + are filtered out) and returns a hashable cache key. + When provided it overrides the default key-building logic, + which serializes Pydantic models automatically and raises a + descriptive TypeError for other unhashable arguments. """ if ttl <= 0: raise ValueError(f"cache ttl must be > 0, got {ttl!r}.") @@ -48,20 +86,30 @@ async def wrapper(*args, **kwargs): continue filtered_kwargs[k] = v - key = (tuple(filtered_args), frozenset(filtered_kwargs.items())) + if key is not None: + cache_key = key(*filtered_args, **filtered_kwargs) + else: + key_args = tuple( + _make_cache_key_part(v, f"args[{i}]") for i, v in enumerate(filtered_args) + ) + key_kwargs = frozenset( + (k, _make_cache_key_part(v, k)) for k, v in filtered_kwargs.items() + ) + cache_key = (key_args, key_kwargs) + cache_lookup_key = cache_key now = time.time() - entry = cache_store.get(key) + entry = cache_store.get(cache_lookup_key) if entry is not None: val, expiry = entry if now < expiry: - cache_store.move_to_end(key) + cache_store.move_to_end(cache_lookup_key) return val - del cache_store[key] + del cache_store[cache_lookup_key] result = await func(*args, **kwargs) - cache_store[key] = (result, now + ttl) - cache_store.move_to_end(key) + cache_store[cache_lookup_key] = (result, now + ttl) + cache_store.move_to_end(cache_lookup_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]: diff --git a/tests/test_cache_unhashable.py b/tests/test_cache_unhashable.py new file mode 100644 index 0000000..9ae7323 --- /dev/null +++ b/tests/test_cache_unhashable.py @@ -0,0 +1,242 @@ +from __future__ import annotations + +import asyncio +import pytest + +from pydantic import BaseModel + +from nitrostack.core.additional_decorators import cache, _make_cache_key_part +from nitrostack.core.context import ExecutionContext + + +def run(coro): + return asyncio.run(coro) + + +class SearchQuery(BaseModel): + q: str + limit: int = 10 + + +class FilterOptions(BaseModel): + tags: list[str] = [] + active: bool = True + + +class TestMakeCacheKeyPart: + def test_pydantic_model_returns_deterministic_tuple(self): + m = SearchQuery(q="hello", limit=5) + part = _make_cache_key_part(m, "arg") + assert isinstance(part, tuple) + assert part[0] == "__pydantic__" + assert part[1] == "SearchQuery" + assert isinstance(part[2], str) + + def test_same_model_state_produces_same_key(self): + a = SearchQuery(q="same", limit=3) + b = SearchQuery(q="same", limit=3) + assert _make_cache_key_part(a, "x") == _make_cache_key_part(b, "x") + + def test_different_model_state_produces_different_key(self): + a = SearchQuery(q="foo") + b = SearchQuery(q="bar") + assert _make_cache_key_part(a, "x") != _make_cache_key_part(b, "x") + + def test_hashable_value_returned_as_is(self): + assert _make_cache_key_part(42, "n") == 42 + assert _make_cache_key_part("text", "s") == "text" + assert _make_cache_key_part(3.14, "f") == 3.14 + assert _make_cache_key_part(None, "n") is None + assert _make_cache_key_part((1, 2), "t") == (1, 2) + + def test_unhashable_list_raises_descriptive_type_error(self): + with pytest.raises(TypeError, match="args\\[0\\]"): + _make_cache_key_part([1, 2, 3], "args[0]") + + def test_unhashable_dict_raises_descriptive_type_error(self): + with pytest.raises(TypeError, match="my_param"): + _make_cache_key_part({"a": 1}, "my_param") + + def test_error_message_contains_type_name(self): + with pytest.raises(TypeError, match="list"): + _make_cache_key_part([1], "x") + + def test_error_message_mentions_key_callable(self): + with pytest.raises(TypeError, match="key="): + _make_cache_key_part({"k": "v"}, "x") + + +class TestCacheWithPydanticModels: + def test_pydantic_positional_arg_caches_correctly(self): + calls = [] + + @cache(ttl=60) + async def search(query: SearchQuery) -> str: + calls.append(query.q) + return f"result:{query.q}" + + q = SearchQuery(q="python") + assert run(search(q)) == "result:python" + assert run(search(q)) == "result:python" + assert calls == ["python"] + + def test_different_pydantic_models_produce_different_cache_entries(self): + calls = [] + + @cache(ttl=60) + async def search(query: SearchQuery) -> str: + calls.append(query.q) + return query.q + + run(search(SearchQuery(q="a"))) + run(search(SearchQuery(q="b"))) + run(search(SearchQuery(q="a"))) + assert calls == ["a", "b"] + + def test_pydantic_keyword_arg_caches_correctly(self): + calls = [] + + @cache(ttl=60) + async def fetch(query: SearchQuery, options: FilterOptions) -> int: + calls.append(1) + return len(query.q) + + q = SearchQuery(q="hello") + f = FilterOptions(tags=["python"]) + assert run(fetch(query=q, options=f)) == 5 + assert run(fetch(query=q, options=f)) == 5 + assert calls == [1] + + def test_pydantic_model_cache_isolation_between_decorated_functions(self): + calls_a, calls_b = [], [] + + @cache(ttl=60) + async def fn_a(q: SearchQuery) -> str: + calls_a.append(q.q) + return "a" + + @cache(ttl=60) + async def fn_b(q: SearchQuery) -> str: + calls_b.append(q.q) + return "b" + + q = SearchQuery(q="x") + run(fn_a(q)) + run(fn_b(q)) + assert calls_a == ["x"] + assert calls_b == ["x"] + + def test_pydantic_positional_and_keyword_mixed(self): + calls = [] + + @cache(ttl=60) + async def op(query: SearchQuery, limit: int = 5) -> str: + calls.append((query.q, limit)) + return f"{query.q}:{limit}" + + q = SearchQuery(q="mix") + run(op(q, limit=3)) + run(op(q, limit=3)) + assert calls == [("mix", 3)] + + run(op(q, limit=10)) + assert len(calls) == 2 + + +class TestCacheWithUnhashableArguments: + def test_unhashable_positional_arg_raises_type_error(self): + @cache(ttl=60) + async def fn(data: list) -> int: + return len(data) + + with pytest.raises(TypeError, match="args\\[0\\]"): + run(fn([1, 2, 3])) + + def test_unhashable_keyword_arg_raises_type_error(self): + @cache(ttl=60) + async def fn(mapping: dict) -> int: + return len(mapping) + + with pytest.raises(TypeError, match="mapping"): + run(fn(mapping={"a": 1})) + + def test_error_identifies_unhashable_type(self): + @cache(ttl=60) + async def fn(x: set) -> int: + return len(x) + + with pytest.raises(TypeError, match="set"): + run(fn({1, 2})) + + +class TestCacheCustomKeyCallable: + def test_custom_key_accepts_pydantic_model(self): + calls = [] + + @cache(ttl=60, key=lambda q: q.q) + async def search(query: SearchQuery) -> str: + calls.append(query.q) + return query.q + + run(search(SearchQuery(q="hi"))) + run(search(SearchQuery(q="hi"))) + assert calls == ["hi"] + + def test_custom_key_discriminates_on_custom_field(self): + calls = [] + + @cache(ttl=60, key=lambda q: q.q) + async def search(query: SearchQuery) -> str: + calls.append(query.q) + return query.q + + run(search(SearchQuery(q="same", limit=1))) + run(search(SearchQuery(q="same", limit=99))) + assert calls == ["same"] + + def test_custom_key_accepts_unhashable_list(self): + calls = [] + + @cache(ttl=60, key=lambda data: tuple(sorted(data))) + async def process(data: list) -> int: + calls.append(len(data)) + return len(data) + + run(process([3, 1, 2])) + run(process([3, 1, 2])) + assert calls == [3] + + def test_custom_key_with_kwargs(self): + calls = [] + + @cache(ttl=60, key=lambda *, q, limit: (q, limit)) + async def search(*, q: str, limit: int) -> str: + calls.append((q, limit)) + return f"{q}:{limit}" + + run(search(q="hello", limit=5)) + run(search(q="hello", limit=5)) + assert calls == [("hello", 5)] + + def test_custom_key_function_called_with_filtered_args(self): + seen = [] + + def my_key(*args, **kwargs): + seen.append((args, kwargs)) + return str(args) + str(sorted(kwargs.items())) + + @cache(ttl=60, key=my_key) + async def fn(value: int, ctx: ExecutionContext | None = None) -> int: + return value + + ctx = ExecutionContext.__new__(ExecutionContext) + run(fn(42, ctx=ctx)) + + assert all( + not any(isinstance(a, ExecutionContext) for a in call_args) + for call_args, _ in seen + ) + assert all( + not any(isinstance(v, ExecutionContext) for v in call_kwargs.values()) + for _, call_kwargs in seen + ) From 928e4bc3cbc0bbfebba1b54460a6f906b67bf087 Mon Sep 17 00:00:00 2001 From: Aadarsh Mahesh K Date: Wed, 4 Nov 2026 18:07:32 +0530 Subject: [PATCH 2/2] fix(cache): include excluded fields in key and scope type by module --- nitrostack/core/additional_decorators.py | 40 +++++++++++++++------ tests/test_cache_unhashable.py | 46 ++++++++++++++++++++++-- 2 files changed, 74 insertions(+), 12 deletions(-) diff --git a/nitrostack/core/additional_decorators.py b/nitrostack/core/additional_decorators.py index 4144813..ae3ee4f 100644 --- a/nitrostack/core/additional_decorators.py +++ b/nitrostack/core/additional_decorators.py @@ -20,15 +20,27 @@ def copy_mcp_attributes(src: Any, dst: Any) -> None: def _make_cache_key_part(value: Any, label: str) -> Any: """Convert *value* into a hashable cache-key component. - Pydantic BaseModel instances are serialized via model_dump(mode="json"). + Pydantic BaseModel instances are keyed from all declared fields (including + those marked exclude=True) via pydantic_core.to_jsonable_python, identified + by their fully-qualified class name (module + qualname) to prevent collisions + between identically-named models from different modules. Natively hashable values are used as-is. Anything else raises a descriptive TypeError. """ try: + import pydantic_core from pydantic import BaseModel as PydanticBaseModel # type: ignore if isinstance(value, PydanticBaseModel): - dumped = value.model_dump(mode="json") - return ("__pydantic__", type(value).__name__, json.dumps(dumped, sort_keys=True)) + cls = type(value) + # Use all declared fields (bypasses exclude=True) so that excluded + all_fields = {name: getattr(value, name) for name in cls.model_fields} + serialized = json.dumps( + pydantic_core.to_jsonable_python(all_fields), sort_keys=True + ) + # Use module + qualname so two classes named "Query" in different + # modules never share a cache entry. + type_id = f"{cls.__module__}.{cls.__qualname__}" + return ("__pydantic__", type_id, serialized) except ImportError: pass @@ -87,7 +99,16 @@ async def wrapper(*args, **kwargs): filtered_kwargs[k] = v if key is not None: - cache_key = key(*filtered_args, **filtered_kwargs) + raw_key = key(*filtered_args, **filtered_kwargs) + try: + hash(raw_key) + except TypeError: + raise TypeError( + f"cache: the key= callable returned a value of type " + f"{type(raw_key).__name__!r} which is not hashable. " + f"Ensure key= returns a hashable object (e.g. a string, int, or tuple)." + ) + cache_key = raw_key else: key_args = tuple( _make_cache_key_part(v, f"args[{i}]") for i, v in enumerate(filtered_args) @@ -97,19 +118,18 @@ async def wrapper(*args, **kwargs): ) cache_key = (key_args, key_kwargs) - cache_lookup_key = cache_key now = time.time() - entry = cache_store.get(cache_lookup_key) + entry = cache_store.get(cache_key) if entry is not None: val, expiry = entry if now < expiry: - cache_store.move_to_end(cache_lookup_key) + cache_store.move_to_end(cache_key) return val - del cache_store[cache_lookup_key] + del cache_store[cache_key] result = await func(*args, **kwargs) - cache_store[cache_lookup_key] = (result, now + ttl) - cache_store.move_to_end(cache_lookup_key) + cache_store[cache_key] = (result, now + ttl) + cache_store.move_to_end(cache_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]: diff --git a/tests/test_cache_unhashable.py b/tests/test_cache_unhashable.py index 9ae7323..e10f2f3 100644 --- a/tests/test_cache_unhashable.py +++ b/tests/test_cache_unhashable.py @@ -3,7 +3,7 @@ import asyncio import pytest -from pydantic import BaseModel +from pydantic import BaseModel, Field from nitrostack.core.additional_decorators import cache, _make_cache_key_part from nitrostack.core.context import ExecutionContext @@ -29,7 +29,7 @@ def test_pydantic_model_returns_deterministic_tuple(self): part = _make_cache_key_part(m, "arg") assert isinstance(part, tuple) assert part[0] == "__pydantic__" - assert part[1] == "SearchQuery" + assert "SearchQuery" in part[1] assert isinstance(part[2], str) def test_same_model_state_produces_same_key(self): @@ -142,6 +142,40 @@ async def op(query: SearchQuery, limit: int = 5) -> str: run(op(q, limit=10)) assert len(calls) == 2 + def test_excluded_field_produces_separate_cache_entries(self): + class UserQuery(BaseModel): + q: str + user_id: str = Field(default="", exclude=True) + + calls = [] + + @cache(ttl=60) + async def fetch(uq: UserQuery) -> str: + calls.append(uq.user_id) + return uq.user_id + + assert run(fetch(UserQuery(q="hello", user_id="alice"))) == "alice" + assert run(fetch(UserQuery(q="hello", user_id="bob"))) == "bob" + assert calls == ["alice", "bob"] + + def test_same_class_name_different_module_no_cache_collision(self): + class Q(BaseModel): + q: str + + OtherQ = type("Q", (BaseModel,), {"__annotations__": {"q": str}, "__module__": "other_module"}) + + calls = [] + + @cache(ttl=60) + async def fetch(v) -> str: + result = type(v).__module__ + calls.append(result) + return result + + r1, r2 = run(fetch(Q(q="x"))), run(fetch(OtherQ(q="x"))) + assert r1 != r2 + assert len(calls) == 2 + class TestCacheWithUnhashableArguments: def test_unhashable_positional_arg_raises_type_error(self): @@ -240,3 +274,11 @@ async def fn(value: int, ctx: ExecutionContext | None = None) -> int: not any(isinstance(v, ExecutionContext) for v in call_kwargs.values()) for _, call_kwargs in seen ) + + def test_custom_key_returning_unhashable_raises_descriptive_error(self): + @cache(ttl=60, key=lambda x: [x]) + async def fn(x: int) -> int: + return x + + with pytest.raises(TypeError, match="key= callable"): + run(fn(1))