From e3df2dac15997aa4c9f2fee85e3bf52881172745 Mon Sep 17 00:00:00 2001
From: Yujong Lee <yujong@berri.ai>
Date: Tue, 3 Nov 2026 17:28:45 +0900
Subject: [PATCH] fix(auth-cache): type UserApiKeyCache values as plain data

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
---
 .../proxy/common_utils/cache_pydantic_utils.py    |  6 +++++-
 litellm/proxy/common_utils/user_api_key_cache.py  |  8 ++++----
 .../proxy/common_utils/test_user_api_key_cache.py | 15 +++++++++++++++
 3 files changed, 24 insertions(+), 5 deletions(-)

diff --git a/litellm/proxy/common_utils/cache_pydantic_utils.py b/litellm/proxy/common_utils/cache_pydantic_utils.py
index 3703cf7c916..34d46be56cd 100644
--- a/litellm/proxy/common_utils/cache_pydantic_utils.py
+++ b/litellm/proxy/common_utils/cache_pydantic_utils.py
@@ -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:
     """
diff --git a/litellm/proxy/common_utils/user_api_key_cache.py b/litellm/proxy/common_utils/user_api_key_cache.py
index a2272c3f73d..d61ab8c8518 100644
--- a/litellm/proxy/common_utils/user_api_key_cache.py
+++ b/litellm/proxy/common_utils/user_api_key_cache.py
@@ -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
diff --git a/tests/unit/proxy/common_utils/test_user_api_key_cache.py b/tests/unit/proxy/common_utils/test_user_api_key_cache.py
index 577e5e5a389..21ae8655b97 100644
--- a/tests/unit/proxy/common_utils/test_user_api_key_cache.py
+++ b/tests/unit/proxy/common_utils/test_user_api_key_cache.py
@@ -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:
     """