fix(auth-cache): type UserApiKeyCache values as plain data
Some checks failed
Cost map pins recheck / rerun-cost-map-guard (push) Has been cancelled
LiteLLM Rust DB / db-queries (push) Has been cancelled
LiteLLM Rust / rust-lint (push) Has been cancelled
LiteLLM Rust / rust-test (push) Has been cancelled
LiteLLM Rust / rust-wheel (push) Has been cancelled
Terraform Modules / fmt, validate, test (aws) (push) Has been cancelled
Terraform Modules / fmt, validate, test (gcp) (push) Has been cancelled

Every cache setter now takes a pydantic model or plain data (str, number, bool,
None, sequences, str-keyed mappings), which is what Redis can hold and what the
Rust auth cache accepts. No production writer passes anything else, and pyright
rejects a new caller that does
This commit is contained in:
Yujong Lee 2026-11-03 17:28:45 +09:00
parent 6438d6a797
commit e3df2dac15
3 changed files with 24 additions and 5 deletions

View file

@ -16,7 +16,8 @@ are encoded; pass a Pydantic model or convert with e.g. ``dataclasses.asdict`` f
from __future__ import annotations
from typing import Any, TypeVar
from collections.abc import Mapping, Sequence
from typing import Any, TypeAlias, TypeVar
from pydantic import BaseModel, ValidationError
@ -24,6 +25,9 @@ from litellm._logging import verbose_proxy_logger
T = TypeVar("T", bound=BaseModel)
CacheData: TypeAlias = "str | int | float | bool | None | Sequence[CacheData] | Mapping[str, CacheData]"
CacheValue: TypeAlias = "BaseModel | CacheData"
class CacheCodec:
"""

View file

@ -12,7 +12,7 @@ from litellm.caching.dual_cache import DualCache
from litellm.caching.in_memory_cache import InMemoryCache
from litellm.caching.redis_cache import RedisCache
from litellm.constants import DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL
from litellm.proxy.common_utils.cache_pydantic_utils import CacheCodec
from litellm.proxy.common_utils.cache_pydantic_utils import CacheCodec, CacheValue
if TYPE_CHECKING:
from opentelemetry.trace import Span
@ -223,14 +223,14 @@ class UserApiKeyCache(DualCache):
return None
return decoded
def set_cache(self, key: str | None, value: object, local_only: bool = False, **kwargs: object):
def set_cache(self, key: str | None, value: CacheValue, local_only: bool = False, **kwargs: object):
model_type: Final = cast(type[BaseModel] | None, kwargs.pop("model_type", None))
payload: Final[object] = CacheCodec.serialize(value, model_type=model_type)
if key is not None and is_user_key_cache_key(key):
return self.key_object_cache.set_cache(key=key, value=payload, local_only=local_only, **kwargs)
return super().set_cache(key=key, value=payload, local_only=local_only, **kwargs)
async def async_set_cache(self, key: str | None, value: object, local_only: bool = False, **kwargs: object):
async def async_set_cache(self, key: str | None, value: CacheValue, local_only: bool = False, **kwargs: object):
"""Inside a request the Redis SET rides the request's pipeline (memory is written at once); anywhere
else, or with options the pipeline does not carry, it goes to Redis directly as before."""
model_type: Final = cast(type[BaseModel] | None, kwargs.pop("model_type", None))
@ -291,7 +291,7 @@ class UserApiKeyCache(DualCache):
self.budget_windows_absent_cache.flush_cache()
async def async_set_cache_pipeline(
self, cache_list: Sequence[tuple[str, object]], local_only: bool = False, **kwargs: object
self, cache_list: Sequence[tuple[str, CacheValue]], local_only: bool = False, **kwargs: object
) -> None:
"""
Batch writes with the same Codec boundary as ``async_set_cache`` without

View file

@ -1,3 +1,4 @@
import datetime
import hashlib
import json
from typing import Any, Final
@ -251,6 +252,20 @@ class TestUserApiKeyCache:
with pytest.raises(TypeError):
fake.set_cache("k2", {"ok": NotSerializable()})
@pytest.mark.asyncio
async def test_a_model_holding_a_datetime_is_cached_as_its_json_dump(self):
cache = UserApiKeyCache()
key_obj = UserAPIKeyAuth(token="abc", expires=datetime.datetime(2026, 1, 1, tzinfo=datetime.timezone.utc))
await cache.async_set_cache("k", key_obj, model_type=UserAPIKeyAuth)
stored = cache.get_cache("k", local_only=True)
assert stored == key_obj.model_dump(mode="json")
assert json.loads(json.dumps(stored)) == stored
cached = await cache.async_get_cache("k", model_type=UserAPIKeyAuth)
assert cached is not None
assert cached.model_dump(mode="json") == key_obj.model_dump(mode="json")
class TestUserKeyObjectPartition:
"""