fix: keep union order and literal types in serializer cache keys, bound the cache

Author: devin-ai-integration[bot]

src/openrouter/_hooks/serializer_cache.py

@@ -8,19 +8,33 @@
 """
 
 import functools
-from typing import Any, Callable, Tuple, Type
+from typing import Any, Callable, Tuple, Type, get_args, get_origin
 
 from pydantic import BaseModel, ConfigDict
 
 from ..utils import serializers
 
 _uncached_create_model: Callable[..., Type[BaseModel]] = serializers.create_model
 
+# Bounded so dynamically created types can't grow the cache without limit.
+_CACHE_SIZE = 1024
 
-@functools.lru_cache(maxsize=None)
+
+def _type_key(typ: Any) -> Any:
+    # Equal annotations can still validate differently: `int | float ==
+    # float | int` and `Literal[1] == Literal[True]`, but pydantic tries union
+    # members in order and checks literal types. Keep order and leaf types.
+    args = get_args(typ)
+    if not args:
+        return (type(typ), typ)
+    return (get_origin(typ), tuple(_type_key(arg) for arg in args))
+
+
+@functools.lru_cache(maxsize=_CACHE_SIZE)
 def _cached_wrapper_model(
-    name: str, typ: Any, config_items: Tuple[Tuple[str, Any], ...]
+    name: str, type_key: Any, config_items: Tuple[Tuple[str, Any], ...], typ: Any
 ) -> Type[BaseModel]:
+    del type_key
     return _uncached_create_model(
         name, body=(typ, ...), __config__=ConfigDict(**dict(config_items))  # type: ignore[typeddict-item]
     )
@@ -36,8 +50,11 @@ def _create_model(name: str, **kwargs: Any) -> Type[BaseModel]:
         or not isinstance(config, dict)
     ):
         return _uncached_create_model(name, **kwargs)
+    typ = body[0]
     try:
-        return _cached_wrapper_model(name, body[0], tuple(sorted(config.items())))
+        return _cached_wrapper_model(
+            name, _type_key(typ), tuple(sorted(config.items())), typ
+        )
     except TypeError:
         return _uncached_create_model(name, **kwargs)
 

tests/test_serializer_cache.py

@@ -1,6 +1,8 @@
-import json
+from typing import Annotated, List, Literal, Optional, Union
 
 import httpx
+import pydantic
+import pytest
 
 from openrouter import OpenRouter
 from openrouter._hooks import serializer_cache
@@ -65,12 +67,94 @@ def counting_create_model(*args, **kwargs):
     assert second.model_dump() == first.model_dump()
 
 
-def test_cached_models_serialize_like_uncached_ones():
+def _uncached(monkeypatch):
+    monkeypatch.setattr(
+        serializers, "create_model", serializer_cache._uncached_create_model
+    )
+
+
+def _outcome(fn):
+    try:
+        result = fn()
+    except pydantic.ValidationError as error:
+        return ("error", [(e["type"], e["loc"]) for e in error.errors()])
+    return ("ok", type(result), result)
+
+
+class _Aliased(pydantic.BaseModel):
+    model_config = pydantic.ConfigDict(populate_by_name=True)
+    max_tokens: Optional[int] = pydantic.Field(default=None, alias="maxTokens")
+
+
+CASES = [
+    ("1", int | float),
+    ("1", float | int),
+    ("1", Union[int, float]),
+    ("1", Union[float, int]),
+    (True, Literal[1]),
+    (True, Literal[True]),
+    (1, Literal[True]),
+    ({"a": [1, 2]}, dict[str, list[int]]),
+    ({"a": ["x"]}, dict[str, list[int]]),
+    ({"maxTokens": 3}, _Aliased),
+    ({"max_tokens": 3}, _Aliased),
+    ({"maxTokens": "x"}, _Aliased),
+    (None, Optional[_Aliased]),
+    ([1, "2"], List[int]),
+]
+
+
+def _all_outcomes():
+    unmarshalled = [_outcome(lambda: serializers.unmarshal(v, t)) for v, t in CASES]
+    marshalled = [
+        _outcome(lambda: serializers.marshal_json(v, t)) for v, t in CASES
+    ]
+    return unmarshalled + marshalled
+
+
+@pytest.mark.parametrize("order", [1, -1], ids=["forward", "reverse"])
+def test_cached_models_behave_like_uncached_ones(monkeypatch, order):
+    with monkeypatch.context() as m:
+        _uncached(m)
+        expected = _all_outcomes()
+
+    _client()
+    serializer_cache._cached_wrapper_model.cache_clear()
+    # Warm the cache in the given order, then compare every case.
+    for value, typ in CASES[::order]:
+        _outcome(lambda: serializers.unmarshal(value, typ))
+        _outcome(lambda: serializers.marshal_json(value, typ))
+
+    assert _all_outcomes() == expected
+
+
+def test_unhashable_types_fall_back_to_uncached_models(monkeypatch):
+    built = []
+    uncached_create_model = serializer_cache._uncached_create_model
+
+    def counting_create_model(*args, **kwargs):
+        built.append(args[0])
+        return uncached_create_model(*args, **kwargs)
+
     _client()
+    serializer_cache._cached_wrapper_model.cache_clear()
+    monkeypatch.setattr(serializer_cache, "_uncached_create_model", counting_create_model)
+    typ = Annotated[int, {"unhashable": []}]
+
+    assert serializers.unmarshal("5", typ) == 5
+    assert serializers.unmarshal("5", typ) == 5
+    assert built == ["Unmarshaller", "Unmarshaller"]
+    assert serializer_cache._cached_wrapper_model.cache_info().currsize == 0
+
+
+def test_cache_size_is_bounded():
+    _client()
+    serializer_cache._cached_wrapper_model.cache_clear()
 
-    value = {"a": [1, 2]}
-    typ = dict[str, list[int]]
+    for i in range(serializer_cache._CACHE_SIZE + 50):
+        dynamic = pydantic.create_model(f"Dynamic{i}", x=(int, ...))
+        serializers.unmarshal({"x": i}, dynamic)
 
-    assert json.loads(serializers.marshal_json(value, typ)) == value
-    assert serializers.unmarshal(value, typ) == value
-    assert serializers.unmarshal(value, typ) == value
+    info = serializer_cache._cached_wrapper_model.cache_info()
+    assert info.maxsize == serializer_cache._CACHE_SIZE
+    assert info.currsize == serializer_cache._CACHE_SIZE