From 30885467ff0e8cb5a609744a521f5c55acd2a554 Mon Sep 17 00:00:00 2001 From: harish-berri Date: Tue, 21 Apr 2026 19:57:02 +0000 Subject: [PATCH 01/28] Add debugger settings to debug single worker proxy_server per request. --- .gitignore | 1 - .vscode/launch.json | 24 ++++++++++++++++++++++++ 2 files changed, 24 insertions(+), 1 deletion(-) create mode 100644 .vscode/launch.json diff --git a/.gitignore b/.gitignore index 38bf9554b5bb..f5591b7f2c6c 100644 --- a/.gitignore +++ b/.gitignore @@ -90,7 +90,6 @@ test.py litellm_config.yaml !.github/observatory/litellm_config.yaml .cursor -.vscode/launch.json litellm/proxy/to_delete_loadtest_work/* update_model_cost_map.py tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py diff --git a/.vscode/launch.json b/.vscode/launch.json new file mode 100644 index 000000000000..e48948319b02 --- /dev/null +++ b/.vscode/launch.json @@ -0,0 +1,24 @@ +{ + "version": "0.2.0", + "configurations": [ + { + "name": "LiteLLM proxy (debug)", + "type": "debugpy", + "request": "launch", + "program": "${workspaceFolder}/litellm/proxy/proxy_cli.py", + "cwd": "${workspaceFolder}", + "args": [ + "--host", + "0.0.0.0", + "--config", + "dev_config.yaml", + "--port", + "4000", + "--num_workers", + "1" + ], + "console": "integratedTerminal", + "justMyCode": false + } + ] +} \ No newline at end of file From d58f657fa281b72b27d0c2596d7e710a673e42c0 Mon Sep 17 00:00:00 2001 From: harish-berri Date: Tue, 21 Apr 2026 20:12:42 +0000 Subject: [PATCH 02/28] Replace assertions with simple if conditions. assertions raise an exception which are not great for performance (specifically repeated throws). Logic remains the same and the tests are still passing --- litellm/proxy/auth/user_api_key_auth.py | 18 ++++++++++-------- 1 file changed, 10 insertions(+), 8 deletions(-) diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index b6366e41d5d5..205e61767508 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -1115,9 +1115,7 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 is_master_key_valid = False ## VALIDATE MASTER KEY ## - try: - assert isinstance(master_key, str) - except Exception: + if not isinstance(master_key, str): raise HTTPException( status_code=500, detail={ @@ -1183,11 +1181,15 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 if len(api_key) > 8 else "****" ) - assert api_key.startswith( - "sk-" - ), "LiteLLM Virtual Key expected. Received={}, expected to start with 'sk-'.".format( - _masked_key - ) # prevent token hashes from being used + if not api_key.startswith("sk-"): + raise HTTPException( + status_code=status.HTTP_401_UNAUTHORIZED, + detail=( + "LiteLLM Virtual Key expected. Received={}, expected to start with 'sk-'.".format( + _masked_key + ) + ), + ) # prevent token hashes from being used else: verbose_logger.warning( "litellm.proxy.proxy_server.user_api_key_auth(): Warning - Key is not a string. Got type={}".format( From 984287daaaf462cff6b4f3b0a861b69f0bd34965 Mon Sep 17 00:00:00 2001 From: harish-berri Date: Tue, 21 Apr 2026 22:53:54 +0000 Subject: [PATCH 03/28] Implement CacheCodec for DualCache serialization and deserialization. Add attach_redis_cache method to DualCache for lazy Redis integration. Update RedisCache to handle None keys and improve logging. Enhance user_api_key_auth caching logic and introduce tests for CacheCodec functionality. --- litellm/caching/dual_cache.py | 19 ++++ litellm/caching/redis_cache.py | 10 ++- litellm/proxy/auth/auth_checks.py | 34 +++++--- litellm/proxy/auth/user_api_key_auth.py | 8 +- .../common_utils/cache_pydantic_utils.py | 86 +++++++++++++++++++ litellm/proxy/proxy_server.py | 41 +++++++-- tests/local_testing/test_dual_cache.py | 69 +++++++++++++++ .../proxy/common_utils/test_cache_codec.py | 85 ++++++++++++++++++ 8 files changed, 330 insertions(+), 22 deletions(-) create mode 100644 litellm/proxy/common_utils/cache_pydantic_utils.py create mode 100644 tests/test_litellm/proxy/common_utils/test_cache_codec.py diff --git a/litellm/caching/dual_cache.py b/litellm/caching/dual_cache.py index 34ae3638a5bf..7436f95f5b40 100644 --- a/litellm/caching/dual_cache.py +++ b/litellm/caching/dual_cache.py @@ -92,6 +92,25 @@ def update_cache_ttl( if default_redis_ttl is not None: self.default_redis_ttl = default_redis_ttl + def attach_redis_cache( + self, + redis_cache: Optional[RedisCache] = None, + *, + default_redis_ttl: Optional[float] = None, + ) -> None: + """ + Attach a Redis backend if this DualCache does not already have one. + + No-op when ``redis_cache`` is None or when Redis was already set (constructor + or a prior attach). Use this for lazy wiring after a shared Redis client exists. + Does not backfill in-memory-only keys to Redis. + """ + if redis_cache is None or self.redis_cache is not None: + return + self.redis_cache = redis_cache + if default_redis_ttl is not None: + self.default_redis_ttl = default_redis_ttl + def set_cache(self, key, value, local_only: bool = False, **kwargs): # Update both Redis and in-memory cache try: diff --git a/litellm/caching/redis_cache.py b/litellm/caching/redis_cache.py index 84a2887f527b..f8abb9e95c80 100644 --- a/litellm/caching/redis_cache.py +++ b/litellm/caching/redis_cache.py @@ -551,6 +551,13 @@ async def script_callable(keys: List[str], args: List[Any]) -> Any: async def async_set_cache(self, key, value, **kwargs): from redis.asyncio import Redis + if key is None: + verbose_logger.debug( + "LiteLLM Redis Caching: async set() skipped — key is None, value=%r", + value, + ) + return None + start_time = time.time() try: _redis_client: Redis = self.init_async_client() # type: ignore @@ -569,8 +576,9 @@ async def async_set_cache(self, key, value, **kwargs): ) ) verbose_logger.error( - "LiteLLM Redis Caching: async set() - Got exception from REDIS %s, Writing value=%s", + "LiteLLM Redis Caching: async set() - Got exception from REDIS %s, key=%r, value=%r", str(e), + key, value, ) raise e diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 290b34be836b..1577b5bda9ce 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -12,7 +12,7 @@ import asyncio import re import time -from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Union, cast +from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Type, Union, cast from fastapi import HTTPException, Request, status from pydantic import BaseModel @@ -65,6 +65,7 @@ TOOL_CAPABLE_CALL_TYPES, extract_request_tool_names, ) +from litellm.proxy.common_utils.cache_pydantic_utils import CacheCodec from litellm.proxy.route_llm_request import route_request from litellm.proxy.utils import PrismaClient, ProxyLogging, log_db_metrics from litellm.router import Router @@ -1435,13 +1436,23 @@ async def get_user_object( async def _cache_management_object( key: str, - value: BaseModel, + value: Union[BaseModel, Dict[str, Any]], user_api_key_cache: DualCache, proxy_logging_obj: Optional[ProxyLogging], + *, + model_type: Type[BaseModel], ): + """ + Persist management objects to DualCache (in-memory + optional Redis). + + Values must be JSON-serializable for the Redis path (``json.dumps``). Payloads are + encoded with ``CacheCodec.serialize(..., model_type=...)`` so writes match reads + via ``CacheCodec.deserialize(..., model_type)``. + """ + cache_payload = CacheCodec.serialize(value, model_type=model_type) await user_api_key_cache.async_set_cache( key=key, - value=value, + value=cache_payload, ttl=DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL, ) @@ -1462,6 +1473,7 @@ async def _cache_team_object( value=team_table, user_api_key_cache=user_api_key_cache, proxy_logging_obj=proxy_logging_obj, + model_type=LiteLLM_TeamTableCachedObj, ) @@ -1481,6 +1493,7 @@ async def _cache_key_object( value=user_api_key_obj, user_api_key_cache=user_api_key_cache, proxy_logging_obj=proxy_logging_obj, + model_type=UserAPIKeyAuth, ) @@ -2196,15 +2209,13 @@ async def get_key_object( # check if in cache key = hashed_token - cached_key_obj: Optional[UserAPIKeyAuth] = await user_api_key_cache.async_get_cache( - key=key - ) - + # Same flow as before: use cache only when we have a hit we can turn into UserAPIKeyAuth + # (dict from Redis / model_dump, or UserAPIKeyAuth from in-memory). Otherwise fall through to DB. + cached_key_obj = await user_api_key_cache.async_get_cache(key=key) if cached_key_obj is not None: - if isinstance(cached_key_obj, dict): - return UserAPIKeyAuth(**cached_key_obj) - elif isinstance(cached_key_obj, UserAPIKeyAuth): - return cached_key_obj + user_api_key_auth = CacheCodec.deserialize(cached_key_obj, UserAPIKeyAuth) + if user_api_key_auth is not None: + return user_api_key_auth if check_cache_only: raise Exception( @@ -3465,6 +3476,7 @@ async def get_project_object( value=project_obj, user_api_key_cache=user_api_key_cache, proxy_logging_obj=proxy_logging_obj, + model_type=LiteLLM_ProjectTableCachedObj, ) return project_obj diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 205e61767508..a6c8fd02ee10 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -1452,9 +1452,11 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 else: _team_obj = None - await user_api_key_cache.async_set_cache( - key=valid_token.team_id, value=_team_obj - ) # save team table in cache - used for tpm/rpm limiting - tpm_rpm_limiter.py + # Only cache when the key is a real team_id (non-team keys must not use key=None). + if valid_token.team_id is not None and _team_obj is not None: + await user_api_key_cache.async_set_cache( + key=valid_token.team_id, value=_team_obj + ) # save team table in cache - used for tpm/rpm limiting - tpm_rpm_limiter.py # Fetch project object if key belongs to a project _project_obj = None diff --git a/litellm/proxy/common_utils/cache_pydantic_utils.py b/litellm/proxy/common_utils/cache_pydantic_utils.py new file mode 100644 index 000000000000..af0c3eef18d7 --- /dev/null +++ b/litellm/proxy/common_utils/cache_pydantic_utils.py @@ -0,0 +1,86 @@ +""" +DualCache presents a single API for reads and writes, but the two backends behave +differently: the in-memory layer can store arbitrary Python objects (including live +``BaseModel`` instances), while Redis persists strings and therefore needs JSON-safe +payloads (``json.dumps`` on the Redis side). + +Call sites therefore see cache ``value`` / ``cached`` as effectively ``Any``: the same +key may deserialize to a model on one process (memory hit) or to a ``dict`` after a +Redis round-trip. ``CacheCodec`` centralizes encode/decode at that boundary: +``CacheCodec.serialize`` before ``set``, ``CacheCodec.deserialize`` after ``get`` +when you need a typed ``BaseModel``. + +``dataclasses`` are not supported: only ``dict`` and Pydantic ``BaseModel`` inputs +are encoded; pass a Pydantic model or convert with e.g. ``dataclasses.asdict`` first. +""" + +from __future__ import annotations + +from typing import Any, Optional, Type, TypeVar + +from pydantic import BaseModel, ValidationError + +from litellm._logging import verbose_proxy_logger + +T = TypeVar("T", bound=BaseModel) + + +class CacheCodec: + """ + Encode/decode Pydantic models for DualCache (memory vs Redis safe payloads). + + Dataclasses are not supported yet (only ``dict`` and ``BaseModel``). + + Use ``serialize`` with ``model_type`` when writing so the same schema is used + as on read (``deserialize``). Pass ``model_type`` whenever you know it + (validates ``dict`` payloads and normalizes ``BaseModel`` instances). + """ + + @staticmethod + def serialize(value: Any, model_type: Optional[Type[T]] = None) -> Any: + """ + Encode a value for DualCache / Redis (``json.dumps``-safe). + + If ``model_type`` is set, the payload is validated with that model, then + ``model_dump(mode="json", exclude_none=True)`` — symmetric with ``deserialize``. + + If ``model_type`` is omitted, any ``BaseModel`` is dumped as above; other + values (e.g. plain ``dict``) are returned unchanged. + """ + if model_type is not None: + if isinstance(value, (dict, BaseModel)): + return model_type.model_validate(value).model_dump( + mode="json", exclude_none=True + ) + return value + if isinstance(value, BaseModel): + return value.model_dump(mode="json", exclude_none=True) + return value + + @staticmethod + def deserialize(cached: Any, model_type: Type[T]) -> Optional[T]: + """ + Decode a cache entry to ``model_type``. + + - ``None`` → ``None`` + - Already an instance of ``model_type`` (including subclasses) → returned as-is + - ``dict`` → ``model_type.model_validate(...)``; on ``ValidationError``, + logs a warning and returns ``None`` (treat as cache miss; avoids serving + malformed or schema-drifted entries) + - Any other type → ``None`` (caller should treat as cache miss or log) + """ + if cached is None: + return None + if isinstance(cached, model_type): + return cached + if isinstance(cached, dict): + try: + return model_type.model_validate(cached) + except ValidationError as e: + verbose_proxy_logger.warning( + "CacheCodec.deserialize: validation failed for %s (%s)", + model_type.__name__, + e, + ) + return None + return None diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index f3ff3f5e23f0..72ec41ef97f3 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -93,6 +93,7 @@ TransformRequestBody, UserAPIKeyAuth, ) +from litellm.proxy.common_utils.cache_pydantic_utils import CacheCodec from litellm.proxy.common_utils.callback_utils import ( normalize_callback_names, process_callback, @@ -1943,14 +1944,22 @@ async def _update_key_cache(token: str, response_cost: float): else: hashed_token = token verbose_proxy_logger.debug("_update_key_cache: hashed_token=%s", hashed_token) - existing_spend_obj: LiteLLM_VerificationTokenView = await user_api_key_cache.async_get_cache(key=hashed_token) # type: ignore + cached_key = await user_api_key_cache.async_get_cache(key=hashed_token) verbose_proxy_logger.debug( - f"_update_key_cache: existing_spend_obj={existing_spend_obj}" + f"_update_key_cache: existing_spend_obj={cached_key}" ) + if cached_key is None: + return + existing_spend_obj = CacheCodec.deserialize(cached_key, UserAPIKeyAuth) if existing_spend_obj is None: + verbose_proxy_logger.warning( + "_update_key_cache: unexpected cached key type %s for hashed_token=%s; skipping spend update", + type(cached_key).__name__, + hashed_token, + ) return - else: - existing_spend = existing_spend_obj.spend + + existing_spend = existing_spend_obj.spend or 0.0 # Calculate the new cost by adding the existing cost and response_cost new_spend = existing_spend + response_cost @@ -2008,9 +2017,14 @@ async def _update_key_cache(token: str, response_cost: float): existing_team_member_spend + response_cost ) - # Update the cost column for the given token + # Update the cost column for the given token (dict for Redis pipeline json.dumps) existing_spend_obj.spend = new_spend - values_to_update_in_cache.append((hashed_token, existing_spend_obj)) + values_to_update_in_cache.append( + ( + hashed_token, + CacheCodec.serialize(existing_spend_obj, model_type=UserAPIKeyAuth), + ) + ) ### UPDATE USER SPEND ### async def _update_user_cache(): @@ -2851,9 +2865,22 @@ def _init_cache( ): ## INIT PROXY REDIS USAGE CLIENT ## redis_usage_cache = litellm.cache.cache - spend_counter_cache.redis_cache = redis_usage_cache + spend_counter_cache.attach_redis_cache( + redis_usage_cache, + default_redis_ttl=litellm.default_redis_ttl, + ) # Note: PKCE verifier storage uses redis_usage_cache directly (not # user_api_key_cache) to avoid routing all API-key lookups through Redis. + # Share the same Redis client for virtual-key lookups (same DualCache as + # model_max_budget_limiter). attach_redis_cache is a no-op if Redis is + # already set (e.g. config reload). + user_api_key_cache.attach_redis_cache( + redis_usage_cache, + default_redis_ttl=litellm.default_redis_ttl, + ) + verbose_proxy_logger.debug( + "Attached redis_usage_cache Redis client to user_api_key_cache" + ) def switch_on_llm_response_caching(self): """ diff --git a/tests/local_testing/test_dual_cache.py b/tests/local_testing/test_dual_cache.py index 5a1cdf864871..30f45033c58e 100644 --- a/tests/local_testing/test_dual_cache.py +++ b/tests/local_testing/test_dual_cache.py @@ -245,3 +245,72 @@ async def test_dual_cache_delete(is_async): result = dual_cache.get_cache(test_key) assert result is None + + +def test_dual_cache_late_attach_redis_wires_writes_and_ttl_sync(): + """ + Typical lazy startup (sync): DualCache runs with in-memory only, then Redis + becomes available and is attached. New writes must reach Redis; keys written + before attach are not backfilled. Optional default_redis_ttl is applied on attach. + """ + in_memory = InMemoryCache() + dual_cache = DualCache(in_memory_cache=in_memory, redis_cache=None) + + mock_redis = MagicMock() + mock_redis.set_cache = MagicMock() + mock_redis.async_set_cache = AsyncMock() + + key_before = f"before_attach_{uuid.uuid4()}" + val_before = {"phase": "memory_only"} + dual_cache.set_cache(key_before, val_before) + + assert in_memory.get_cache(key_before) == val_before + + dual_cache.attach_redis_cache(mock_redis, default_redis_ttl=99.0) + assert dual_cache.redis_cache is mock_redis + assert dual_cache.default_redis_ttl == 99.0 + + mock_redis.set_cache.assert_not_called() + + key_after = f"after_attach_{uuid.uuid4()}" + val_after = {"phase": "memory_and_redis"} + dual_cache.set_cache(key_after, val_after) + mock_redis.set_cache.assert_called_once() + assert mock_redis.set_cache.call_args[0][:2] == (key_after, val_after) + + assert in_memory.get_cache(key_after) == val_after + + +@pytest.mark.asyncio +async def test_dual_cache_late_attach_redis_wires_writes_and_ttl_async(): + """ + Typical lazy startup (async): DualCache runs with in-memory only, then Redis + becomes available and is attached. New writes must reach Redis; keys written + before attach are not backfilled. Optional default_redis_ttl is applied on attach. + """ + in_memory = InMemoryCache() + dual_cache = DualCache(in_memory_cache=in_memory, redis_cache=None) + + mock_redis = MagicMock() + mock_redis.set_cache = MagicMock() + mock_redis.async_set_cache = AsyncMock() + + key_before = f"before_attach_{uuid.uuid4()}" + val_before = {"phase": "memory_only"} + await dual_cache.async_set_cache(key_before, val_before) + + assert in_memory.get_cache(key_before) == val_before + + dual_cache.attach_redis_cache(mock_redis, default_redis_ttl=99.0) + assert dual_cache.redis_cache is mock_redis + assert dual_cache.default_redis_ttl == 99.0 + + mock_redis.async_set_cache.assert_not_called() + + key_after = f"after_attach_{uuid.uuid4()}" + val_after = {"phase": "memory_and_redis"} + await dual_cache.async_set_cache(key_after, val_after) + mock_redis.async_set_cache.assert_called_once() + assert mock_redis.async_set_cache.call_args[0][:2] == (key_after, val_after) + + assert in_memory.get_cache(key_after) == val_after diff --git a/tests/test_litellm/proxy/common_utils/test_cache_codec.py b/tests/test_litellm/proxy/common_utils/test_cache_codec.py new file mode 100644 index 000000000000..e99f22a75d4b --- /dev/null +++ b/tests/test_litellm/proxy/common_utils/test_cache_codec.py @@ -0,0 +1,85 @@ +import logging +from typing import Optional + +import pytest +from pydantic import BaseModel, ValidationError + +from litellm.proxy.common_utils.cache_pydantic_utils import CacheCodec + + +class _SampleModel(BaseModel): + name: str + count: Optional[int] = None + + +class _SampleSubModel(_SampleModel): + pass + + +class TestCacheCodecSerialize: + def test_without_model_type_base_model_dumped_json_safe(self): + m = _SampleModel(name="a", count=1) + out = CacheCodec.serialize(m) + assert out == {"name": "a", "count": 1} + + def test_without_model_type_dict_unchanged(self): + d = {"name": "x"} + assert CacheCodec.serialize(d) is d + + def test_without_model_type_primitive_unchanged(self): + assert CacheCodec.serialize(42) == 42 + + def test_with_model_type_dict_validated_and_dumped(self): + out = CacheCodec.serialize({"name": "b", "count": 2}, model_type=_SampleModel) + assert out == {"name": "b", "count": 2} + + def test_with_model_type_base_model_validated_and_dumped(self): + m = _SampleModel(name="c", count=None) + out = CacheCodec.serialize(m, model_type=_SampleModel) + assert out == {"name": "c"} + + def test_with_model_type_exclude_none_on_dump(self): + out = CacheCodec.serialize({"name": "d"}, model_type=_SampleModel) + assert out == {"name": "d"} + assert "count" not in out + + def test_with_model_type_non_dict_non_model_passthrough(self): + assert CacheCodec.serialize("raw", model_type=_SampleModel) == "raw" + + def test_with_model_type_invalid_dict_raises(self): + with pytest.raises(ValidationError): + CacheCodec.serialize({"count": 1}, model_type=_SampleModel) + + +class TestCacheCodecDeserialize: + def test_none_returns_none(self): + assert CacheCodec.deserialize(None, _SampleModel) is None + + def test_dict_validates_to_model(self): + m = CacheCodec.deserialize({"name": "e", "count": 3}, _SampleModel) + assert isinstance(m, _SampleModel) + assert m.name == "e" + assert m.count == 3 + + def test_instance_same_type_returned_as_is(self): + original = _SampleModel(name="f") + m = CacheCodec.deserialize(original, _SampleModel) + assert m is original + + def test_subclass_instance_accepted(self): + sub = _SampleSubModel(name="g") + m = CacheCodec.deserialize(sub, _SampleModel) + assert m is sub + + def test_wrong_type_returns_none(self): + assert CacheCodec.deserialize("not-a-dict", _SampleModel) is None + + def test_invalid_dict_returns_none_and_logs_warning(self, caplog): + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + out = CacheCodec.deserialize({"count": 1}, _SampleModel) + assert out is None + assert any( + "CacheCodec.deserialize" in r.message and "_SampleModel" in r.message + for r in caplog.records + if r.levelno >= logging.WARNING + ), f"Expected deserialize validation warning. Records: {[r.message for r in caplog.records]}" From 9cfc5fc20951ddf7d2aae0e409fe9c1b5d458cc4 Mon Sep 17 00:00:00 2001 From: harish-berri Date: Tue, 21 Apr 2026 23:30:46 +0000 Subject: [PATCH 04/28] update tests --- tests/local_testing/test_dual_cache.py | 69 ------------------ tests/test_litellm/caching/test_dual_cache.py | 70 +++++++++++++++++++ 2 files changed, 70 insertions(+), 69 deletions(-) diff --git a/tests/local_testing/test_dual_cache.py b/tests/local_testing/test_dual_cache.py index 30f45033c58e..5a1cdf864871 100644 --- a/tests/local_testing/test_dual_cache.py +++ b/tests/local_testing/test_dual_cache.py @@ -245,72 +245,3 @@ async def test_dual_cache_delete(is_async): result = dual_cache.get_cache(test_key) assert result is None - - -def test_dual_cache_late_attach_redis_wires_writes_and_ttl_sync(): - """ - Typical lazy startup (sync): DualCache runs with in-memory only, then Redis - becomes available and is attached. New writes must reach Redis; keys written - before attach are not backfilled. Optional default_redis_ttl is applied on attach. - """ - in_memory = InMemoryCache() - dual_cache = DualCache(in_memory_cache=in_memory, redis_cache=None) - - mock_redis = MagicMock() - mock_redis.set_cache = MagicMock() - mock_redis.async_set_cache = AsyncMock() - - key_before = f"before_attach_{uuid.uuid4()}" - val_before = {"phase": "memory_only"} - dual_cache.set_cache(key_before, val_before) - - assert in_memory.get_cache(key_before) == val_before - - dual_cache.attach_redis_cache(mock_redis, default_redis_ttl=99.0) - assert dual_cache.redis_cache is mock_redis - assert dual_cache.default_redis_ttl == 99.0 - - mock_redis.set_cache.assert_not_called() - - key_after = f"after_attach_{uuid.uuid4()}" - val_after = {"phase": "memory_and_redis"} - dual_cache.set_cache(key_after, val_after) - mock_redis.set_cache.assert_called_once() - assert mock_redis.set_cache.call_args[0][:2] == (key_after, val_after) - - assert in_memory.get_cache(key_after) == val_after - - -@pytest.mark.asyncio -async def test_dual_cache_late_attach_redis_wires_writes_and_ttl_async(): - """ - Typical lazy startup (async): DualCache runs with in-memory only, then Redis - becomes available and is attached. New writes must reach Redis; keys written - before attach are not backfilled. Optional default_redis_ttl is applied on attach. - """ - in_memory = InMemoryCache() - dual_cache = DualCache(in_memory_cache=in_memory, redis_cache=None) - - mock_redis = MagicMock() - mock_redis.set_cache = MagicMock() - mock_redis.async_set_cache = AsyncMock() - - key_before = f"before_attach_{uuid.uuid4()}" - val_before = {"phase": "memory_only"} - await dual_cache.async_set_cache(key_before, val_before) - - assert in_memory.get_cache(key_before) == val_before - - dual_cache.attach_redis_cache(mock_redis, default_redis_ttl=99.0) - assert dual_cache.redis_cache is mock_redis - assert dual_cache.default_redis_ttl == 99.0 - - mock_redis.async_set_cache.assert_not_called() - - key_after = f"after_attach_{uuid.uuid4()}" - val_after = {"phase": "memory_and_redis"} - await dual_cache.async_set_cache(key_after, val_after) - mock_redis.async_set_cache.assert_called_once() - assert mock_redis.async_set_cache.call_args[0][:2] == (key_after, val_after) - - assert in_memory.get_cache(key_after) == val_after diff --git a/tests/test_litellm/caching/test_dual_cache.py b/tests/test_litellm/caching/test_dual_cache.py index 8e502175761a..647747262019 100644 --- a/tests/test_litellm/caching/test_dual_cache.py +++ b/tests/test_litellm/caching/test_dual_cache.py @@ -1,5 +1,6 @@ import asyncio import time +import uuid from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -260,3 +261,72 @@ async def test_async_increment_cache_returns_none_when_no_in_memory_cache_and_re f"Expected None when in_memory_cache is absent and Redis fails, got {result!r}. " "Returning the delta (1.0) would silently miscalculate rate-limit counters." ) + + +def test_dual_cache_late_attach_redis_wires_writes_and_ttl_sync(): + """ + Typical lazy startup (sync): DualCache runs with in-memory only, then Redis + becomes available and is attached. New writes must reach Redis; keys written + before attach are not backfilled. Optional default_redis_ttl is applied on attach. + """ + in_memory = InMemoryCache() + dual_cache = DualCache(in_memory_cache=in_memory, redis_cache=None) + + mock_redis = MagicMock() + mock_redis.set_cache = MagicMock() + mock_redis.async_set_cache = AsyncMock() + + key_before = f"before_attach_{uuid.uuid4()}" + val_before = {"phase": "memory_only"} + dual_cache.set_cache(key_before, val_before) + + assert in_memory.get_cache(key_before) == val_before + + dual_cache.attach_redis_cache(mock_redis, default_redis_ttl=99.0) + assert dual_cache.redis_cache is mock_redis + assert dual_cache.default_redis_ttl == 99.0 + + mock_redis.set_cache.assert_not_called() + + key_after = f"after_attach_{uuid.uuid4()}" + val_after = {"phase": "memory_and_redis"} + dual_cache.set_cache(key_after, val_after) + mock_redis.set_cache.assert_called_once() + assert mock_redis.set_cache.call_args[0][:2] == (key_after, val_after) + + assert in_memory.get_cache(key_after) == val_after + + +@pytest.mark.asyncio +async def test_dual_cache_late_attach_redis_wires_writes_and_ttl_async(): + """ + Typical lazy startup (async): DualCache runs with in-memory only, then Redis + becomes available and is attached. New writes must reach Redis; keys written + before attach are not backfilled. Optional default_redis_ttl is applied on attach. + """ + in_memory = InMemoryCache() + dual_cache = DualCache(in_memory_cache=in_memory, redis_cache=None) + + mock_redis = MagicMock() + mock_redis.set_cache = MagicMock() + mock_redis.async_set_cache = AsyncMock() + + key_before = f"before_attach_{uuid.uuid4()}" + val_before = {"phase": "memory_only"} + await dual_cache.async_set_cache(key_before, val_before) + + assert in_memory.get_cache(key_before) == val_before + + dual_cache.attach_redis_cache(mock_redis, default_redis_ttl=99.0) + assert dual_cache.redis_cache is mock_redis + assert dual_cache.default_redis_ttl == 99.0 + + mock_redis.async_set_cache.assert_not_called() + + key_after = f"after_attach_{uuid.uuid4()}" + val_after = {"phase": "memory_and_redis"} + await dual_cache.async_set_cache(key_after, val_after) + mock_redis.async_set_cache.assert_called_once() + assert mock_redis.async_set_cache.call_args[0][:2] == (key_after, val_after) + + assert in_memory.get_cache(key_after) == val_after From 595f42d22a81321d46ffe335f2a70347e77d7bcf Mon Sep 17 00:00:00 2001 From: harish-berri Date: Thu, 23 Apr 2026 23:42:03 +0000 Subject: [PATCH 05/28] Refactor caching logic in auth_checks and user_api_key_auth to utilize CacheCodec for serialization and deserialization. Simplify cache retrieval by removing unnecessary type checks and streamline cache storage with consistent key formatting. --- litellm/proxy/auth/auth_checks.py | 25 +++++++++++-------------- litellm/proxy/auth/user_api_key_auth.py | 9 ++++++++- 2 files changed, 19 insertions(+), 15 deletions(-) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 1577b5bda9ce..760c0fb12d6a 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -1627,13 +1627,10 @@ async def _get_team_object_from_cache( if cached_team_obj is None: cached_team_obj = await user_api_key_cache.async_get_cache(key=key) - if cached_team_obj is not None: - if isinstance(cached_team_obj, dict): - return LiteLLM_TeamTableCachedObj(**cached_team_obj) - elif isinstance(cached_team_obj, LiteLLM_TeamTableCachedObj): - return cached_team_obj - - return None + if cached_team_obj is None: + return None + + return CacheCodec.deserialize(cached_team_obj, LiteLLM_TeamTableCachedObj) async def get_team_object( @@ -1709,9 +1706,12 @@ async def _cache_access_object( proxy_logging_obj: Optional[ProxyLogging] = None, ): key = "access_group_id:{}".format(access_group_id) + cache_payload = CacheCodec.serialize( + access_group_table, model_type=LiteLLM_AccessGroupTable + ) await user_api_key_cache.async_set_cache( key=key, - value=access_group_table, + value=cache_payload, ttl=DEFAULT_ACCESS_GROUP_CACHE_TTL, ) @@ -1758,13 +1758,10 @@ async def get_access_object( key = "access_group_id:{}".format(access_group_id) - # Always check cache first cached_access_obj = await user_api_key_cache.async_get_cache(key=key) - if cached_access_obj is not None: - if isinstance(cached_access_obj, dict): - return LiteLLM_AccessGroupTable(**cached_access_obj) - elif isinstance(cached_access_obj, LiteLLM_AccessGroupTable): - return cached_access_obj + deserialized = CacheCodec.deserialize(cached_access_obj, LiteLLM_AccessGroupTable) + if deserialized is not None: + return deserialized # Not in cache - fetch from DB try: diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index a6c8fd02ee10..855ce50ec0cf 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -58,6 +58,7 @@ from litellm.proxy.auth.oauth2_proxy_hook import handle_oauth2_proxy_request from litellm.proxy.auth.route_checks import RouteChecks from litellm.proxy.common_utils.cache_coordinator import EventDrivenCacheCoordinator +from litellm.proxy.common_utils.cache_pydantic_utils import CacheCodec from litellm.proxy.common_utils.http_parsing_utils import ( _read_request_body, _safe_get_request_headers, @@ -1454,8 +1455,14 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 # Only cache when the key is a real team_id (non-team keys must not use key=None). if valid_token.team_id is not None and _team_obj is not None: + # Match get_team_object / spend counters: "team_id:{id}". Serialize for Redis + # (json.dumps) — same as _cache_team_object in auth_checks. + team_cache_key = f"team_id:{valid_token.team_id}" await user_api_key_cache.async_set_cache( - key=valid_token.team_id, value=_team_obj + key=team_cache_key, + value=CacheCodec.serialize( + _team_obj, model_type=LiteLLM_TeamTableCachedObj + ), ) # save team table in cache - used for tpm/rpm limiting - tpm_rpm_limiter.py # Fetch project object if key belongs to a project From d4a26ff364a85980cdef5cae86beb77c6628b167 Mon Sep 17 00:00:00 2001 From: harish-berri Date: Fri, 24 Apr 2026 01:30:01 +0000 Subject: [PATCH 06/28] Enhance caching mechanism by integrating CacheCodec for serialization across various components. Introduce the enable_redis_auth_cache flag to control Redis integration for user_api_key_cache, improving performance in multi-worker deployments. Update documentation and tests to reflect these changes. --- .gitignore | 3 +- docs/my-website/docs/proxy/config_settings.md | 61 +++++++ litellm/proxy/_types.py | 13 ++ litellm/proxy/auth/auth_checks.py | 156 ++++++++++------- litellm/proxy/auth/handle_jwt.py | 9 +- litellm/proxy/auth/user_api_key_auth.py | 26 ++- .../common_utils/cache_pydantic_utils.py | 7 + .../access_group_endpoints.py | 19 +- litellm/proxy/management_endpoints/ui_sso.py | 7 +- litellm/proxy/proxy_server.py | 164 ++++++++++-------- .../proxy/common_utils/test_cache_codec.py | 36 ++++ .../proxy/test_redis_auth_cache_flag.py | 130 ++++++++++++++ 12 files changed, 465 insertions(+), 166 deletions(-) create mode 100644 tests/test_litellm/proxy/test_redis_auth_cache_flag.py diff --git a/.gitignore b/.gitignore index f5591b7f2c6c..59812ed6ed49 100644 --- a/.gitignore +++ b/.gitignore @@ -99,4 +99,5 @@ STABILIZATION_TODO.md **/test-results **/playwright-report **/*.storageState.json -**/coverage \ No newline at end of file +**/coverage +test-config \ No newline at end of file diff --git a/docs/my-website/docs/proxy/config_settings.md b/docs/my-website/docs/proxy/config_settings.md index be819a04ca11..7187690375db 100644 --- a/docs/my-website/docs/proxy/config_settings.md +++ b/docs/my-website/docs/proxy/config_settings.md @@ -307,9 +307,70 @@ router_settings: | token_rate_limit_type | string | Rate limit counting method: "total", "output", or "input" tokens | | use_redis_transaction_buffer | boolean | If true, buffers database transactions in Redis before writing | | use_shared_health_check | boolean | If true, uses Redis-backed shared health check state across multiple proxy instances | +| enable_redis_auth_cache | boolean | **[Beta]** When `true`, attaches Redis to the virtual-key auth cache (`user_api_key_cache`) so all proxy workers/pods share the same cache instead of each pod resolving keys independently against the database. Requires `litellm_settings.cache: true` with a Redis backend. Significantly reduces database load in multi-worker deployments by eliminating per-pod cache misses on the `combined_view` query. Off by default for a safe phased rollout — enable once your Redis cluster is healthy. Will become the default in a future release. See [Redis Auth Cache](#redis-auth-cache-multi-worker-db-load-reduction). | | user_header_mappings | dict | Map custom request headers to user IDs using lookup rules | | user_header_name | string | HTTP header name to extract user identity from requests | +## Redis Auth Cache — Multi-Worker DB Load Reduction + +### Problem + +In multi-worker or multi-pod deployments each worker process keeps its own **in-memory** virtual-key cache (`user_api_key_cache`). When a request is routed to a worker that has not seen a key before, it runs a `combined_view` SQL query against the database (8+ JOINs). With many workers and a large number of unique keys this causes: + +- Redundant DB queries on every pod's cold-start +- Sustained CPU spikes on the database (especially visible in RDS Performance Insights) +- Increased p99 latency for requests that miss the local cache + +### Solution + +Setting `enable_redis_auth_cache: true` attaches Redis to `user_api_key_cache` so the resolved key object is stored in a **shared** Redis cache. A cache hit on any worker prevents the DB query entirely. + +```yaml +# config.yaml + +litellm_settings: + cache: true + cache_params: + type: redis + host: os.environ/REDIS_HOST + port: os.environ/REDIS_PORT + +general_settings: + master_key: sk-1234 + enable_redis_auth_cache: true # ← share the auth cache across workers +``` + +### Requirements + +| Requirement | Notes | +|---|---| +| `litellm_settings.cache: true` | Redis must be configured as the cache backend | +| Redis cluster is healthy | Auth lookups now depend on Redis availability — monitor it | +| LiteLLM ≥ version with CacheCodec support | All read/write paths now use `CacheCodec` for safe serialisation across Redis round-trips | + +### Rollout recommendation + +1. Deploy with `enable_redis_auth_cache: false` (default) and baseline DB CPU. +2. Enable on a canary pod first; confirm cache-hit rate via proxy debug logs (`LITELLM_LOG=DEBUG`). +3. Roll out to all pods; monitor DB CPU drop. +4. The flag will be removed and the behaviour made permanent in a future LiteLLM release. + +### Debugging + +Set `LITELLM_LOG=DEBUG` and look for: + +``` +enable_redis_auth_cache=True: attached Redis to user_api_key_cache — virtual-key lookups are now shared across all proxy workers. +``` + +If the flag is off you will see: + +``` +enable_redis_auth_cache is not set: user_api_key_cache remains in-memory only (per-worker). Set general_settings.enable_redis_auth_cache: true to share the auth cache across workers and reduce DB load. +``` + +--- + ### router_settings - Reference :::info diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index a592a81ed868..960261d5cf10 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -2358,6 +2358,19 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase): None, description="List of MCP server fields that must be filled in for a submission to pass standards checks (e.g. ['description', 'source_url', 'alias']).", ) + enable_redis_auth_cache: Optional[bool] = Field( + None, + description=( + "When True, attaches Redis to user_api_key_cache so virtual-key lookups " + "are shared across all proxy workers/pods instead of being resolved " + "per-process. Requires a Redis cache to be configured under litellm_settings. " + "Reduces DB load significantly in multi-worker deployments by eliminating " + "redundant combined_view SQL queries caused by per-pod cache misses. " + "Off by default for a safe phased rollout — set to True once your Redis " + "cluster is healthy and the CacheCodec serialisation has been validated in " + "your environment. Will be enabled by default in a future release." + ), + ) class ConfigYAML(LiteLLMPydanticObjectBase): diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 760c0fb12d6a..02d41041e86b 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -823,7 +823,9 @@ async def get_default_end_user_budget( # Check cache first cached_budget = await user_api_key_cache.async_get_cache(key=cache_key) if cached_budget is not None: - return LiteLLM_BudgetTable(**cached_budget) + deserialized = CacheCodec.deserialize(cached_budget, LiteLLM_BudgetTable) + if deserialized is not None: + return deserialized # Fetch from database try: @@ -837,14 +839,15 @@ async def get_default_end_user_budget( ) return None + _budget_obj = LiteLLM_BudgetTable(**budget_record.dict()) # Cache the budget for 60 seconds await user_api_key_cache.async_set_cache( key=cache_key, - value=budget_record.dict(), + value=CacheCodec.serialize(_budget_obj, model_type=LiteLLM_BudgetTable), ttl=DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL, ) - return LiteLLM_BudgetTable(**budget_record.dict()) + return _budget_obj except Exception as e: verbose_proxy_logger.error(f"Error fetching default end user budget: {str(e)}") @@ -960,20 +963,20 @@ async def get_end_user_object( # Check cache first cached_user_obj = await user_api_key_cache.async_get_cache(key=_key) if cached_user_obj is not None: - return_obj = LiteLLM_EndUserTable(**cached_user_obj) - - # Apply default budget if needed - return_obj = await _apply_default_budget_to_end_user( - end_user_obj=return_obj, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - parent_otel_span=parent_otel_span, - ) + return_obj = CacheCodec.deserialize(cached_user_obj, LiteLLM_EndUserTable) + if return_obj is not None: + # Apply default budget if needed + return_obj = await _apply_default_budget_to_end_user( + end_user_obj=return_obj, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=parent_otel_span, + ) - # Check budget limits - _check_end_user_budget(end_user_obj=return_obj, route=route) + # Check budget limits + _check_end_user_budget(end_user_obj=return_obj, route=route) - return return_obj + return return_obj # Fetch from database try: @@ -996,9 +999,10 @@ async def get_end_user_object( parent_otel_span=parent_otel_span, ) - # Save to cache (always store as dict for consistency) + # Save to cache await user_api_key_cache.async_set_cache( - key="end_user_id:{}".format(end_user_id), value=_response.dict() + key="end_user_id:{}".format(end_user_id), + value=CacheCodec.serialize(_response, model_type=LiteLLM_EndUserTable), ) # Check budget limits @@ -1051,10 +1055,11 @@ async def get_tag_objects_batch( cache_key = f"tag:{tag_name}" cached_tag = await user_api_key_cache.async_get_cache(key=cache_key) if cached_tag is not None: - if isinstance(cached_tag, dict): - tag_objects[tag_name] = LiteLLM_TagTable(**cached_tag) + deserialized_tag = CacheCodec.deserialize(cached_tag, LiteLLM_TagTable) + if deserialized_tag is not None: + tag_objects[tag_name] = deserialized_tag else: - tag_objects[tag_name] = cached_tag + uncached_tags.append(tag_name) else: uncached_tags.append(tag_name) @@ -1070,11 +1075,12 @@ async def get_tag_objects_batch( for db_tag in db_tags: tag_name = db_tag.tag_name cache_key = f"tag:{tag_name}" - # Cache with default TTL (same as end_user objects) + _tag_obj = LiteLLM_TagTable(**db_tag.dict()) await user_api_key_cache.async_set_cache( - key=cache_key, value=db_tag.dict() + key=cache_key, + value=CacheCodec.serialize(_tag_obj, model_type=LiteLLM_TagTable), ) - tag_objects[tag_name] = LiteLLM_TagTable(**db_tag.dict()) + tag_objects[tag_name] = _tag_obj except Exception as e: verbose_proxy_logger.debug(f"Error batch fetching tags from database: {e}") @@ -1146,7 +1152,11 @@ async def get_team_membership( # check if in cache cached_membership_obj = await user_api_key_cache.async_get_cache(key=_key) if cached_membership_obj is not None: - return LiteLLM_TeamMembership(**cached_membership_obj) + deserialized_membership = CacheCodec.deserialize( + cached_membership_obj, LiteLLM_TeamMembership + ) + if deserialized_membership is not None: + return deserialized_membership # else, check db try: @@ -1158,10 +1168,11 @@ async def get_team_membership( if response is None: return None - # save the team membership object to cache (store as dict) - await user_api_key_cache.async_set_cache(key=_key, value=response.dict()) - _response = LiteLLM_TeamMembership(**response.dict()) + await user_api_key_cache.async_set_cache( + key=_key, + value=CacheCodec.serialize(_response, model_type=LiteLLM_TeamMembership), + ) return _response except Exception: @@ -1350,10 +1361,11 @@ async def get_user_object( if not check_db_only: cached_user_obj = await user_api_key_cache.async_get_cache(key=user_id) if cached_user_obj is not None: - if isinstance(cached_user_obj, dict): - return LiteLLM_UserTable(**cached_user_obj) - elif isinstance(cached_user_obj, LiteLLM_UserTable): - return cached_user_obj + deserialized_user = CacheCodec.deserialize( + cached_user_obj, LiteLLM_UserTable + ) + if deserialized_user is not None: + return deserialized_user # else, check db if prisma_client is None: raise Exception("No db connected") @@ -1415,7 +1427,7 @@ async def get_user_object( # save the user object to cache await user_api_key_cache.async_set_cache( key=user_id, - value=response_dict, + value=CacheCodec.serialize(_response, model_type=LiteLLM_UserTable), ttl=DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL, ) @@ -1886,16 +1898,19 @@ async def get_team_object_by_alias( ) # Cache the result by both alias and team_id + _serialized_team = CacheCodec.serialize( + team_obj, model_type=LiteLLM_TeamTableCachedObj + ) await user_api_key_cache.async_set_cache( key=cache_key, - value=team_obj, + value=_serialized_team, ttl=DEFAULT_IN_MEMORY_TTL, ) # Also cache by team_id for consistency team_id_cache_key = "team_id:{}".format(team_obj.team_id) await user_api_key_cache.async_set_cache( key=team_id_cache_key, - value=team_obj, + value=_serialized_team, ttl=DEFAULT_IN_MEMORY_TTL, ) @@ -1946,10 +1961,11 @@ async def get_org_object_by_alias( cache_key = "org_alias:{}".format(org_alias) cached_org_obj = await user_api_key_cache.async_get_cache(key=cache_key) if cached_org_obj is not None: - if isinstance(cached_org_obj, dict): - return LiteLLM_OrganizationTable(**cached_org_obj) - elif isinstance(cached_org_obj, LiteLLM_OrganizationTable): - return cached_org_obj + deserialized_org = CacheCodec.deserialize( + cached_org_obj, LiteLLM_OrganizationTable + ) + if deserialized_org is not None: + return deserialized_org # Query database by organization_alias try: @@ -1976,16 +1992,19 @@ async def get_org_object_by_alias( org = orgs[0] org_obj = LiteLLM_OrganizationTable(**org.model_dump()) + _serialized_org = CacheCodec.serialize( + org_obj, model_type=LiteLLM_OrganizationTable + ) # Cache the result await user_api_key_cache.async_set_cache( key=cache_key, - value=org_obj.model_dump(), + value=_serialized_org, ttl=DEFAULT_IN_MEMORY_TTL, ) # Also cache by org_id for consistency await user_api_key_cache.async_set_cache( key="org_id:{}".format(org_obj.organization_id), - value=org_obj.model_dump(), + value=_serialized_org, ttl=DEFAULT_IN_MEMORY_TTL, ) @@ -2287,10 +2306,11 @@ async def get_object_permission( key = "object_permission_id:{}".format(object_permission_id) cached_obj_permission = await user_api_key_cache.async_get_cache(key=key) if cached_obj_permission is not None: - if isinstance(cached_obj_permission, dict): - return LiteLLM_ObjectPermissionTable(**cached_obj_permission) - elif isinstance(cached_obj_permission, LiteLLM_ObjectPermissionTable): - return cached_obj_permission + deserialized_perm = CacheCodec.deserialize( + cached_obj_permission, LiteLLM_ObjectPermissionTable + ) + if deserialized_perm is not None: + return deserialized_perm # else, check db try: @@ -2301,14 +2321,16 @@ async def get_object_permission( if response is None: return None - # save the object permission to cache + _perm_obj = LiteLLM_ObjectPermissionTable(**response.dict()) await user_api_key_cache.async_set_cache( key=key, - value=response.model_dump(), + value=CacheCodec.serialize( + _perm_obj, model_type=LiteLLM_ObjectPermissionTable + ), ttl=DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL, ) - return LiteLLM_ObjectPermissionTable(**response.dict()) + return _perm_obj except Exception: return None @@ -2339,10 +2361,11 @@ async def get_managed_vector_store_rows_by_uuids( key = "managed_vector_store_id:{}".format(uuid) cached = await user_api_key_cache.async_get_cache(key=key) if cached is not None: - if isinstance(cached, dict): - result.append(LiteLLM_ManagedVectorStoresTable(**cached)) - elif isinstance(cached, LiteLLM_ManagedVectorStoresTable): - result.append(cached) + deserialized_vs = CacheCodec.deserialize( + cached, LiteLLM_ManagedVectorStoresTable + ) + if deserialized_vs is not None: + result.append(deserialized_vs) else: cache_misses.append(uuid) else: @@ -2370,7 +2393,9 @@ async def get_managed_vector_store_rows_by_uuids( key = "managed_vector_store_id:{}".format(cached_obj.vector_store_id) await user_api_key_cache.async_set_cache( key=key, - value=row_dict, + value=CacheCodec.serialize( + cached_obj, model_type=LiteLLM_ManagedVectorStoresTable + ), ttl=DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL, ) result.append(cached_obj) @@ -2413,12 +2438,13 @@ async def get_org_object( cache_key = "org_id:{}:with_budget".format(org_id) # check if in cache - cached_org_obj = user_api_key_cache.async_get_cache(key=cache_key) + cached_org_obj = await user_api_key_cache.async_get_cache(key=cache_key) if cached_org_obj is not None: - if isinstance(cached_org_obj, dict): - return LiteLLM_OrganizationTable(**cached_org_obj) - elif isinstance(cached_org_obj, LiteLLM_OrganizationTable): - return cached_org_obj + deserialized_org = CacheCodec.deserialize( + cached_org_obj, LiteLLM_OrganizationTable + ) + if deserialized_org is not None: + return deserialized_org # else, check db try: query_kwargs: Dict[str, Any] = {"where": {"organization_id": org_id}} @@ -2432,16 +2458,15 @@ async def get_org_object( if response is None: raise Exception + _org_obj = LiteLLM_OrganizationTable(**response.model_dump()) # Cache the result await user_api_key_cache.async_set_cache( key=cache_key, - value=( - response.model_dump() if hasattr(response, "model_dump") else response - ), + value=CacheCodec.serialize(_org_obj, model_type=LiteLLM_OrganizationTable), ttl=DEFAULT_IN_MEMORY_TTL, ) - return response + return _org_obj except Exception: raise Exception( f"Organization doesn't exist in db. Organization={org_id}. Create organization via `/organization/new` call." @@ -3451,10 +3476,11 @@ async def get_project_object( cache_key = "project_id:{}".format(project_id) cached_obj = await user_api_key_cache.async_get_cache(key=cache_key) if cached_obj is not None: - if isinstance(cached_obj, dict): - return LiteLLM_ProjectTableCachedObj(**cached_obj) - elif isinstance(cached_obj, LiteLLM_ProjectTableCachedObj): - return cached_obj + deserialized_project = CacheCodec.deserialize( + cached_obj, LiteLLM_ProjectTableCachedObj + ) + if deserialized_project is not None: + return deserialized_project # Fetch from DB project_row = await prisma_client.db.litellm_projecttable.find_unique( diff --git a/litellm/proxy/auth/handle_jwt.py b/litellm/proxy/auth/handle_jwt.py index 4a6856b6d14f..07d944ae9319 100644 --- a/litellm/proxy/auth/handle_jwt.py +++ b/litellm/proxy/auth/handle_jwt.py @@ -45,6 +45,7 @@ UserAPIKeyAuth, ) from litellm.proxy.auth.auth_checks import can_team_access_model +from litellm.proxy.common_utils.cache_pydantic_utils import CacheCodec from litellm.proxy.utils import PrismaClient, ProxyLogging from .auth_checks import ( @@ -1376,7 +1377,9 @@ async def sync_user_role_and_teams( if user_api_key_cache is not None: await user_api_key_cache.async_set_cache( key=user_object.user_id, - value=user_object.model_dump(), + value=CacheCodec.serialize( + user_object, model_type=LiteLLM_UserTable + ), ttl=DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL, ) @@ -1399,7 +1402,9 @@ async def sync_user_role_and_teams( if user_api_key_cache is not None: await user_api_key_cache.async_set_cache( key=user_object.user_id, - value=user_object.model_dump(), + value=CacheCodec.serialize( + user_object, model_type=LiteLLM_UserTable + ), ttl=DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL, ) return None diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 855ce50ec0cf..3fe8863e2672 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -1297,27 +1297,39 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 if prisma_client is not None: _cache_key = f"{valid_token.team_id}_{valid_token.user_id}" - team_member_info = await user_api_key_cache.async_get_cache( + _cached_member = await user_api_key_cache.async_get_cache( key=_cache_key ) + team_member_info: Optional[LiteLLM_TeamMembership] = ( + CacheCodec.deserialize(_cached_member, LiteLLM_TeamMembership) + if _cached_member is not None + else None + ) if team_member_info is None: # read from DB _user_id = valid_token.user_id _team_id = valid_token.team_id if _user_id is not None and _team_id is not None: - team_member_info = await prisma_client.db.litellm_teammembership.find_first( + _db_member = await prisma_client.db.litellm_teammembership.find_first( where={ "user_id": _user_id, "team_id": _team_id, }, # type: ignore include={"litellm_budget_table": True}, ) - await user_api_key_cache.async_set_cache( - key=_cache_key, - value=team_member_info, - ttl=5, - ) + if _db_member is not None: + team_member_info = LiteLLM_TeamMembership( + **_db_member.dict() + ) + await user_api_key_cache.async_set_cache( + key=_cache_key, + value=CacheCodec.serialize( + team_member_info, + model_type=LiteLLM_TeamMembership, + ), + ttl=5, + ) if ( team_member_info is not None diff --git a/litellm/proxy/common_utils/cache_pydantic_utils.py b/litellm/proxy/common_utils/cache_pydantic_utils.py index af0c3eef18d7..80a8d6281a12 100644 --- a/litellm/proxy/common_utils/cache_pydantic_utils.py +++ b/litellm/proxy/common_utils/cache_pydantic_utils.py @@ -44,10 +44,17 @@ def serialize(value: Any, model_type: Optional[Type[T]] = None) -> Any: If ``model_type`` is set, the payload is validated with that model, then ``model_dump(mode="json", exclude_none=True)`` — symmetric with ``deserialize``. + If the value is already an instance of ``model_type`` (or a subclass), + ``model_validate`` is skipped to avoid an unnecessary Pydantic copy — the + value is dumped directly. + If ``model_type`` is omitted, any ``BaseModel`` is dumped as above; other values (e.g. plain ``dict``) are returned unchanged. """ if model_type is not None: + if isinstance(value, model_type): + # Already the right type: dump directly, skip re-validation copy. + return value.model_dump(mode="json", exclude_none=True) if isinstance(value, (dict, BaseModel)): return model_type.model_validate(value).model_dump( mode="json", exclude_none=True diff --git a/litellm/proxy/management_endpoints/access_group_endpoints.py b/litellm/proxy/management_endpoints/access_group_endpoints.py index caaec12f7a39..30fed0962324 100644 --- a/litellm/proxy/management_endpoints/access_group_endpoints.py +++ b/litellm/proxy/management_endpoints/access_group_endpoints.py @@ -9,6 +9,7 @@ LitellmUserRoles, UserAPIKeyAuth, ) +from litellm.proxy.common_utils.cache_pydantic_utils import CacheCodec from litellm.proxy.auth.auth_checks import ( _cache_access_object, _cache_key_object, @@ -236,12 +237,11 @@ async def _patch_key_caches_add_access_group( ) -> None: """Patch cached key objects to include access_group_id.""" for token in key_tokens: - cached_key = await user_api_key_cache.async_get_cache(key=token) - if cached_key is None: + raw = await user_api_key_cache.async_get_cache(key=token) + if raw is None: continue - if isinstance(cached_key, dict): - cached_key = UserAPIKeyAuth(**cached_key) - if not isinstance(cached_key, UserAPIKeyAuth): + cached_key = CacheCodec.deserialize(raw, UserAPIKeyAuth) + if cached_key is None: continue if cached_key.access_group_ids is None: cached_key.access_group_ids = [access_group_id] @@ -267,12 +267,11 @@ async def _patch_key_caches_remove_access_group( ) -> None: """Patch cached key objects to remove access_group_id.""" for token in key_tokens: - cached_key = await user_api_key_cache.async_get_cache(key=token) - if cached_key is None: + raw = await user_api_key_cache.async_get_cache(key=token) + if raw is None: continue - if isinstance(cached_key, dict): - cached_key = UserAPIKeyAuth(**cached_key) - if isinstance(cached_key, UserAPIKeyAuth) and cached_key.access_group_ids: + cached_key = CacheCodec.deserialize(raw, UserAPIKeyAuth) + if cached_key is not None and cached_key.access_group_ids: cached_key.access_group_ids = [ ag for ag in cached_key.access_group_ids if ag != access_group_id ] diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index 46e7963da7c5..0362c3ff6b2a 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -70,6 +70,7 @@ UserAPIKeyAuth, ) from litellm.proxy.auth.auth_checks import ExperimentalUIJWTToken, get_user_object +from litellm.proxy.common_utils.cache_pydantic_utils import CacheCodec from litellm.proxy.auth.auth_utils import _has_user_setup_sso from litellm.proxy.auth.handle_jwt import JWTHandler from litellm.proxy.auth.user_api_key_auth import user_api_key_auth @@ -1233,11 +1234,7 @@ async def _sync_user_role_from_jwt_role_map( user_info.user_role = mapped_role.value await user_api_key_cache.async_set_cache( key=user_info.user_id, - value=( - user_info.model_dump() - if hasattr(user_info, "model_dump") - else dict(user_info) - ), + value=CacheCodec.serialize(user_info, model_type=LiteLLM_UserTable), ) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 72ec41ef97f3..f8c52a5289b9 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -78,8 +78,11 @@ InvitationNew, InvitationUpdate, Litellm_EntityType, + LiteLLM_EndUserTable, LiteLLM_JWTAuth, + LiteLLM_TagTable, LiteLLM_TeamTable, + LiteLLM_TeamTableCachedObj, LiteLLM_UserTable, LitellmUserRoles, PassThroughGenericEndpoint, @@ -2029,34 +2032,38 @@ async def _update_key_cache(token: str, response_cost: float): ### UPDATE USER SPEND ### async def _update_user_cache(): ## UPDATE CACHE FOR USER ID + GLOBAL PROXY + if response_cost is None: + return user_ids = [user_id] try: for _id in user_ids: # Fetch the existing cost for the given user if _id is None: continue - existing_spend_obj = await user_api_key_cache.async_get_cache(key=_id) - if existing_spend_obj is None: + cached_user = await user_api_key_cache.async_get_cache(key=_id) + if cached_user is None: # do nothing if there is no cache value return + existing_spend_obj = CacheCodec.deserialize(cached_user, LiteLLM_UserTable) + if existing_spend_obj is None: + return verbose_proxy_logger.debug( f"_update_user_db: existing spend: {existing_spend_obj}; response_cost: {response_cost}" ) - if isinstance(existing_spend_obj, dict): - existing_spend = existing_spend_obj["spend"] - else: - existing_spend = existing_spend_obj.spend + existing_spend = existing_spend_obj.spend or 0.0 # Calculate the new cost by adding the existing cost and response_cost new_spend = existing_spend + response_cost - # Update the cost column for the given user - if isinstance(existing_spend_obj, dict): - existing_spend_obj["spend"] = new_spend - values_to_update_in_cache.append((_id, existing_spend_obj)) - else: - existing_spend_obj.spend = new_spend - values_to_update_in_cache.append((_id, existing_spend_obj.json())) + existing_spend_obj.spend = new_spend + values_to_update_in_cache.append( + ( + _id, + CacheCodec.serialize( + existing_spend_obj, model_type=LiteLLM_UserTable + ), + ) + ) ## UPDATE GLOBAL PROXY ## global_proxy_spend = await user_api_key_cache.async_get_cache( key="{}:spend".format(litellm_proxy_admin_name) @@ -2088,31 +2095,33 @@ async def _update_end_user_cache(): _id = "end_user_id:{}".format(end_user_id) try: # Fetch the existing cost for the given user - existing_spend_obj = await user_api_key_cache.async_get_cache(key=_id) - if existing_spend_obj is None: + cached_end_user = await user_api_key_cache.async_get_cache(key=_id) + if cached_end_user is None: # if user does not exist in LiteLLM_UserTable, create a new user # do nothing if end-user not in api key cache return + existing_spend_obj = CacheCodec.deserialize( + cached_end_user, LiteLLM_EndUserTable + ) + if existing_spend_obj is None: + return verbose_proxy_logger.debug( f"_update_end_user_db: existing spend: {existing_spend_obj}; response_cost: {response_cost}" ) - if existing_spend_obj is None: - existing_spend = 0 - else: - if isinstance(existing_spend_obj, dict): - existing_spend = existing_spend_obj["spend"] - else: - existing_spend = existing_spend_obj.spend + + existing_spend = existing_spend_obj.spend or 0.0 # Calculate the new cost by adding the existing cost and response_cost new_spend = existing_spend + response_cost - # Update the cost column for the given user - if isinstance(existing_spend_obj, dict): - existing_spend_obj["spend"] = new_spend - values_to_update_in_cache.append((_id, existing_spend_obj)) - else: - existing_spend_obj.spend = new_spend - values_to_update_in_cache.append((_id, existing_spend_obj.json())) + existing_spend_obj.spend = new_spend + values_to_update_in_cache.append( + ( + _id, + CacheCodec.serialize( + existing_spend_obj, model_type=LiteLLM_EndUserTable + ), + ) + ) except Exception as e: verbose_proxy_logger.warning( "Spend tracking - failed to update end user spend in cache. " @@ -2131,36 +2140,32 @@ async def _update_team_cache(): _id = "team_id:{}".format(team_id) try: - # Fetch the existing cost for the given user - existing_spend_obj: Optional[LiteLLM_TeamTable] = ( - await user_api_key_cache.async_get_cache(key=_id) + cached_team = await user_api_key_cache.async_get_cache(key=_id) + if cached_team is None: + # do nothing if team not in api key cache + return + existing_spend_obj: Optional[LiteLLM_TeamTableCachedObj] = ( + CacheCodec.deserialize(cached_team, LiteLLM_TeamTableCachedObj) ) if existing_spend_obj is None: - # do nothing if team not in api key cache return verbose_proxy_logger.debug( f"_update_team_db: existing spend: {existing_spend_obj}; response_cost: {response_cost}" ) - if existing_spend_obj is None: - existing_spend: Optional[float] = 0.0 - else: - if isinstance(existing_spend_obj, dict): - existing_spend = existing_spend_obj["spend"] - else: - existing_spend = existing_spend_obj.spend - if existing_spend is None: - existing_spend = 0.0 + existing_spend: float = existing_spend_obj.spend or 0.0 # Calculate the new cost by adding the existing cost and response_cost new_spend = existing_spend + response_cost - # Update the cost column for the given user - if isinstance(existing_spend_obj, dict): - existing_spend_obj["spend"] = new_spend - values_to_update_in_cache.append((_id, existing_spend_obj)) - else: - existing_spend_obj.spend = new_spend - values_to_update_in_cache.append((_id, existing_spend_obj)) + existing_spend_obj.spend = new_spend + values_to_update_in_cache.append( + ( + _id, + CacheCodec.serialize( + existing_spend_obj, model_type=LiteLLM_TeamTableCachedObj + ), + ) + ) except Exception as e: verbose_proxy_logger.warning( "Spend tracking - failed to update team spend in cache. " @@ -2187,32 +2192,32 @@ async def _update_tag_cache(): cache_key = f"tag:{tag_name}" # Fetch the existing tag object from cache - existing_tag_obj = await user_api_key_cache.async_get_cache( - key=cache_key - ) - if existing_tag_obj is None: + cached_tag = await user_api_key_cache.async_get_cache(key=cache_key) + if cached_tag is None: # do nothing if tag not in api key cache continue + existing_tag_obj = CacheCodec.deserialize(cached_tag, LiteLLM_TagTable) + if existing_tag_obj is None: + continue + verbose_proxy_logger.debug( f"_update_tag_cache: existing spend for tag={tag_name}: {existing_tag_obj}; response_cost: {response_cost}" ) - if isinstance(existing_tag_obj, dict): - existing_spend = existing_tag_obj.get("spend", 0) or 0 - else: - existing_spend = getattr(existing_tag_obj, "spend", 0) or 0 - + existing_spend = existing_tag_obj.spend or 0.0 # Calculate the new cost by adding the existing cost and response_cost new_spend = existing_spend + response_cost - # Update the spend column for the given tag - if isinstance(existing_tag_obj, dict): - existing_tag_obj["spend"] = new_spend - values_to_update_in_cache.append((cache_key, existing_tag_obj)) - else: - existing_tag_obj.spend = new_spend - values_to_update_in_cache.append((cache_key, existing_tag_obj)) + existing_tag_obj.spend = new_spend + values_to_update_in_cache.append( + ( + cache_key, + CacheCodec.serialize( + existing_tag_obj, model_type=LiteLLM_TagTable + ), + ) + ) except Exception as e: verbose_proxy_logger.warning( "Spend tracking - failed to update tag spend in cache. " @@ -2849,7 +2854,7 @@ def _init_cache( self, cache_params: dict, ): - global redis_usage_cache, llm_router + global redis_usage_cache, llm_router, general_settings from litellm import Cache if "default_in_memory_ttl" in cache_params: @@ -2871,16 +2876,23 @@ def _init_cache( ) # Note: PKCE verifier storage uses redis_usage_cache directly (not # user_api_key_cache) to avoid routing all API-key lookups through Redis. - # Share the same Redis client for virtual-key lookups (same DualCache as - # model_max_budget_limiter). attach_redis_cache is a no-op if Redis is - # already set (e.g. config reload). - user_api_key_cache.attach_redis_cache( - redis_usage_cache, - default_redis_ttl=litellm.default_redis_ttl, - ) - verbose_proxy_logger.debug( - "Attached redis_usage_cache Redis client to user_api_key_cache" - ) + if general_settings.get("enable_redis_auth_cache") is True: + user_api_key_cache.attach_redis_cache( + redis_usage_cache, + default_redis_ttl=litellm.default_redis_ttl, + ) + verbose_proxy_logger.info( + "enable_redis_auth_cache=True: attached Redis to " + "user_api_key_cache — virtual-key lookups are now " + "shared across all proxy workers." + ) + else: + verbose_proxy_logger.info( + "enable_redis_auth_cache is not set: user_api_key_cache " + "remains in-memory only (per-worker). Set " + "general_settings.enable_redis_auth_cache: true to share " + "the auth cache across workers and reduce DB load." + ) def switch_on_llm_response_caching(self): """ diff --git a/tests/test_litellm/proxy/common_utils/test_cache_codec.py b/tests/test_litellm/proxy/common_utils/test_cache_codec.py index e99f22a75d4b..ef52eab4b595 100644 --- a/tests/test_litellm/proxy/common_utils/test_cache_codec.py +++ b/tests/test_litellm/proxy/common_utils/test_cache_codec.py @@ -1,5 +1,6 @@ import logging from typing import Optional +from unittest.mock import patch import pytest from pydantic import BaseModel, ValidationError @@ -50,6 +51,41 @@ def test_with_model_type_invalid_dict_raises(self): with pytest.raises(ValidationError): CacheCodec.serialize({"count": 1}, model_type=_SampleModel) + def test_with_model_type_already_correct_instance_skips_revalidation(self): + """Fast-path: value is already model_type — model_validate must NOT be called.""" + m = _SampleModel(name="fast", count=7) + with patch.object(_SampleModel, "model_validate", wraps=_SampleModel.model_validate) as mock_validate: + out = CacheCodec.serialize(m, model_type=_SampleModel) + assert out == {"name": "fast", "count": 7} + mock_validate.assert_not_called() + + def test_with_model_type_subclass_instance_skips_revalidation(self): + """Subclass is isinstance of base → should also take the fast path.""" + sub = _SampleSubModel(name="sub", count=2) + with patch.object(_SampleModel, "model_validate", wraps=_SampleModel.model_validate) as mock_validate: + out = CacheCodec.serialize(sub, model_type=_SampleModel) + assert out == {"name": "sub", "count": 2} + mock_validate.assert_not_called() + + def test_with_model_type_dict_input_goes_through_model_validate(self): + """A dict value (not yet an instance) must still go through model_validate.""" + raw = {"name": "via-dict", "count": 5} + with patch.object( + _SampleModel, "model_validate", wraps=_SampleModel.model_validate + ) as mock_validate: + out = CacheCodec.serialize(raw, model_type=_SampleModel) + assert out == {"name": "via-dict", "count": 5} + mock_validate.assert_called_once() + + def test_with_model_type_incompatible_model_raises_validation_error(self): + """Passing an instance of a completely different model is a caller error and raises.""" + + class _IncompatibleModel(BaseModel): + name: str + + with pytest.raises(Exception): + CacheCodec.serialize(_IncompatibleModel(name="x"), model_type=_SampleModel) + class TestCacheCodecDeserialize: def test_none_returns_none(self): diff --git a/tests/test_litellm/proxy/test_redis_auth_cache_flag.py b/tests/test_litellm/proxy/test_redis_auth_cache_flag.py new file mode 100644 index 000000000000..f824fd2f9c0c --- /dev/null +++ b/tests/test_litellm/proxy/test_redis_auth_cache_flag.py @@ -0,0 +1,130 @@ +""" +Tests for the enable_redis_auth_cache general_settings flag. + +Verifies that _init_cache attaches Redis to user_api_key_cache only when +the flag is explicitly set to True, and leaves it in-memory-only otherwise. +""" + +from contextlib import contextmanager +from unittest.mock import MagicMock, patch + +import pytest + +import litellm +import litellm.proxy.proxy_server as ps +from litellm.caching.caching import RedisCache +from litellm.caching.dual_cache import DualCache + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +class _FakeRedisCache(RedisCache): + """ + Minimal RedisCache subclass that passes isinstance checks without + requiring a real Redis connection. __init__ is bypassed so no + network calls are made. + """ + + def __init__(self): # noqa: super().__init__ skipped intentionally + pass + + +@contextmanager +def _patched_init_cache(general_settings: dict, cache_params: dict): + """ + Context manager that: + 1. Replaces the module-level globals with fresh DualCache instances. + 2. Patches ``litellm.Cache`` (locally imported inside _init_cache) so + it returns a fake cache whose ``.cache`` attribute is a + _FakeRedisCache (passes the isinstance guard in _init_cache). + 3. Yields (user_api_key_cache, spend_counter_cache) after calling + _init_cache, then restores everything. + """ + fake_redis = _FakeRedisCache() + + mock_litellm_cache = MagicMock() + mock_litellm_cache.cache = fake_redis + + fresh_user_cache = DualCache() + fresh_spend_cache = DualCache() + + with ( + patch.object(ps, "general_settings", general_settings), + patch.object(ps, "user_api_key_cache", fresh_user_cache), + patch.object(ps, "spend_counter_cache", fresh_spend_cache), + patch.object(ps, "llm_router", None), + # Cache is locally imported inside _init_cache: patch it at source. + patch("litellm.Cache", return_value=mock_litellm_cache), + ): + litellm.cache = None + ps.ProxyConfig()._init_cache(cache_params) + yield fresh_user_cache, fresh_spend_cache + + +# --------------------------------------------------------------------------- +# Tests +# --------------------------------------------------------------------------- + + +class TestRedisAuthCacheFlag: + def test_flag_true_attaches_redis_to_user_api_key_cache(self): + """When enable_redis_auth_cache=True, user_api_key_cache.redis_cache must be set.""" + with _patched_init_cache( + general_settings={"enable_redis_auth_cache": True}, + cache_params={"type": "redis", "host": "localhost", "port": 6379}, + ) as (user_cache, _): + assert user_cache.redis_cache is not None, ( + "Redis should be attached to user_api_key_cache when " + "enable_redis_auth_cache=True" + ) + + def test_flag_false_leaves_user_api_key_cache_in_memory_only(self): + """When enable_redis_auth_cache=False, user_api_key_cache must stay in-memory.""" + with _patched_init_cache( + general_settings={"enable_redis_auth_cache": False}, + cache_params={"type": "redis", "host": "localhost", "port": 6379}, + ) as (user_cache, _): + assert user_cache.redis_cache is None, ( + "user_api_key_cache must remain in-memory-only when " + "enable_redis_auth_cache=False" + ) + + def test_flag_absent_leaves_user_api_key_cache_in_memory_only(self): + """When enable_redis_auth_cache is not set at all, default is in-memory-only.""" + with _patched_init_cache( + general_settings={}, + cache_params={"type": "redis", "host": "localhost", "port": 6379}, + ) as (user_cache, _): + assert user_cache.redis_cache is None, ( + "user_api_key_cache must remain in-memory-only when " + "enable_redis_auth_cache is absent from general_settings" + ) + + def test_spend_counter_cache_always_gets_redis_regardless_of_flag(self): + """spend_counter_cache must receive Redis regardless of the auth-cache flag.""" + for flag_value in (True, False, None): + gs = ( + {"enable_redis_auth_cache": flag_value} + if flag_value is not None + else {} + ) + with _patched_init_cache( + general_settings=gs, + cache_params={"type": "redis", "host": "localhost", "port": 6379}, + ) as (_, spend_cache): + assert spend_cache.redis_cache is not None, ( + f"spend_counter_cache must always get Redis " + f"(enable_redis_auth_cache={flag_value!r})" + ) + + def test_flag_false_spend_gets_redis_but_user_cache_does_not(self): + """Explicit False: spend cache wired, auth cache left in-memory.""" + with _patched_init_cache( + general_settings={"enable_redis_auth_cache": False}, + cache_params={"type": "redis", "host": "localhost", "port": 6379}, + ) as (user_cache, spend_cache): + assert spend_cache.redis_cache is not None + assert user_cache.redis_cache is None From 655e75276e156cb28d0df8c214523182076a28af Mon Sep 17 00:00:00 2001 From: harish-berri Date: Fri, 24 Apr 2026 01:44:08 +0000 Subject: [PATCH 07/28] Enhance team endpoint tests by integrating AsyncMock for cache methods, ensuring proper asynchronous behavior in test_update_team_guardrails_with_org_id. This improves test reliability and aligns with recent caching improvements. --- .../proxy/management_endpoints/test_team_endpoints.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py index f1e077b5e4a4..65015b9d6b69 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -5151,6 +5151,9 @@ async def test_update_team_guardrails_with_org_id(): return_value=mock_updated_team ) mock_prisma.jsonify_team_object = MagicMock(side_effect=lambda db_data: db_data) + # async_get_cache must be an AsyncMock so `await` in get_org_object works + mock_cache.async_get_cache = AsyncMock(return_value=None) + mock_cache.async_set_cache = AsyncMock() # Mock llm_router mock_router = MagicMock() From d9292e7bcfdf984fcf94a6ffdbe2919e06c964a7 Mon Sep 17 00:00:00 2001 From: harish-berri Date: Fri, 24 Apr 2026 16:13:42 +0000 Subject: [PATCH 08/28] Update test for CacheCodec serialization to clarify validation error handling. --- .../proxy/common_utils/test_cache_codec.py | 13 +++++++++---- 1 file changed, 9 insertions(+), 4 deletions(-) diff --git a/tests/test_litellm/proxy/common_utils/test_cache_codec.py b/tests/test_litellm/proxy/common_utils/test_cache_codec.py index ef52eab4b595..044d4c2d1a7f 100644 --- a/tests/test_litellm/proxy/common_utils/test_cache_codec.py +++ b/tests/test_litellm/proxy/common_utils/test_cache_codec.py @@ -78,13 +78,18 @@ def test_with_model_type_dict_input_goes_through_model_validate(self): mock_validate.assert_called_once() def test_with_model_type_incompatible_model_raises_validation_error(self): - """Passing an instance of a completely different model is a caller error and raises.""" + """Passing a BaseModel whose fields don't satisfy model_type's required fields raises. + + _IncompatibleModel only has `foo: int`, so when Pydantic v2 extracts its + data and validates it against _SampleModel (which requires `name: str`), + a ValidationError is raised. + """ class _IncompatibleModel(BaseModel): - name: str + foo: int # missing required 'name' field of _SampleModel - with pytest.raises(Exception): - CacheCodec.serialize(_IncompatibleModel(name="x"), model_type=_SampleModel) + with pytest.raises(ValidationError): + CacheCodec.serialize(_IncompatibleModel(foo=1), model_type=_SampleModel) class TestCacheCodecDeserialize: From ce087aa796f05bced22e297953f99baabc4fba08 Mon Sep 17 00:00:00 2001 From: harish-berri Date: Fri, 24 Apr 2026 23:44:57 +0000 Subject: [PATCH 09/28] refactor(proxy): update documentation and code references for enable_redis_auth_cache - Removed the deprecated enable_redis_auth_cache field from ConfigGeneralSettings. - Updated references in proxy_server.py to use litellm_settings for enable_redis_auth_cache. - Clarified documentation to reflect the new configuration structure and its impact on Redis integration for user_api_key_cache. --- docs/my-website/docs/proxy/config_settings.md | 8 ++++---- litellm/proxy/_types.py | 13 ------------- litellm/proxy/proxy_server.py | 13 ++++++++++--- 3 files changed, 14 insertions(+), 20 deletions(-) diff --git a/docs/my-website/docs/proxy/config_settings.md b/docs/my-website/docs/proxy/config_settings.md index 33b6006ac2c8..cc9c9de58f55 100644 --- a/docs/my-website/docs/proxy/config_settings.md +++ b/docs/my-website/docs/proxy/config_settings.md @@ -307,7 +307,6 @@ router_settings: | token_rate_limit_type | string | Rate limit counting method: "total", "output", or "input" tokens | | use_redis_transaction_buffer | boolean | If true, buffers database transactions in Redis before writing | | use_shared_health_check | boolean | If true, uses Redis-backed shared health check state across multiple proxy instances | -| enable_redis_auth_cache | boolean | **[Beta]** When `true`, attaches Redis to the virtual-key auth cache (`user_api_key_cache`) so all proxy workers/pods share the same cache instead of each pod resolving keys independently against the database. Requires `litellm_settings.cache: true` with a Redis backend. Significantly reduces database load in multi-worker deployments by eliminating per-pod cache misses on the `combined_view` query. Off by default for a safe phased rollout — enable once your Redis cluster is healthy. Will become the default in a future release. See [Redis Auth Cache](#redis-auth-cache-multi-worker-db-load-reduction). | | user_header_mappings | dict | Map custom request headers to user IDs using lookup rules | | user_header_name | string | HTTP header name to extract user identity from requests | @@ -323,21 +322,22 @@ In multi-worker or multi-pod deployments each worker process keeps its own **in- ### Solution -Setting `enable_redis_auth_cache: true` attaches Redis to `user_api_key_cache` so the resolved key object is stored in a **shared** Redis cache. A cache hit on any worker prevents the DB query entirely. +Setting `enable_redis_auth_cache: true` under `litellm_settings` attaches Redis to `user_api_key_cache` so the resolved key object is stored in a **shared** Redis cache. A cache hit on any worker prevents the DB query entirely. ```yaml # config.yaml litellm_settings: + enable_redis_auth_cache: true # ← share the auth cache across workers cache: true cache_params: type: redis host: os.environ/REDIS_HOST port: os.environ/REDIS_PORT + max_connections: 100 # optional — tune per-worker pool size (default: 50) general_settings: master_key: sk-1234 - enable_redis_auth_cache: true # ← share the auth cache across workers ``` ### Requirements @@ -366,7 +366,7 @@ enable_redis_auth_cache=True: attached Redis to user_api_key_cache — virtual-k If the flag is off you will see: ``` -enable_redis_auth_cache is not set: user_api_key_cache remains in-memory only (per-worker). Set general_settings.enable_redis_auth_cache: true to share the auth cache across workers and reduce DB load. +enable_redis_auth_cache is not set: user_api_key_cache remains in-memory only (per-worker). Set litellm_settings.enable_redis_auth_cache: true to share the auth cache across workers and reduce DB load. ``` --- diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index ee7e7ce94638..b8e0aa0d1281 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -2383,19 +2383,6 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase): None, description="List of MCP server fields that must be filled in for a submission to pass standards checks (e.g. ['description', 'source_url', 'alias']).", ) - enable_redis_auth_cache: Optional[bool] = Field( - None, - description=( - "When True, attaches Redis to user_api_key_cache so virtual-key lookups " - "are shared across all proxy workers/pods instead of being resolved " - "per-process. Requires a Redis cache to be configured under litellm_settings. " - "Reduces DB load significantly in multi-worker deployments by eliminating " - "redundant combined_view SQL queries caused by per-pod cache misses. " - "Off by default for a safe phased rollout — set to True once your Redis " - "cluster is healthy and the CacheCodec serialisation has been validated in " - "your environment. Will be enabled by default in a future release." - ), - ) class ConfigYAML(LiteLLMPydanticObjectBase): diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 73bc46ba391f..8d199d2c092c 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -2979,6 +2979,7 @@ def load_team_config(self, team_id: str): def _init_cache( self, cache_params: dict, + enable_redis_auth_cache: bool = False, ): global redis_usage_cache, llm_router, general_settings from litellm import Cache @@ -3002,7 +3003,7 @@ def _init_cache( ) # Note: PKCE verifier storage uses redis_usage_cache directly (not # user_api_key_cache) to avoid routing all API-key lookups through Redis. - if general_settings.get("enable_redis_auth_cache") is True: + if enable_redis_auth_cache is True: user_api_key_cache.attach_redis_cache( redis_usage_cache, default_redis_ttl=litellm.default_redis_ttl, @@ -3016,7 +3017,7 @@ def _init_cache( verbose_proxy_logger.info( "enable_redis_auth_cache is not set: user_api_key_cache " "remains in-memory only (per-worker). Set " - "general_settings.enable_redis_auth_cache: true to share " + "litellm_settings.enable_redis_auth_cache: true to share " "the auth cache across workers and reduce DB load." ) @@ -3337,7 +3338,13 @@ async def load_config( # noqa: PLR0915 cache_params[key] = get_secret(value) ## to pass a complete url, or set ssl=True, etc. just set it as `os.environ[REDIS_URL] = `, _redis.py checks for REDIS specific environment variables - self._init_cache(cache_params=cache_params) + self._init_cache( + cache_params=cache_params, + enable_redis_auth_cache=litellm_settings.get( + "enable_redis_auth_cache", False + ) + is True, + ) if litellm.cache is not None: verbose_proxy_logger.debug( f"{blue_color_code}Set Cache on LiteLLM Proxy{reset_color_code}" From 77847752fafa253aeac2e34b4b3ceacc3ba16c6b Mon Sep 17 00:00:00 2001 From: harish-berri Date: Sat, 25 Apr 2026 00:21:33 +0000 Subject: [PATCH 10/28] feat(auth): add Prometheus metrics for combined_view SQL queries - Introduced `litellm_auth_combined_view_queries_total` metric to track the number of combined_view SQL queries issued for virtual-key authentication. - Added `AuthMetrics` class to encapsulate metric incrementing logic. - Updated relevant code to increment the new metric on database lookups, aiding in the validation of Redis cache effectiveness. - Updated Prometheus metric definitions to include the new auth diagnostic metrics. --- litellm/integrations/prometheus.py | 14 +++++++ litellm/proxy/auth/auth_checks.py | 2 + litellm/proxy/auth/auth_metrics.py | 52 ++++++++++++++++++++++++ litellm/types/integrations/prometheus.py | 7 ++++ 4 files changed, 75 insertions(+) create mode 100644 litellm/proxy/auth/auth_metrics.py diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index 723b142dfadb..1a899cf0e838 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -265,6 +265,20 @@ def __init__( # noqa: PLR0915 ######################################## # LiteLLM Virtual API KEY metrics ######################################## + + # Auth DB load diagnostic: count direct combined_view SQL queries. + # Each increment means a virtual-key cache miss that hit the DB. + # Useful for validating that enable_redis_auth_cache is working. + self.litellm_auth_combined_view_queries_total = self._counter_factory( + "litellm_auth_combined_view_queries_total", + "Number of times the combined_view SQL query was issued for virtual-key auth. " + "Each count is a cache miss that hit the database. Use to validate " + "enable_redis_auth_cache is reducing DB load.", + labelnames=self.get_labels_for_metric( + "litellm_auth_combined_view_queries_total" + ), + ) + # Remaining MODEL RPM limit for API Key self.litellm_remaining_api_key_requests_for_model = self._gauge_factory( "litellm_remaining_api_key_requests_for_model", diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 7b1d1c234d10..3cce911cd0cb 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -60,6 +60,7 @@ SpecialModelNames, UserAPIKeyAuth, ) +from litellm.proxy.auth.auth_metrics import AuthMetrics from litellm.proxy.auth.route_checks import RouteChecks from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler from litellm.proxy.guardrails.tool_name_extraction import ( @@ -2257,6 +2258,7 @@ async def _fetch_key_object_from_db_with_reconnect( Fetch key object from DB and retry once if a DB connection error can be healed. """ try: + AuthMetrics.inc_combined_view_query(hashed_token) return await prisma_client.get_data( token=hashed_token, table_name="combined_view", diff --git a/litellm/proxy/auth/auth_metrics.py b/litellm/proxy/auth/auth_metrics.py new file mode 100644 index 000000000000..0b463cc920ff --- /dev/null +++ b/litellm/proxy/auth/auth_metrics.py @@ -0,0 +1,52 @@ +""" +Prometheus metric helpers for the auth layer. + +All metrics are thin wrappers around the shared ``PrometheusLogger`` instance so +that every counter follows the same registration path (``_counter_factory``, +label-filter config) as the rest of LiteLLM's metrics. + +Usage:: + + from litellm.proxy.auth.auth_metrics import AuthMetrics + + AuthMetrics.inc_combined_view_query(hashed_token="sk-xxx") +""" + +from litellm._logging import verbose_proxy_logger + + +class AuthMetrics: + """Static helpers for incrementing auth-layer Prometheus counters.""" + + @staticmethod + def _get_prom(): + """Return the active PrometheusLogger, or None if Prometheus is not configured.""" + try: + from litellm.router_utils.cooldown_callbacks import ( + _get_prometheus_logger_from_callbacks, + ) + + return _get_prometheus_logger_from_callbacks() + except Exception: + return None + + @staticmethod + def inc_combined_view_query(hashed_token: str) -> None: + """ + Increment ``litellm_auth_combined_view_queries_total``. + + Called once per virtual-key DB lookup (combined_view query). Each + increment represents a cache miss that hit the database — use this to + validate that ``enable_redis_auth_cache`` is reducing DB load. + """ + try: + prom = AuthMetrics._get_prom() + if prom is not None: + prom.litellm_auth_combined_view_queries_total.labels( + hashed_api_key=hashed_token + ).inc() + except Exception as e: + verbose_proxy_logger.debug( + "AuthMetrics.inc_combined_view_query: failed to increment counter: %s", + e, + ) diff --git a/litellm/types/integrations/prometheus.py b/litellm/types/integrations/prometheus.py index 43a287f29bc1..a4573535f725 100644 --- a/litellm/types/integrations/prometheus.py +++ b/litellm/types/integrations/prometheus.py @@ -228,6 +228,8 @@ class UserAPIKeyLabelNames(Enum): "litellm_guardrail_latency_seconds", "litellm_guardrail_errors_total", "litellm_guardrail_requests_total", + # Auth DB diagnostic metrics + "litellm_auth_combined_view_queries_total", # Cache metrics "litellm_cache_hits_metric", "litellm_cache_misses_metric", @@ -307,6 +309,11 @@ class PrometheusMetricLabels: litellm_guardrail_errors_total: List[str] = [] litellm_guardrail_requests_total: List[str] = [] + # Auth DB diagnostic - label by key so you can see which virtual key causes DB hits + litellm_auth_combined_view_queries_total = [ + UserAPIKeyLabelNames.API_KEY_HASH.value, + ] + litellm_proxy_total_requests_metric = [ UserAPIKeyLabelNames.END_USER.value, UserAPIKeyLabelNames.API_KEY_HASH.value, From f08a8534ab199e3adfb1cb190a3a6d994e987a0a Mon Sep 17 00:00:00 2001 From: harish-berri Date: Sat, 25 Apr 2026 00:44:59 +0000 Subject: [PATCH 11/28] fix failing test --- .../proxy/test_redis_auth_cache_flag.py | 27 ++++++++++--------- 1 file changed, 15 insertions(+), 12 deletions(-) diff --git a/tests/test_litellm/proxy/test_redis_auth_cache_flag.py b/tests/test_litellm/proxy/test_redis_auth_cache_flag.py index f824fd2f9c0c..261222be9a24 100644 --- a/tests/test_litellm/proxy/test_redis_auth_cache_flag.py +++ b/tests/test_litellm/proxy/test_redis_auth_cache_flag.py @@ -1,5 +1,5 @@ """ -Tests for the enable_redis_auth_cache general_settings flag. +Tests for the enable_redis_auth_cache litellm_settings flag. Verifies that _init_cache attaches Redis to user_api_key_cache only when the flag is explicitly set to True, and leaves it in-memory-only otherwise. @@ -33,14 +33,16 @@ def __init__(self): # noqa: super().__init__ skipped intentionally @contextmanager -def _patched_init_cache(general_settings: dict, cache_params: dict): +def _patched_init_cache(litellm_settings: dict, cache_params: dict): """ Context manager that: 1. Replaces the module-level globals with fresh DualCache instances. 2. Patches ``litellm.Cache`` (locally imported inside _init_cache) so it returns a fake cache whose ``.cache`` attribute is a _FakeRedisCache (passes the isinstance guard in _init_cache). - 3. Yields (user_api_key_cache, spend_counter_cache) after calling + 3. Extracts enable_redis_auth_cache from litellm_settings and passes it + as the second argument to _init_cache (matching production behaviour). + 4. Yields (user_api_key_cache, spend_counter_cache) after calling _init_cache, then restores everything. """ fake_redis = _FakeRedisCache() @@ -51,8 +53,9 @@ def _patched_init_cache(general_settings: dict, cache_params: dict): fresh_user_cache = DualCache() fresh_spend_cache = DualCache() + enable_redis_auth_cache = litellm_settings.get("enable_redis_auth_cache", False) + with ( - patch.object(ps, "general_settings", general_settings), patch.object(ps, "user_api_key_cache", fresh_user_cache), patch.object(ps, "spend_counter_cache", fresh_spend_cache), patch.object(ps, "llm_router", None), @@ -60,7 +63,7 @@ def _patched_init_cache(general_settings: dict, cache_params: dict): patch("litellm.Cache", return_value=mock_litellm_cache), ): litellm.cache = None - ps.ProxyConfig()._init_cache(cache_params) + ps.ProxyConfig()._init_cache(cache_params, enable_redis_auth_cache) yield fresh_user_cache, fresh_spend_cache @@ -73,7 +76,7 @@ class TestRedisAuthCacheFlag: def test_flag_true_attaches_redis_to_user_api_key_cache(self): """When enable_redis_auth_cache=True, user_api_key_cache.redis_cache must be set.""" with _patched_init_cache( - general_settings={"enable_redis_auth_cache": True}, + litellm_settings={"enable_redis_auth_cache": True}, cache_params={"type": "redis", "host": "localhost", "port": 6379}, ) as (user_cache, _): assert user_cache.redis_cache is not None, ( @@ -84,7 +87,7 @@ def test_flag_true_attaches_redis_to_user_api_key_cache(self): def test_flag_false_leaves_user_api_key_cache_in_memory_only(self): """When enable_redis_auth_cache=False, user_api_key_cache must stay in-memory.""" with _patched_init_cache( - general_settings={"enable_redis_auth_cache": False}, + litellm_settings={"enable_redis_auth_cache": False}, cache_params={"type": "redis", "host": "localhost", "port": 6379}, ) as (user_cache, _): assert user_cache.redis_cache is None, ( @@ -95,24 +98,24 @@ def test_flag_false_leaves_user_api_key_cache_in_memory_only(self): def test_flag_absent_leaves_user_api_key_cache_in_memory_only(self): """When enable_redis_auth_cache is not set at all, default is in-memory-only.""" with _patched_init_cache( - general_settings={}, + litellm_settings={}, cache_params={"type": "redis", "host": "localhost", "port": 6379}, ) as (user_cache, _): assert user_cache.redis_cache is None, ( "user_api_key_cache must remain in-memory-only when " - "enable_redis_auth_cache is absent from general_settings" + "enable_redis_auth_cache is absent from litellm_settings" ) def test_spend_counter_cache_always_gets_redis_regardless_of_flag(self): """spend_counter_cache must receive Redis regardless of the auth-cache flag.""" for flag_value in (True, False, None): - gs = ( + ls = ( {"enable_redis_auth_cache": flag_value} if flag_value is not None else {} ) with _patched_init_cache( - general_settings=gs, + litellm_settings=ls, cache_params={"type": "redis", "host": "localhost", "port": 6379}, ) as (_, spend_cache): assert spend_cache.redis_cache is not None, ( @@ -123,7 +126,7 @@ def test_spend_counter_cache_always_gets_redis_regardless_of_flag(self): def test_flag_false_spend_gets_redis_but_user_cache_does_not(self): """Explicit False: spend cache wired, auth cache left in-memory.""" with _patched_init_cache( - general_settings={"enable_redis_auth_cache": False}, + litellm_settings={"enable_redis_auth_cache": False}, cache_params={"type": "redis", "host": "localhost", "port": 6379}, ) as (user_cache, spend_cache): assert spend_cache.redis_cache is not None From 3c2c61e1e4c1d4f54ef2d796710505d56436faa8 Mon Sep 17 00:00:00 2001 From: harish-berri Date: Tue, 28 Apr 2026 19:15:03 +0000 Subject: [PATCH 12/28] refactor(proxy): replace DualCache with UserApiKeyCache for user API key management - Updated instances of DualCache to UserApiKeyCache across multiple files to enhance cache handling for user API keys. - Adjusted cache retrieval and storage methods to ensure proper serialization and deserialization of cached objects. - Introduced a new UserApiKeyCache class to streamline caching logic and improve type safety. - Updated relevant tests to reflect changes in caching behavior and ensure compatibility with the new cache implementation. --- litellm/proxy/auth/auth_checks.py | 310 +++++++++--------- litellm/proxy/auth/handle_jwt.py | 29 +- litellm/proxy/auth/user_api_key_auth.py | 30 +- .../proxy/common_utils/user_api_key_cache.py | 156 +++++++++ .../access_group_endpoints.py | 17 +- litellm/proxy/management_endpoints/ui_sso.py | 12 +- litellm/proxy/proxy_server.py | 27 +- .../proxy/auth/test_auth_checks.py | 36 +- .../common_utils/test_user_api_key_cache.py | 113 +++++++ .../proxy/test_redis_auth_cache_flag.py | 14 +- 10 files changed, 502 insertions(+), 242 deletions(-) create mode 100644 litellm/proxy/common_utils/user_api_key_cache.py create mode 100644 tests/test_litellm/proxy/common_utils/test_user_api_key_cache.py diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 3cce911cd0cb..8bd8eb794e02 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -68,6 +68,7 @@ extract_request_tool_names, ) from litellm.proxy.common_utils.cache_pydantic_utils import CacheCodec +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.proxy.route_llm_request import route_request from litellm.proxy.utils import PrismaClient, ProxyLogging, log_db_metrics from litellm.router import Router @@ -854,7 +855,7 @@ def get_actual_routes(allowed_routes: list) -> list: async def get_default_end_user_budget( prisma_client: Optional[PrismaClient], - user_api_key_cache: DualCache, + user_api_key_cache: UserApiKeyCache, parent_otel_span: Optional[Span] = None, ) -> Optional[LiteLLM_BudgetTable]: """ @@ -877,11 +878,12 @@ async def get_default_end_user_budget( cache_key = f"default_end_user_budget:{litellm.max_end_user_budget_id}" # Check cache first - cached_budget = await user_api_key_cache.async_get_cache(key=cache_key) + cached_budget = await user_api_key_cache.async_get_cache( + key=cache_key, + model_type=LiteLLM_BudgetTable, + ) if cached_budget is not None: - deserialized = CacheCodec.deserialize(cached_budget, LiteLLM_BudgetTable) - if deserialized is not None: - return deserialized + return cached_budget # Fetch from database try: @@ -899,7 +901,8 @@ async def get_default_end_user_budget( # Cache the budget for 60 seconds await user_api_key_cache.async_set_cache( key=cache_key, - value=CacheCodec.serialize(_budget_obj, model_type=LiteLLM_BudgetTable), + value=_budget_obj, + model_type=LiteLLM_BudgetTable, ttl=DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL, ) @@ -914,7 +917,7 @@ async def get_default_end_user_budget( async def get_team_member_default_budget( budget_id: str, prisma_client: Optional[PrismaClient], - user_api_key_cache: DualCache, + user_api_key_cache: UserApiKeyCache, ) -> Optional[LiteLLM_BudgetTable]: """ Fetches the team-level default per-member budget referenced by team.metadata["team_member_budget_id"]. @@ -970,7 +973,7 @@ async def get_team_member_default_budget( async def _apply_default_budget_to_end_user( end_user_obj: LiteLLM_EndUserTable, prisma_client: PrismaClient, - user_api_key_cache: DualCache, + user_api_key_cache: UserApiKeyCache, parent_otel_span: Optional[Span] = None, ) -> LiteLLM_EndUserTable: """ @@ -1043,7 +1046,7 @@ def _check_end_user_budget( async def get_end_user_object( end_user_id: Optional[str], prisma_client: Optional[PrismaClient], - user_api_key_cache: DualCache, + user_api_key_cache: UserApiKeyCache, route: str, parent_otel_span: Optional[Span] = None, proxy_logging_obj: Optional[ProxyLogging] = None, @@ -1074,22 +1077,24 @@ async def get_end_user_object( _key = "end_user_id:{}".format(end_user_id) # Check cache first - cached_user_obj = await user_api_key_cache.async_get_cache(key=_key) + cached_user_obj = await user_api_key_cache.async_get_cache( + key=_key, + model_type=LiteLLM_EndUserTable, + ) if cached_user_obj is not None: - return_obj = CacheCodec.deserialize(cached_user_obj, LiteLLM_EndUserTable) - if return_obj is not None: - # Apply default budget if needed - return_obj = await _apply_default_budget_to_end_user( - end_user_obj=return_obj, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - parent_otel_span=parent_otel_span, - ) + return_obj = cached_user_obj + # Apply default budget if needed + return_obj = await _apply_default_budget_to_end_user( + end_user_obj=return_obj, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=parent_otel_span, + ) - # Check budget limits - _check_end_user_budget(end_user_obj=return_obj, route=route) + # Check budget limits + _check_end_user_budget(end_user_obj=return_obj, route=route) - return return_obj + return return_obj # Fetch from database try: @@ -1115,7 +1120,8 @@ async def get_end_user_object( # Save to cache await user_api_key_cache.async_set_cache( key="end_user_id:{}".format(end_user_id), - value=CacheCodec.serialize(_response, model_type=LiteLLM_EndUserTable), + value=_response, + model_type=LiteLLM_EndUserTable, ) # Check budget limits @@ -1133,7 +1139,7 @@ async def get_end_user_object( async def get_tag_objects_batch( tag_names: List[str], prisma_client: Optional[PrismaClient], - user_api_key_cache: DualCache, + user_api_key_cache: UserApiKeyCache, parent_otel_span: Optional[Span] = None, proxy_logging_obj: Optional[ProxyLogging] = None, ) -> Dict[str, LiteLLM_TagTable]: @@ -1166,13 +1172,12 @@ async def get_tag_objects_batch( # Try to get all tags from cache first for tag_name in tag_names: cache_key = f"tag:{tag_name}" - cached_tag = await user_api_key_cache.async_get_cache(key=cache_key) + cached_tag = await user_api_key_cache.async_get_cache( + key=cache_key, + model_type=LiteLLM_TagTable, + ) if cached_tag is not None: - deserialized_tag = CacheCodec.deserialize(cached_tag, LiteLLM_TagTable) - if deserialized_tag is not None: - tag_objects[tag_name] = deserialized_tag - else: - uncached_tags.append(tag_name) + tag_objects[tag_name] = cached_tag else: uncached_tags.append(tag_name) @@ -1191,7 +1196,8 @@ async def get_tag_objects_batch( _tag_obj = LiteLLM_TagTable(**db_tag.dict()) await user_api_key_cache.async_set_cache( key=cache_key, - value=CacheCodec.serialize(_tag_obj, model_type=LiteLLM_TagTable), + value=_tag_obj, + model_type=LiteLLM_TagTable, ) tag_objects[tag_name] = _tag_obj except Exception as e: @@ -1204,7 +1210,7 @@ async def get_tag_objects_batch( async def get_tag_object( tag_name: Optional[str], prisma_client: Optional[PrismaClient], - user_api_key_cache: DualCache, + user_api_key_cache: UserApiKeyCache, parent_otel_span: Optional[Span] = None, proxy_logging_obj: Optional[ProxyLogging] = None, ) -> Optional[LiteLLM_TagTable]: @@ -1243,7 +1249,7 @@ async def get_team_membership( user_id: str, team_id: str, prisma_client: Optional[PrismaClient], - user_api_key_cache: DualCache, + user_api_key_cache: UserApiKeyCache, parent_otel_span: Optional[Span] = None, proxy_logging_obj: Optional[ProxyLogging] = None, ) -> Optional["LiteLLM_TeamMembership"]: @@ -1263,13 +1269,12 @@ async def get_team_membership( _key = "team_membership:{}:{}".format(user_id, team_id) # check if in cache - cached_membership_obj = await user_api_key_cache.async_get_cache(key=_key) + cached_membership_obj = await user_api_key_cache.async_get_cache( + key=_key, + model_type=LiteLLM_TeamMembership, + ) if cached_membership_obj is not None: - deserialized_membership = CacheCodec.deserialize( - cached_membership_obj, LiteLLM_TeamMembership - ) - if deserialized_membership is not None: - return deserialized_membership + return cached_membership_obj # else, check db try: @@ -1284,7 +1289,8 @@ async def get_team_membership( _response = LiteLLM_TeamMembership(**response.dict()) await user_api_key_cache.async_set_cache( key=_key, - value=CacheCodec.serialize(_response, model_type=LiteLLM_TeamMembership), + value=_response, + model_type=LiteLLM_TeamMembership, ) return _response @@ -1453,7 +1459,7 @@ async def _get_fuzzy_user_object( async def get_user_object( user_id: Optional[str], prisma_client: Optional[PrismaClient], - user_api_key_cache: DualCache, + user_api_key_cache: UserApiKeyCache, user_id_upsert: bool, parent_otel_span: Optional[Span] = None, proxy_logging_obj: Optional[ProxyLogging] = None, @@ -1472,13 +1478,12 @@ async def get_user_object( # check if in cache if not check_db_only: - cached_user_obj = await user_api_key_cache.async_get_cache(key=user_id) + cached_user_obj = await user_api_key_cache.async_get_cache( + key=user_id, + model_type=LiteLLM_UserTable, + ) if cached_user_obj is not None: - deserialized_user = CacheCodec.deserialize( - cached_user_obj, LiteLLM_UserTable - ) - if deserialized_user is not None: - return deserialized_user + return cached_user_obj # else, check db if prisma_client is None: raise Exception("No db connected") @@ -1540,7 +1545,8 @@ async def get_user_object( # save the user object to cache await user_api_key_cache.async_set_cache( key=user_id, - value=CacheCodec.serialize(_response, model_type=LiteLLM_UserTable), + value=_response, + model_type=LiteLLM_UserTable, ttl=DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL, ) @@ -1562,22 +1568,20 @@ async def get_user_object( async def _cache_management_object( key: str, value: Union[BaseModel, Dict[str, Any]], - user_api_key_cache: DualCache, + user_api_key_cache: UserApiKeyCache, proxy_logging_obj: Optional[ProxyLogging], *, model_type: Type[BaseModel], ): """ - Persist management objects to DualCache (in-memory + optional Redis). + Persist management objects via ``UserApiKeyCache`` (in-memory + optional Redis). - Values must be JSON-serializable for the Redis path (``json.dumps``). Payloads are - encoded with ``CacheCodec.serialize(..., model_type=...)`` so writes match reads - via ``CacheCodec.deserialize(..., model_type)``. + ``UserApiKeyCache`` serializes with ``model_type`` so Redis and in-memory stay aligned. """ - cache_payload = CacheCodec.serialize(value, model_type=model_type) await user_api_key_cache.async_set_cache( key=key, - value=cache_payload, + value=value, + model_type=model_type, ttl=DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL, ) @@ -1585,7 +1589,7 @@ async def _cache_management_object( async def _cache_team_object( team_id: str, team_table: LiteLLM_TeamTableCachedObj, - user_api_key_cache: DualCache, + user_api_key_cache: UserApiKeyCache, proxy_logging_obj: Optional[ProxyLogging], ): key = "team_id:{}".format(team_id) @@ -1605,7 +1609,7 @@ async def _cache_team_object( async def _cache_key_object( hashed_token: str, user_api_key_obj: UserAPIKeyAuth, - user_api_key_cache: DualCache, + user_api_key_cache: UserApiKeyCache, proxy_logging_obj: Optional[ProxyLogging], ): key = hashed_token @@ -1624,7 +1628,7 @@ async def _cache_key_object( async def _delete_cache_key_object( hashed_token: str, - user_api_key_cache: DualCache, + user_api_key_cache: UserApiKeyCache, proxy_logging_obj: Optional[ProxyLogging], ): key = hashed_token @@ -1672,7 +1676,7 @@ async def _get_team_object_from_db(team_id: str, prisma_client: PrismaClient): async def _get_team_object_from_user_api_key_cache( team_id: str, prisma_client: PrismaClient, - user_api_key_cache: DualCache, + user_api_key_cache: UserApiKeyCache, last_db_access_time: LimitedSizeOrderedDict, db_cache_expiry: int, proxy_logging_obj: Optional[ProxyLogging], @@ -1733,35 +1737,38 @@ async def _get_team_object_from_user_api_key_cache( async def _get_team_object_from_cache( key: str, proxy_logging_obj: Optional[ProxyLogging], - user_api_key_cache: DualCache, + user_api_key_cache: UserApiKeyCache, parent_otel_span: Optional[Span], ) -> Optional[LiteLLM_TeamTableCachedObj]: - cached_team_obj: Optional[LiteLLM_TeamTableCachedObj] = None - - ## CHECK REDIS CACHE ## + ## INTERNAL USAGE CACHE (plain DualCache) — checked before UserApiKeyCache stores ## if ( proxy_logging_obj is not None and proxy_logging_obj.internal_usage_cache.dual_cache ): - cached_team_obj = ( + cached_raw = ( await proxy_logging_obj.internal_usage_cache.dual_cache.async_get_cache( key=key, parent_otel_span=parent_otel_span ) ) + if cached_raw is not None: + from_internal = CacheCodec.deserialize( + cached_raw, LiteLLM_TeamTableCachedObj + ) + if from_internal is not None: + return from_internal - if cached_team_obj is None: - cached_team_obj = await user_api_key_cache.async_get_cache(key=key) - - if cached_team_obj is None: - return None - - return CacheCodec.deserialize(cached_team_obj, LiteLLM_TeamTableCachedObj) + decoded = await user_api_key_cache.async_get_cache( + key=key, + parent_otel_span=parent_otel_span, + model_type=LiteLLM_TeamTableCachedObj, + ) + return decoded async def get_team_object( team_id: str, prisma_client: Optional[PrismaClient], - user_api_key_cache: DualCache, + user_api_key_cache: UserApiKeyCache, parent_otel_span: Optional[Span] = None, proxy_logging_obj: Optional[ProxyLogging] = None, check_cache_only: Optional[bool] = None, @@ -1827,23 +1834,21 @@ async def get_team_object( async def _cache_access_object( access_group_id: str, access_group_table: LiteLLM_AccessGroupTable, - user_api_key_cache: DualCache, + user_api_key_cache: UserApiKeyCache, proxy_logging_obj: Optional[ProxyLogging] = None, ): key = "access_group_id:{}".format(access_group_id) - cache_payload = CacheCodec.serialize( - access_group_table, model_type=LiteLLM_AccessGroupTable - ) await user_api_key_cache.async_set_cache( key=key, - value=cache_payload, + value=access_group_table, + model_type=LiteLLM_AccessGroupTable, ttl=DEFAULT_ACCESS_GROUP_CACHE_TTL, ) async def _delete_cache_access_object( access_group_id: str, - user_api_key_cache: DualCache, + user_api_key_cache: UserApiKeyCache, proxy_logging_obj: Optional[ProxyLogging] = None, ): key = "access_group_id:{}".format(access_group_id) @@ -1861,7 +1866,7 @@ async def _delete_cache_access_object( async def get_access_object( access_group_id: str, prisma_client: Optional[PrismaClient], - user_api_key_cache: DualCache, + user_api_key_cache: UserApiKeyCache, proxy_logging_obj: Optional[ProxyLogging] = None, ) -> LiteLLM_AccessGroupTable: """ @@ -1883,10 +1888,12 @@ async def get_access_object( key = "access_group_id:{}".format(access_group_id) - cached_access_obj = await user_api_key_cache.async_get_cache(key=key) - deserialized = CacheCodec.deserialize(cached_access_obj, LiteLLM_AccessGroupTable) - if deserialized is not None: - return deserialized + cached_access_obj = await user_api_key_cache.async_get_cache( + key=key, + model_type=LiteLLM_AccessGroupTable, + ) + if cached_access_obj is not None: + return cached_access_obj # Not in cache - fetch from DB try: @@ -1932,7 +1939,7 @@ async def get_access_object( async def get_team_object_by_alias( team_alias: str, prisma_client: Optional[PrismaClient], - user_api_key_cache: DualCache, + user_api_key_cache: UserApiKeyCache, parent_otel_span: Optional["Span"] = None, proxy_logging_obj: Optional[ProxyLogging] = None, ) -> LiteLLM_TeamTableCachedObj: @@ -2011,19 +2018,18 @@ async def get_team_object_by_alias( ) # Cache the result by both alias and team_id - _serialized_team = CacheCodec.serialize( - team_obj, model_type=LiteLLM_TeamTableCachedObj - ) await user_api_key_cache.async_set_cache( key=cache_key, - value=_serialized_team, + value=team_obj, + model_type=LiteLLM_TeamTableCachedObj, ttl=DEFAULT_IN_MEMORY_TTL, ) # Also cache by team_id for consistency team_id_cache_key = "team_id:{}".format(team_obj.team_id) await user_api_key_cache.async_set_cache( key=team_id_cache_key, - value=_serialized_team, + value=team_obj, + model_type=LiteLLM_TeamTableCachedObj, ttl=DEFAULT_IN_MEMORY_TTL, ) @@ -2045,7 +2051,7 @@ async def get_team_object_by_alias( async def get_org_object_by_alias( org_alias: str, prisma_client: Optional[PrismaClient], - user_api_key_cache: DualCache, + user_api_key_cache: UserApiKeyCache, parent_otel_span: Optional["Span"] = None, proxy_logging_obj: Optional[ProxyLogging] = None, ) -> Optional[LiteLLM_OrganizationTable]: @@ -2072,13 +2078,12 @@ async def get_org_object_by_alias( # Check cache first (keyed by alias) cache_key = "org_alias:{}".format(org_alias) - cached_org_obj = await user_api_key_cache.async_get_cache(key=cache_key) + cached_org_obj = await user_api_key_cache.async_get_cache( + key=cache_key, + model_type=LiteLLM_OrganizationTable, + ) if cached_org_obj is not None: - deserialized_org = CacheCodec.deserialize( - cached_org_obj, LiteLLM_OrganizationTable - ) - if deserialized_org is not None: - return deserialized_org + return cached_org_obj # Query database by organization_alias try: @@ -2105,19 +2110,18 @@ async def get_org_object_by_alias( org = orgs[0] org_obj = LiteLLM_OrganizationTable(**org.model_dump()) - _serialized_org = CacheCodec.serialize( - org_obj, model_type=LiteLLM_OrganizationTable - ) # Cache the result await user_api_key_cache.async_set_cache( key=cache_key, - value=_serialized_org, + value=org_obj, + model_type=LiteLLM_OrganizationTable, ttl=DEFAULT_IN_MEMORY_TTL, ) # Also cache by org_id for consistency await user_api_key_cache.async_set_cache( key="org_id:{}".format(org_obj.organization_id), - value=_serialized_org, + value=org_obj, + model_type=LiteLLM_OrganizationTable, ttl=DEFAULT_IN_MEMORY_TTL, ) @@ -2321,7 +2325,7 @@ async def get_jwt_key_mapping_object( async def get_key_object( hashed_token: str, prisma_client: Optional[PrismaClient], - user_api_key_cache: DualCache, + user_api_key_cache: UserApiKeyCache, parent_otel_span: Optional[Span] = None, proxy_logging_obj: Optional[ProxyLogging] = None, check_cache_only: Optional[bool] = None, @@ -2341,11 +2345,12 @@ async def get_key_object( # Same flow as before: use cache only when we have a hit we can turn into UserAPIKeyAuth # (dict from Redis / model_dump, or UserAPIKeyAuth from in-memory). Otherwise fall through to DB. - cached_key_obj = await user_api_key_cache.async_get_cache(key=key) - if cached_key_obj is not None: - user_api_key_auth = CacheCodec.deserialize(cached_key_obj, UserAPIKeyAuth) - if user_api_key_auth is not None: - return user_api_key_auth + user_api_key_auth = await user_api_key_cache.async_get_cache( + key=key, + model_type=UserAPIKeyAuth, + ) + if user_api_key_auth is not None: + return user_api_key_auth if check_cache_only: raise Exception( @@ -2402,7 +2407,7 @@ async def get_key_object( async def get_object_permission( object_permission_id: str, prisma_client: Optional[PrismaClient], - user_api_key_cache: DualCache, + user_api_key_cache: UserApiKeyCache, parent_otel_span: Optional[Span] = None, proxy_logging_obj: Optional[ProxyLogging] = None, ) -> Optional[LiteLLM_ObjectPermissionTable]: @@ -2418,13 +2423,12 @@ async def get_object_permission( # check if in cache key = "object_permission_id:{}".format(object_permission_id) - cached_obj_permission = await user_api_key_cache.async_get_cache(key=key) - if cached_obj_permission is not None: - deserialized_perm = CacheCodec.deserialize( - cached_obj_permission, LiteLLM_ObjectPermissionTable - ) - if deserialized_perm is not None: - return deserialized_perm + deserialized_perm = await user_api_key_cache.async_get_cache( + key=key, + model_type=LiteLLM_ObjectPermissionTable, + ) + if deserialized_perm is not None: + return deserialized_perm # else, check db try: @@ -2438,9 +2442,8 @@ async def get_object_permission( _perm_obj = LiteLLM_ObjectPermissionTable(**response.dict()) await user_api_key_cache.async_set_cache( key=key, - value=CacheCodec.serialize( - _perm_obj, model_type=LiteLLM_ObjectPermissionTable - ), + value=_perm_obj, + model_type=LiteLLM_ObjectPermissionTable, ttl=DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL, ) @@ -2453,7 +2456,7 @@ async def get_object_permission( async def get_managed_vector_store_rows_by_uuids( uuids: List[str], prisma_client: Optional[PrismaClient], - user_api_key_cache: DualCache, + user_api_key_cache: UserApiKeyCache, parent_otel_span: Optional[Span] = None, proxy_logging_obj: Optional[ProxyLogging] = None, ) -> List[LiteLLM_ManagedVectorStoresTable]: @@ -2473,15 +2476,12 @@ async def get_managed_vector_store_rows_by_uuids( for uuid in uuids: key = "managed_vector_store_id:{}".format(uuid) - cached = await user_api_key_cache.async_get_cache(key=key) - if cached is not None: - deserialized_vs = CacheCodec.deserialize( - cached, LiteLLM_ManagedVectorStoresTable - ) - if deserialized_vs is not None: - result.append(deserialized_vs) - else: - cache_misses.append(uuid) + deserialized_vs = await user_api_key_cache.async_get_cache( + key=key, + model_type=LiteLLM_ManagedVectorStoresTable, + ) + if deserialized_vs is not None: + result.append(deserialized_vs) else: cache_misses.append(uuid) @@ -2507,9 +2507,8 @@ async def get_managed_vector_store_rows_by_uuids( key = "managed_vector_store_id:{}".format(cached_obj.vector_store_id) await user_api_key_cache.async_set_cache( key=key, - value=CacheCodec.serialize( - cached_obj, model_type=LiteLLM_ManagedVectorStoresTable - ), + value=cached_obj, + model_type=LiteLLM_ManagedVectorStoresTable, ttl=DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL, ) result.append(cached_obj) @@ -2521,7 +2520,7 @@ async def get_managed_vector_store_rows_by_uuids( async def get_org_object( org_id: str, prisma_client: Optional[PrismaClient], - user_api_key_cache: DualCache, + user_api_key_cache: UserApiKeyCache, parent_otel_span: Optional[Span] = None, proxy_logging_obj: Optional[ProxyLogging] = None, include_budget_table: bool = False, @@ -2552,13 +2551,12 @@ async def get_org_object( cache_key = "org_id:{}:with_budget".format(org_id) # check if in cache - cached_org_obj = await user_api_key_cache.async_get_cache(key=cache_key) - if cached_org_obj is not None: - deserialized_org = CacheCodec.deserialize( - cached_org_obj, LiteLLM_OrganizationTable - ) - if deserialized_org is not None: - return deserialized_org + deserialized_org = await user_api_key_cache.async_get_cache( + key=cache_key, + model_type=LiteLLM_OrganizationTable, + ) + if deserialized_org is not None: + return deserialized_org # else, check db try: query_kwargs: Dict[str, Any] = {"where": {"organization_id": org_id}} @@ -2576,7 +2574,8 @@ async def get_org_object( # Cache the result await user_api_key_cache.async_set_cache( key=cache_key, - value=CacheCodec.serialize(_org_obj, model_type=LiteLLM_OrganizationTable), + value=_org_obj, + model_type=LiteLLM_OrganizationTable, ttl=DEFAULT_IN_MEMORY_TTL, ) @@ -2593,7 +2592,7 @@ async def _get_resources_from_access_groups( "access_model_names", "access_mcp_server_ids", "access_agent_ids" ], prisma_client: Optional[PrismaClient] = None, - user_api_key_cache: Optional[DualCache] = None, + user_api_key_cache: Optional[UserApiKeyCache] = None, proxy_logging_obj: Optional[ProxyLogging] = None, ) -> List[str]: """ @@ -2651,7 +2650,7 @@ async def _get_resources_from_access_groups( async def _get_models_from_access_groups( access_group_ids: List[str], prisma_client: Optional[PrismaClient] = None, - user_api_key_cache: Optional[DualCache] = None, + user_api_key_cache: Optional[UserApiKeyCache] = None, proxy_logging_obj: Optional[ProxyLogging] = None, ) -> List[str]: """ @@ -2670,7 +2669,7 @@ async def _get_models_from_access_groups( async def _get_mcp_server_ids_from_access_groups( access_group_ids: List[str], prisma_client: Optional[PrismaClient] = None, - user_api_key_cache: Optional[DualCache] = None, + user_api_key_cache: Optional[UserApiKeyCache] = None, proxy_logging_obj: Optional[ProxyLogging] = None, ) -> List[str]: """ @@ -2689,7 +2688,7 @@ async def _get_mcp_server_ids_from_access_groups( async def _get_agent_ids_from_access_groups( access_group_ids: List[str], prisma_client: Optional[PrismaClient] = None, - user_api_key_cache: Optional[DualCache] = None, + user_api_key_cache: Optional[UserApiKeyCache] = None, proxy_logging_obj: Optional[ProxyLogging] = None, ) -> List[str]: """ @@ -3303,7 +3302,7 @@ async def _check_team_member_budget( user_object: Optional[LiteLLM_UserTable], valid_token: Optional[UserAPIKeyAuth], prisma_client: Optional[PrismaClient], - user_api_key_cache: DualCache, + user_api_key_cache: UserApiKeyCache, proxy_logging_obj: ProxyLogging, ): """Check if team member is over their max budget within the team.""" @@ -3370,7 +3369,7 @@ async def _check_team_member_model_access( valid_token: UserAPIKeyAuth, llm_router: Optional[Router], prisma_client: Optional["PrismaClient"], - user_api_key_cache: DualCache, + user_api_key_cache: UserApiKeyCache, proxy_logging_obj: ProxyLogging, ) -> None: """ @@ -3677,7 +3676,7 @@ async def _project_soft_budget_check( async def get_project_object( project_id: str, prisma_client: Optional[PrismaClient], - user_api_key_cache: DualCache, + user_api_key_cache: UserApiKeyCache, proxy_logging_obj: Optional[ProxyLogging] = None, ) -> Optional[LiteLLM_ProjectTableCachedObj]: """ @@ -3692,13 +3691,12 @@ async def get_project_object( # Check cache first cache_key = "project_id:{}".format(project_id) - cached_obj = await user_api_key_cache.async_get_cache(key=cache_key) - if cached_obj is not None: - deserialized_project = CacheCodec.deserialize( - cached_obj, LiteLLM_ProjectTableCachedObj - ) - if deserialized_project is not None: - return deserialized_project + deserialized_project = await user_api_key_cache.async_get_cache( + key=cache_key, + model_type=LiteLLM_ProjectTableCachedObj, + ) + if deserialized_project is not None: + return deserialized_project # Fetch from DB project_row = await prisma_client.db.litellm_projecttable.find_unique( @@ -3727,7 +3725,7 @@ async def _organization_max_budget_check( valid_token: Optional[UserAPIKeyAuth], team_object: Optional[LiteLLM_TeamTable], prisma_client: Optional[PrismaClient], - user_api_key_cache: DualCache, + user_api_key_cache: UserApiKeyCache, proxy_logging_obj: ProxyLogging, ): """ @@ -3821,7 +3819,7 @@ async def _organization_max_budget_check( async def _tag_max_budget_check( request_body: dict, prisma_client: Optional[PrismaClient], - user_api_key_cache: DualCache, + user_api_key_cache: UserApiKeyCache, proxy_logging_obj: ProxyLogging, valid_token: Optional[UserAPIKeyAuth], ): diff --git a/litellm/proxy/auth/handle_jwt.py b/litellm/proxy/auth/handle_jwt.py index fb86485ffad7..9e7c0abd983d 100644 --- a/litellm/proxy/auth/handle_jwt.py +++ b/litellm/proxy/auth/handle_jwt.py @@ -20,7 +20,6 @@ from jwt.api_jwk import PyJWK from litellm._logging import verbose_proxy_logger -from litellm.caching.caching import DualCache from litellm.constants import DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL from litellm.litellm_core_utils.dot_notation_indexing import get_nested_value from litellm.llms.custom_httpx.httpx_handler import HTTPHandler @@ -46,7 +45,7 @@ ) from litellm.proxy.auth.auth_checks import can_team_access_model from litellm.proxy.auth.route_checks import RouteChecks -from litellm.proxy.common_utils.cache_pydantic_utils import CacheCodec +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.proxy.utils import PrismaClient, ProxyLogging from .auth_checks import ( @@ -74,7 +73,7 @@ class JWTHandler: """ prisma_client: Optional[PrismaClient] - user_api_key_cache: DualCache + user_api_key_cache: UserApiKeyCache # Supported algos: https://pyjwt.readthedocs.io/en/stable/algorithms.html # "Warning: Make sure not to mix symmetric and asymmetric algorithms that interpret # the key in different ways (e.g. HS* and RS*)." @@ -100,7 +99,7 @@ def __init__( def update_environment( self, prisma_client: Optional[PrismaClient], - user_api_key_cache: DualCache, + user_api_key_cache: UserApiKeyCache, litellm_jwtauth: LiteLLM_JWTAuth, leeway: int = 0, ) -> None: @@ -953,7 +952,7 @@ async def find_and_validate_specific_team_id( jwt_handler: JWTHandler, jwt_valid_token: dict, prisma_client: Optional[PrismaClient], - user_api_key_cache: DualCache, + user_api_key_cache: UserApiKeyCache, parent_otel_span: Optional[Span], proxy_logging_obj: ProxyLogging, ) -> Tuple[Optional[str], Optional[LiteLLM_TeamTable]]: @@ -1046,7 +1045,7 @@ async def find_team_with_model_access( route: str, jwt_handler: JWTHandler, prisma_client: Optional[PrismaClient], - user_api_key_cache: DualCache, + user_api_key_cache: UserApiKeyCache, parent_otel_span: Optional[Span], proxy_logging_obj: ProxyLogging, ) -> Tuple[Optional[str], Optional[LiteLLM_TeamTable]]: @@ -1134,7 +1133,7 @@ async def get_objects( valid_user_email: Optional[bool], jwt_handler: JWTHandler, prisma_client: Optional[PrismaClient], - user_api_key_cache: DualCache, + user_api_key_cache: UserApiKeyCache, parent_otel_span: Optional[Span], proxy_logging_obj: ProxyLogging, route: str, @@ -1350,7 +1349,7 @@ async def sync_user_role_and_teams( jwt_valid_token: dict, user_object: Optional[LiteLLM_UserTable], prisma_client: Optional[PrismaClient], - user_api_key_cache: Optional[DualCache] = None, + user_api_key_cache: Optional[UserApiKeyCache] = None, ) -> None: """ Sync user role and team memberships with JWT claims @@ -1378,9 +1377,8 @@ async def sync_user_role_and_teams( if user_api_key_cache is not None: await user_api_key_cache.async_set_cache( key=user_object.user_id, - value=CacheCodec.serialize( - user_object, model_type=LiteLLM_UserTable - ), + value=user_object, + model_type=LiteLLM_UserTable, ttl=DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL, ) @@ -1403,9 +1401,8 @@ async def sync_user_role_and_teams( if user_api_key_cache is not None: await user_api_key_cache.async_set_cache( key=user_object.user_id, - value=CacheCodec.serialize( - user_object, model_type=LiteLLM_UserTable - ), + value=user_object, + model_type=LiteLLM_UserTable, ttl=DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL, ) return None @@ -1417,7 +1414,7 @@ async def _attach_team_from_header_for_admin( request_headers: Optional[dict], jwt_handler: JWTHandler, prisma_client: Optional[PrismaClient], - user_api_key_cache: DualCache, + user_api_key_cache: UserApiKeyCache, parent_otel_span: Optional[Span], proxy_logging_obj: ProxyLogging, ) -> None: @@ -1464,7 +1461,7 @@ async def auth_builder( general_settings: dict, route: str, prisma_client: Optional[PrismaClient], - user_api_key_cache: DualCache, + user_api_key_cache: UserApiKeyCache, parent_otel_span: Optional[Span], proxy_logging_obj: ProxyLogging, request_headers: Optional[dict] = None, diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index f77fdc47398a..30f78a50813b 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -20,7 +20,6 @@ import litellm from litellm._logging import verbose_logger, verbose_proxy_logger from litellm._service_logger import ServiceLogging -from litellm.caching import DualCache from litellm.litellm_core_utils.dd_tracing import tracer from litellm.litellm_core_utils.dot_notation_indexing import get_nested_value from litellm.proxy._types import * @@ -58,7 +57,7 @@ from litellm.proxy.auth.oauth2_proxy_hook import handle_oauth2_proxy_request from litellm.proxy.auth.route_checks import RouteChecks from litellm.proxy.common_utils.cache_coordinator import EventDrivenCacheCoordinator -from litellm.proxy.common_utils.cache_pydantic_utils import CacheCodec +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.proxy.common_utils.http_parsing_utils import ( _read_request_body, _safe_get_request_headers, @@ -328,7 +327,7 @@ def update_valid_token_with_end_user_params( async def _fetch_global_spend_with_event_coordination( cache_key: str, - user_api_key_cache: DualCache, + user_api_key_cache: UserApiKeyCache, prisma_client: PrismaClient, ) -> Optional[float]: """ @@ -351,7 +350,7 @@ async def _load_global_spend() -> Optional[float]: async def get_global_proxy_spend( litellm_proxy_admin_name: str, - user_api_key_cache: DualCache, + user_api_key_cache: UserApiKeyCache, prisma_client: Optional[PrismaClient], token: str, proxy_logging_obj: ProxyLogging, @@ -504,7 +503,7 @@ async def _resolve_jwt_to_virtual_key( jwt_claims: dict, jwt_handler: JWTHandler, prisma_client: Optional[PrismaClient], - user_api_key_cache: DualCache, + user_api_key_cache: UserApiKeyCache, parent_otel_span: Optional[Span], proxy_logging_obj: ProxyLogging, ) -> Optional[UserAPIKeyAuth]: @@ -1315,13 +1314,9 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 if prisma_client is not None: _cache_key = f"{valid_token.team_id}_{valid_token.user_id}" - _cached_member = await user_api_key_cache.async_get_cache( - key=_cache_key - ) - team_member_info: Optional[LiteLLM_TeamMembership] = ( - CacheCodec.deserialize(_cached_member, LiteLLM_TeamMembership) - if _cached_member is not None - else None + team_member_info = await user_api_key_cache.async_get_cache( + key=_cache_key, + model_type=LiteLLM_TeamMembership, ) if team_member_info is None: # read from DB @@ -1342,10 +1337,8 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 ) await user_api_key_cache.async_set_cache( key=_cache_key, - value=CacheCodec.serialize( - team_member_info, - model_type=LiteLLM_TeamMembership, - ), + value=team_member_info, + model_type=LiteLLM_TeamMembership, ttl=5, ) @@ -1497,9 +1490,8 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 team_cache_key = f"team_id:{valid_token.team_id}" await user_api_key_cache.async_set_cache( key=team_cache_key, - value=CacheCodec.serialize( - _team_obj, model_type=LiteLLM_TeamTableCachedObj - ), + value=_team_obj, + model_type=LiteLLM_TeamTableCachedObj, ) # save team table in cache - used for tpm/rpm limiting - tpm_rpm_limiter.py # Fetch project object if key belongs to a project diff --git a/litellm/proxy/common_utils/user_api_key_cache.py b/litellm/proxy/common_utils/user_api_key_cache.py new file mode 100644 index 000000000000..2eebc291bb0f --- /dev/null +++ b/litellm/proxy/common_utils/user_api_key_cache.py @@ -0,0 +1,156 @@ +from __future__ import annotations + +from typing import Any, Optional, Type, TypeVar, Union, cast, overload + +from pydantic import BaseModel + +from litellm._logging import verbose_proxy_logger +from litellm.caching.dual_cache import DualCache +from litellm.proxy.common_utils.cache_pydantic_utils import CacheCodec + +T = TypeVar("T", bound=BaseModel) + + +class UserApiKeyCache(DualCache): + """ + DualCache wrapper for UserAPIKeyAuth-like payloads. + + Stores a Redis-safe JSON payload in BOTH in-memory and Redis to avoid + "memory returns BaseModel, Redis returns dict" format drift. + + When ``model_type`` is provided: + - writes are serialized via ``CacheCodec.serialize(..., model_type=...)`` + - reads are deserialized via ``CacheCodec.deserialize(..., model_type)`` + and return ``Optional[T]``: the model on success, ``None`` on cache miss + **or** if the cached payload fails validation (schema drift). On + validation failure after a cache hit, an error line is emitted via + ``verbose_proxy_logger``. + + When ``model_type`` is omitted, the interface behaves like ``DualCache``: + raw cached payload is returned (dict/str/etc.). + + ``async_set_cache_pipeline`` applies the same untyped Codec pass as omitting + ``model_type`` on ``async_set_cache`` (so ``BaseModel`` rows are dumped before Redis). + """ + + # Overloads: `model_type` must be a real parameter (not only via **kwargs) so + # the untyped branch cannot match calls that pass `model_type=...`. + + @overload + def get_cache( + self, + key, + parent_otel_span=None, + local_only: bool = False, + ) -> Any: ... + + @overload + def get_cache( + self, + key, + parent_otel_span=None, + local_only: bool = False, + *, + model_type: Type[T], + ) -> Optional[T]: ... + + def set_cache(self, key, value, local_only: bool = False, **kwargs): # type: ignore[override] + model_type = cast(Optional[Type[BaseModel]], kwargs.pop("model_type", None)) + payload = CacheCodec.serialize(value, model_type=model_type) + return super().set_cache(key=key, value=payload, local_only=local_only, **kwargs) + + async def async_set_cache(self, key, value, local_only: bool = False, **kwargs): # type: ignore[override] + model_type = cast(Optional[Type[BaseModel]], kwargs.pop("model_type", None)) + payload = CacheCodec.serialize(value, model_type=model_type) + return await super().async_set_cache( + key=key, value=payload, local_only=local_only, **kwargs + ) + + async def async_set_cache_pipeline( # type: ignore[override] + self, cache_list: list, local_only: bool = False, **kwargs + ) -> None: + """ + Batch writes with the same Codec boundary as ``async_set_cache`` without + ``model_type``: ``BaseModel`` values become JSON-safe dicts; dicts/scalars unchanged. + """ + normalized = [ + (key, CacheCodec.serialize(value, model_type=None)) + for key, value in cache_list + ] + return await super().async_set_cache_pipeline( + cache_list=normalized, local_only=local_only, **kwargs + ) + + @overload + async def async_get_cache( + self, + key, + parent_otel_span=None, + local_only: bool = False, + ) -> Any: ... + + @overload + async def async_get_cache( + self, + key, + parent_otel_span=None, + local_only: bool = False, + *, + model_type: Type[T], + ) -> Optional[T]: ... + + def get_cache( # type: ignore[override] + self, + key, + parent_otel_span=None, + local_only: bool = False, + model_type: Optional[Type[BaseModel]] = None, + **kwargs, + ) -> Union[Any, Optional[BaseModel]]: + if model_type is None and "model_type" in kwargs: + model_type = cast(Optional[Type[BaseModel]], kwargs.pop("model_type", None)) + cached = super().get_cache( + key=key, parent_otel_span=parent_otel_span, local_only=local_only, **kwargs + ) + if model_type is None: + return cached + if cached is None: + return None + decoded = CacheCodec.deserialize(cached, model_type=model_type) + if decoded is None: + verbose_proxy_logger.error( + "UserApiKeyCache.get_cache failed to deserialize cached value for " + "key=%r model_type=%s", + key, + getattr(model_type, "__name__", str(model_type)), + ) + return None + return decoded + + async def async_get_cache( # type: ignore[override] + self, + key, + parent_otel_span=None, + local_only: bool = False, + model_type: Optional[Type[BaseModel]] = None, + **kwargs, + ) -> Union[Any, Optional[BaseModel]]: + if model_type is None and "model_type" in kwargs: + model_type = cast(Optional[Type[BaseModel]], kwargs.pop("model_type", None)) + cached = await super().async_get_cache( + key=key, parent_otel_span=parent_otel_span, local_only=local_only, **kwargs + ) + if model_type is None: + return cached + if cached is None: + return None + decoded = CacheCodec.deserialize(cached, model_type=model_type) + if decoded is None: + verbose_proxy_logger.error( + "UserApiKeyCache.async_get_cache failed to deserialize cached value for " + "key=%r model_type=%s", + key, + getattr(model_type, "__name__", str(model_type)), + ) + return None + return decoded diff --git a/litellm/proxy/management_endpoints/access_group_endpoints.py b/litellm/proxy/management_endpoints/access_group_endpoints.py index 30fed0962324..ceaef20a8d0e 100644 --- a/litellm/proxy/management_endpoints/access_group_endpoints.py +++ b/litellm/proxy/management_endpoints/access_group_endpoints.py @@ -9,7 +9,6 @@ LitellmUserRoles, UserAPIKeyAuth, ) -from litellm.proxy.common_utils.cache_pydantic_utils import CacheCodec from litellm.proxy.auth.auth_checks import ( _cache_access_object, _cache_key_object, @@ -237,10 +236,10 @@ async def _patch_key_caches_add_access_group( ) -> None: """Patch cached key objects to include access_group_id.""" for token in key_tokens: - raw = await user_api_key_cache.async_get_cache(key=token) - if raw is None: - continue - cached_key = CacheCodec.deserialize(raw, UserAPIKeyAuth) + cached_key = await user_api_key_cache.async_get_cache( + key=token, + model_type=UserAPIKeyAuth, + ) if cached_key is None: continue if cached_key.access_group_ids is None: @@ -267,10 +266,10 @@ async def _patch_key_caches_remove_access_group( ) -> None: """Patch cached key objects to remove access_group_id.""" for token in key_tokens: - raw = await user_api_key_cache.async_get_cache(key=token) - if raw is None: - continue - cached_key = CacheCodec.deserialize(raw, UserAPIKeyAuth) + cached_key = await user_api_key_cache.async_get_cache( + key=token, + model_type=UserAPIKeyAuth, + ) if cached_key is not None and cached_key.access_group_ids: cached_key.access_group_ids = [ ag for ag in cached_key.access_group_ids if ag != access_group_id diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index 0362c3ff6b2a..c00f7269bbd6 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -39,7 +39,6 @@ import litellm from litellm._logging import verbose_proxy_logger from litellm._uuid import uuid -from litellm.caching import DualCache from litellm.constants import ( LITELLM_UI_SESSION_DURATION, MAX_SPENDLOG_ROWS_TO_QUERY, @@ -70,7 +69,7 @@ UserAPIKeyAuth, ) from litellm.proxy.auth.auth_checks import ExperimentalUIJWTToken, get_user_object -from litellm.proxy.common_utils.cache_pydantic_utils import CacheCodec +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.proxy.auth.auth_utils import _has_user_setup_sso from litellm.proxy.auth.handle_jwt import JWTHandler from litellm.proxy.auth.user_api_key_auth import user_api_key_auth @@ -1051,7 +1050,7 @@ async def get_existing_user_info_from_db( user_id: Optional[str], user_email: Optional[str], prisma_client: PrismaClient, - user_api_key_cache: DualCache, + user_api_key_cache: UserApiKeyCache, proxy_logging_obj: ProxyLogging, ) -> Optional[LiteLLM_UserTable]: try: @@ -1075,7 +1074,7 @@ async def get_existing_user_info_from_db( async def get_user_info_from_db( result: Union[CustomOpenID, OpenID, dict], prisma_client: PrismaClient, - user_api_key_cache: DualCache, + user_api_key_cache: UserApiKeyCache, proxy_logging_obj: ProxyLogging, user_email: Optional[str], user_defined_values: Optional[SSOUserDefinedValues], @@ -1195,7 +1194,7 @@ async def _sync_user_role_from_jwt_role_map( received_response: Optional[dict], user_info: Optional[Union[LiteLLM_UserTable, NewUserResponse]], prisma_client: PrismaClient, - user_api_key_cache: DualCache, + user_api_key_cache: UserApiKeyCache, user_defined_values: Optional[SSOUserDefinedValues], ) -> None: """ @@ -1234,7 +1233,8 @@ async def _sync_user_role_from_jwt_role_map( user_info.user_role = mapped_role.value await user_api_key_cache.async_set_cache( key=user_info.user_id, - value=CacheCodec.serialize(user_info, model_type=LiteLLM_UserTable), + value=user_info, + model_type=LiteLLM_UserTable, ) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 8d199d2c092c..865a46a54780 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -209,6 +209,7 @@ def generate_feedback_box(): from litellm._logging import verbose_proxy_logger, verbose_router_logger from litellm.caching.caching import DualCache, RedisCache from litellm.caching.redis_cluster_cache import RedisClusterCache +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.constants import ( _REALTIME_BODY_CACHE_SIZE, APSCHEDULER_COALESCE, @@ -1598,7 +1599,7 @@ async def root_redirect(): shared_aiohttp_session: Optional["ClientSession"] = ( None # Global shared session for connection reuse ) -user_api_key_cache = DualCache( +user_api_key_cache = UserApiKeyCache( default_in_memory_ttl=UserAPIKeyCacheTTLEnum.in_memory_cache_ttl.value ) spend_counter_cache = DualCache( @@ -2041,19 +2042,13 @@ async def _update_key_cache(token: str, response_cost: float): else: hashed_token = token verbose_proxy_logger.debug("_update_key_cache: hashed_token=%s", hashed_token) - cached_key = await user_api_key_cache.async_get_cache(key=hashed_token) + existing_spend_obj = await user_api_key_cache.async_get_cache( + key=hashed_token, model_type=UserAPIKeyAuth + ) verbose_proxy_logger.debug( - f"_update_key_cache: existing_spend_obj={cached_key}" + f"_update_key_cache: existing_spend_obj={existing_spend_obj}" ) - if cached_key is None: - return - existing_spend_obj = CacheCodec.deserialize(cached_key, UserAPIKeyAuth) if existing_spend_obj is None: - verbose_proxy_logger.warning( - "_update_key_cache: unexpected cached key type %s for hashed_token=%s; skipping spend update", - type(cached_key).__name__, - hashed_token, - ) return existing_spend = existing_spend_obj.spend or 0.0 @@ -2114,14 +2109,10 @@ async def _update_key_cache(token: str, response_cost: float): existing_team_member_spend + response_cost ) - # Update the cost column for the given token (dict for Redis pipeline json.dumps) + # Existing spend_obj is mutated; UserApiKeyCache.async_set_cache_pipeline turns + # BaseModel values into dicts for Redis (same Codec path as async_set_cache). existing_spend_obj.spend = new_spend - values_to_update_in_cache.append( - ( - hashed_token, - CacheCodec.serialize(existing_spend_obj, model_type=UserAPIKeyAuth), - ) - ) + values_to_update_in_cache.append((hashed_token, existing_spend_obj)) ### UPDATE USER SPEND ### async def _update_user_cache(): diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index 8612d243c41f..37646373a039 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -922,19 +922,21 @@ async def test_get_tag_objects_batch(): # Simulate 5 tags: 2 cached, 3 uncached tag_names = ["cached-1", "uncached-1", "cached-2", "uncached-2", "uncached-3"] - # Mock cached tags - cached_tag_1 = { - "tag_name": "cached-1", - "spend": 10.0, - "models": [], - "litellm_budget_table": None, - } - cached_tag_2 = { - "tag_name": "cached-2", - "spend": 20.0, - "models": [], - "litellm_budget_table": None, - } + # Mock cached tags — must be LiteLLM_TagTable instances: the mocked async_get_cache + # bypasses UserApiKeyCache deserialization, so returning plain dicts would flow through + # as dict (production returns models after Codec.deserialize inside the cache). + cached_tag_1 = LiteLLM_TagTable( + tag_name="cached-1", + spend=10.0, + models=[], + litellm_budget_table=None, + ) + cached_tag_2 = LiteLLM_TagTable( + tag_name="cached-2", + spend=20.0, + models=[], + litellm_budget_table=None, + ) # Mock DB response for uncached tags uncached_tag_1 = MagicMock() @@ -980,13 +982,13 @@ async def test_get_tag_objects_batch(): ) # Mock cache behavior - return cached tags, None for uncached - async def mock_get_cache(key): + async def mock_get_cache(*args, **kwargs): + key = kwargs.get("key") if key == "tag:cached-1": return cached_tag_1 - elif key == "tag:cached-2": + if key == "tag:cached-2": return cached_tag_2 - else: - return None + return None mock_cache.async_get_cache = AsyncMock(side_effect=mock_get_cache) mock_cache.async_set_cache = AsyncMock() diff --git a/tests/test_litellm/proxy/common_utils/test_user_api_key_cache.py b/tests/test_litellm/proxy/common_utils/test_user_api_key_cache.py new file mode 100644 index 000000000000..cff331d9b192 --- /dev/null +++ b/tests/test_litellm/proxy/common_utils/test_user_api_key_cache.py @@ -0,0 +1,113 @@ +import json +from typing import Any + +import pytest + +from litellm.caching.redis_cache import RedisCache +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + +class FakeRedisCache(RedisCache): + """ + In-memory fake that enforces the UserApiKeyCache Redis payload contract. + + For user_api_key_cache entries we expect Redis to store a JSON object (dict) + produced by `CacheCodec.serialize(..., model_type=...)`. + + This fake: + - raises TypeError if the value is not a dict + - raises TypeError if the dict is not JSON-serializable + """ + + def __init__(self): # noqa: super().__init__ skipped intentionally + self._store: dict[str, str] = {} + + def set_cache(self, key: str, value: Any, **kwargs): # type: ignore[override] + if not isinstance(value, dict): + raise TypeError("FakeRedisCache only accepts dict payloads") + self._store[key] = json.dumps(value) + return True + + def get_cache(self, key: str, **kwargs): # type: ignore[override] + raw = self._store.get(key) + if raw is None: + return None + return json.loads(raw) + + async def async_set_cache(self, key: str, value: Any, **kwargs): # type: ignore[override] + if not isinstance(value, dict): + raise TypeError("FakeRedisCache only accepts dict payloads") + self._store[key] = json.dumps(value) + return True + + async def async_get_cache(self, key: str, **kwargs): # type: ignore[override] + raw = self._store.get(key) + if raw is None: + return None + return json.loads(raw) + + def delete_cache(self, key: str): # type: ignore[override] + self._store.pop(key, None) + + async def async_delete_cache(self, key: str): # type: ignore[override] + self._store.pop(key, None) + + +def _make_key_obj(token: str = "tok") -> UserAPIKeyAuth: + # Minimal object (UserAPIKeyAuth inherits token from base view). + return UserAPIKeyAuth(token=token) + + +class TestUserApiKeyCache: + @pytest.mark.asyncio + async def test_async_set_typed_stores_serialized_payload_in_memory_and_redis(self): + cache = UserApiKeyCache(redis_cache=FakeRedisCache()) + obj = _make_key_obj("abc") + + await cache.async_set_cache("k", obj, model_type=UserAPIKeyAuth) + + # In-memory hit should still be raw dict (not BaseModel) because wrapper + # stores the serialized payload into both layers. + raw = await cache.in_memory_cache.async_get_cache("k") # type: ignore[union-attr] + assert isinstance(raw, dict) + assert raw["token"] == "abc" + + # Redis should also hold the same serialized dict + redis_raw = await cache.redis_cache.async_get_cache("k") # type: ignore[union-attr] + assert redis_raw == raw + + @pytest.mark.asyncio + async def test_async_get_typed_returns_model_on_valid_hit(self): + cache = UserApiKeyCache(redis_cache=FakeRedisCache()) + await cache.async_set_cache("k", {"token": "abc"}, model_type=UserAPIKeyAuth) + + value = await cache.async_get_cache("k", model_type=UserAPIKeyAuth) + assert value is not None + assert isinstance(value, UserAPIKeyAuth) + assert value.token == "abc" + + @pytest.mark.asyncio + async def test_async_get_typed_returns_none_on_validation_failure_after_hit(self): + cache = UserApiKeyCache(redis_cache=FakeRedisCache()) + + # Bypass UserApiKeyCache.serialize: CacheCodec rejects non-dict cached values + # for dict-based models (deserialize returns None). + await cache.in_memory_cache.async_set_cache( + key="k", value="invalid-payload-not-a-dict" + ) + + value = await cache.async_get_cache("k", model_type=UserAPIKeyAuth) + assert value is None + + def test_fake_redis_cache_rejects_non_json_serializable_values(self): + fake = FakeRedisCache() + + class NotSerializable: + pass + + with pytest.raises(TypeError): + fake.set_cache("k", NotSerializable()) + + with pytest.raises(TypeError): + fake.set_cache("k2", {"ok": NotSerializable()}) diff --git a/tests/test_litellm/proxy/test_redis_auth_cache_flag.py b/tests/test_litellm/proxy/test_redis_auth_cache_flag.py index 261222be9a24..d0cb5ec54657 100644 --- a/tests/test_litellm/proxy/test_redis_auth_cache_flag.py +++ b/tests/test_litellm/proxy/test_redis_auth_cache_flag.py @@ -6,6 +6,7 @@ """ from contextlib import contextmanager +import json from unittest.mock import MagicMock, patch import pytest @@ -29,7 +30,18 @@ class _FakeRedisCache(RedisCache): """ def __init__(self): # noqa: super().__init__ skipped intentionally - pass + self._store = {} + + def set_cache(self, key, value, **kwargs): # type: ignore[override] + # Enforce Redis JSON-serializable payload contract. + self._store[key] = json.dumps(value) + return True + + def get_cache(self, key, **kwargs): # type: ignore[override] + raw = self._store.get(key) + if raw is None: + return None + return json.loads(raw) @contextmanager From 354ad40d15f504dd84ad7c116f561c274287d215 Mon Sep 17 00:00:00 2001 From: harish-berri Date: Tue, 28 Apr 2026 19:30:56 +0000 Subject: [PATCH 13/28] fix DualCache import error in CI and update docstrings for cache coordinator --- litellm/proxy/auth/handle_jwt.py | 2 ++ .../proxy/common_utils/cache_coordinator.py | 22 ++++++++++++++++--- 2 files changed, 21 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/auth/handle_jwt.py b/litellm/proxy/auth/handle_jwt.py index 9e7c0abd983d..8cb871c9221f 100644 --- a/litellm/proxy/auth/handle_jwt.py +++ b/litellm/proxy/auth/handle_jwt.py @@ -6,6 +6,8 @@ JWT token must have 'litellm_proxy_admin' in scope. """ +from __future__ import annotations + import fnmatch import hashlib import os diff --git a/litellm/proxy/common_utils/cache_coordinator.py b/litellm/proxy/common_utils/cache_coordinator.py index 24da9450ab80..92d9916098d8 100644 --- a/litellm/proxy/common_utils/cache_coordinator.py +++ b/litellm/proxy/common_utils/cache_coordinator.py @@ -20,11 +20,27 @@ class AsyncCacheProtocol(Protocol): - """Protocol for cache backends used by EventDrivenCacheCoordinator.""" + """Protocol for cache backends used by EventDrivenCacheCoordinator. - async def async_get_cache(self, key: str, **kwargs: Any) -> Any: ... + Matches ``DualCache`` / ``UserApiKeyCache`` call shapes (explicit optional params + before ``**kwargs``), not only ``(key, **kwargs)``, so overloads validate. + """ + + async def async_get_cache( + self, + key: str, + parent_otel_span: Any = None, + local_only: bool = False, + **kwargs: Any, + ) -> Any: ... - async def async_set_cache(self, key: str, value: Any, **kwargs: Any) -> Any: ... + async def async_set_cache( + self, + key: str, + value: Any, + local_only: bool = False, + **kwargs: Any, + ) -> Any: ... class EventDrivenCacheCoordinator: From 84b6bd60af93158dd7b5219315920222869eb518 Mon Sep 17 00:00:00 2001 From: harish-berri Date: Tue, 28 Apr 2026 21:08:46 +0000 Subject: [PATCH 14/28] update test cases to match new behaviour. The earlier test cases assumed the cache stores a pydantic object --- tests/proxy_unit_tests/test_auth_checks.py | 11 +++- .../test_user_api_key_auth.py | 7 ++- .../test_access_group_endpoints.py | 57 ++++++++++++++----- 3 files changed, 58 insertions(+), 17 deletions(-) diff --git a/tests/proxy_unit_tests/test_auth_checks.py b/tests/proxy_unit_tests/test_auth_checks.py index 86cd5c0c413c..5636a55c95aa 100644 --- a/tests/proxy_unit_tests/test_auth_checks.py +++ b/tests/proxy_unit_tests/test_auth_checks.py @@ -16,6 +16,7 @@ from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.auth.auth_checks import get_end_user_object from litellm.caching.caching import DualCache +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.proxy._types import ( LiteLLM_EndUserTable, LiteLLM_BudgetTable, @@ -48,9 +49,15 @@ async def test_get_end_user_object(customer_spend, customer_budget): litellm_budget_table=_budget, blocked=False, ) - _cache = DualCache() + # UserApiKeyCache applies model_type on get/set; plain DualCache returns raw dicts + # and breaks get_end_user_object's typed async_get_cache path. + _cache = UserApiKeyCache() _key = "end_user_id:{}".format(end_user_id) - _cache.set_cache(key=_key, value=end_user_obj.model_dump()) + await _cache.async_set_cache( + key=_key, + value=end_user_obj, + model_type=LiteLLM_EndUserTable, + ) try: await get_end_user_object( end_user_id=end_user_id, diff --git a/tests/proxy_unit_tests/test_user_api_key_auth.py b/tests/proxy_unit_tests/test_user_api_key_auth.py index f847b8aa6563..4e0bdf074f23 100644 --- a/tests/proxy_unit_tests/test_user_api_key_auth.py +++ b/tests/proxy_unit_tests/test_user_api_key_auth.py @@ -268,7 +268,12 @@ async def test_aaauser_personal_budgets(key_ownership): test_user_cache = getattr(litellm.proxy.proxy_server, "user_api_key_cache") - assert test_user_cache.get_cache(key=hash_token(user_key)) == valid_token + assert ( + test_user_cache.get_cache( + key=hash_token(user_key), model_type=UserAPIKeyAuth + ) + == valid_token + ) try: await user_api_key_auth(request=request, api_key="Bearer " + user_key) diff --git a/tests/test_litellm/proxy/management_endpoints/test_access_group_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_access_group_endpoints.py index cd2eb7895896..016e10859b62 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_access_group_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_access_group_endpoints.py @@ -738,15 +738,32 @@ def test_delete_access_group_patches_cached_team_and_key( return_value=None ) - # Build cached key object (returned from user_api_key_cache) - if key_cache_group_ids is not None: - cached_key = UserAPIKeyAuth( - token="hashed-key-1", - access_group_ids=list(key_cache_group_ids), + # user_api_key_cache is queried both for teams (fallback after dual_cache) and + # hashed keys — return the right stub per ``key``. A single AsyncMock(return_value=key) + # would wrongly serve the key blob for ``team_id:team-1`` and trigger team patching. + # Use a synchronous side_effect (not async def): AsyncMock awaits coroutine side_effects + # inconsistently across Python/unittest versions; sync returns are awaited as immediate results. + def user_cache_get_side_effect(*args, **kwargs): + cache_key = ( + kwargs.get("key") if "key" in kwargs else (args[0] if args else None) ) - mock_cache.async_get_cache = AsyncMock(return_value=cached_key) - else: - mock_cache.async_get_cache = AsyncMock(return_value=None) + if cache_key == "team_id:team-1": + if team_cache_group_ids is None: + return None + return LiteLLM_TeamTableCachedObj( + team_id="team-1", + access_group_ids=list(team_cache_group_ids), + ) + if cache_key == "hashed-key-1": + if key_cache_group_ids is None: + return None + return UserAPIKeyAuth( + token="hashed-key-1", + access_group_ids=list(key_cache_group_ids), + ) + return None + + mock_cache.async_get_cache = AsyncMock(side_effect=user_cache_get_side_effect) resp = client.delete("/v1/access_group/ag-to-delete") assert resp.status_code == 204 @@ -803,7 +820,7 @@ def test_delete_access_group_patches_cached_team_and_key( def test_delete_access_group_patches_key_cached_as_dict(client_and_mocks): - """Delete correctly patches a key cached as a raw dict (not UserAPIKeyAuth).""" + """Delete patches key cache — mock returns UserAPIKeyAuth (what UserApiKeyCache emits after deserialize).""" client, mock_prisma, mock_access_group_table, mock_cache, mock_proxy_logging = ( client_and_mocks ) @@ -826,12 +843,24 @@ def test_delete_access_group_patches_key_cached_as_dict(client_and_mocks): return_value=None ) - # Key cached as a plain dict (as can happen with Redis serialization) + # Serialized shape from Redis dict; UserApiKeyCache.async_get_cache(model_type=...) yields a model — simulate that. + cached_key_payload = { + "token": "hashed-key-dict", + "access_group_ids": ["ag-to-delete", "ag-other"], + } + + def user_cache_get_dict_when_key_matches(*args, **kwargs): + cache_key = ( + kwargs.get("key") if "key" in kwargs else (args[0] if args else None) + ) + if cache_key == "team_id:team-1": + return None + if cache_key == "hashed-key-dict": + return UserAPIKeyAuth.model_validate(cached_key_payload) + return None + mock_cache.async_get_cache = AsyncMock( - return_value={ - "token": "hashed-key-dict", - "access_group_ids": ["ag-to-delete", "ag-other"], - } + side_effect=user_cache_get_dict_when_key_matches ) resp = client.delete("/v1/access_group/ag-to-delete") From 18996326ef0450139d4c20b29366230b249bbc80 Mon Sep 17 00:00:00 2001 From: harish-berri Date: Tue, 28 Apr 2026 21:17:53 +0000 Subject: [PATCH 15/28] update test cases to fix handle_jwt test cases --- tests/test_litellm/proxy/auth/test_handle_jwt.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/tests/test_litellm/proxy/auth/test_handle_jwt.py b/tests/test_litellm/proxy/auth/test_handle_jwt.py index cdb9ae4d9ab2..a5764f408091 100644 --- a/tests/test_litellm/proxy/auth/test_handle_jwt.py +++ b/tests/test_litellm/proxy/auth/test_handle_jwt.py @@ -403,9 +403,9 @@ async def test_sync_user_role_and_teams_cache_invalidation_on_role_change(): mock_cache.async_set_cache.assert_called_once() call_kwargs = mock_cache.async_set_cache.call_args assert call_kwargs.kwargs["key"] == "u1" - assert ( - call_kwargs.kwargs["value"]["user_role"] == LitellmUserRoles.PROXY_ADMIN.value - ) + assert isinstance(call_kwargs.kwargs["value"], LiteLLM_UserTable) + assert call_kwargs.kwargs["value"].user_role == LitellmUserRoles.PROXY_ADMIN.value + assert call_kwargs.kwargs["model_type"] == LiteLLM_UserTable @pytest.mark.asyncio @@ -450,7 +450,9 @@ async def test_sync_user_role_and_teams_cache_invalidation_on_team_change(): mock_cache.async_set_cache.assert_called_once() call_kwargs = mock_cache.async_set_cache.call_args assert call_kwargs.kwargs["key"] == "u1" - assert set(call_kwargs.kwargs["value"]["teams"]) == {"team1", "team2"} + assert isinstance(call_kwargs.kwargs["value"], LiteLLM_UserTable) + assert set(call_kwargs.kwargs["value"].teams) == {"team1", "team2"} + assert call_kwargs.kwargs["model_type"] == LiteLLM_UserTable @pytest.mark.asyncio From 13ffb9aed5c4bb2f5aa99f07a235aa07a344f0d4 Mon Sep 17 00:00:00 2001 From: harish-berri Date: Tue, 28 Apr 2026 21:27:55 +0000 Subject: [PATCH 16/28] update auth metric increment comment from greptile --- litellm/proxy/auth/auth_metrics.py | 14 +++++++++++--- 1 file changed, 11 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/auth/auth_metrics.py b/litellm/proxy/auth/auth_metrics.py index 0b463cc920ff..7284ba427366 100644 --- a/litellm/proxy/auth/auth_metrics.py +++ b/litellm/proxy/auth/auth_metrics.py @@ -13,6 +13,7 @@ """ from litellm._logging import verbose_proxy_logger +from litellm.types.integrations.prometheus import UserAPIKeyLabelValues class AuthMetrics: @@ -42,9 +43,16 @@ def inc_combined_view_query(hashed_token: str) -> None: try: prom = AuthMetrics._get_prom() if prom is not None: - prom.litellm_auth_combined_view_queries_total.labels( - hashed_api_key=hashed_token - ).inc() + # Counter labelnames include ``hashed_api_key`` plus any + # ``custom_prometheus_metadata_labels`` / ``custom_prometheus_tags`` + # (see ``PrometheusMetricLabels.get_labels``). Use the same + # ``_inc_labeled_counter`` + ``prometheus_label_factory`` path as + # other metrics so label cardinality always matches registration. + prom._inc_labeled_counter( + prom.litellm_auth_combined_view_queries_total, + "litellm_auth_combined_view_queries_total", + UserAPIKeyLabelValues(hashed_api_key=hashed_token), + ) except Exception as e: verbose_proxy_logger.debug( "AuthMetrics.inc_combined_view_query: failed to increment counter: %s", From 90396431cb6e1f6d9d4256fb168e6925e2203c3d Mon Sep 17 00:00:00 2001 From: harish-berri Date: Tue, 28 Apr 2026 21:54:36 +0000 Subject: [PATCH 17/28] lint styling update --- litellm/proxy/proxy_server.py | 16 ++++++++++------ 1 file changed, 10 insertions(+), 6 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 865a46a54780..f54df469d954 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -2129,7 +2129,9 @@ async def _update_user_cache(): if cached_user is None: # do nothing if there is no cache value return - existing_spend_obj = CacheCodec.deserialize(cached_user, LiteLLM_UserTable) + existing_spend_obj = CacheCodec.deserialize( + cached_user, LiteLLM_UserTable + ) if existing_spend_obj is None: return verbose_proxy_logger.debug( @@ -3606,21 +3608,23 @@ async def load_config( # noqa: PLR0915 verbose_proxy_logger.critical( "LITELLM_MASTER_KEY is not set! All requests will be treated as INTERNAL_USER with no admin access. Set LITELLM_MASTER_KEY for production use." ) - ### USER API KEY CACHE IN-MEMORY TTL ### + ### USER API KEY CACHE TTL (in-memory + Redis when Redis auth sharing is enabled) ### user_api_key_cache_ttl = general_settings.get( "user_api_key_cache_ttl", None ) if user_api_key_cache_ttl is not None: + ttl = float(user_api_key_cache_ttl) + # Mirror TTL on Redis as well when ``litellm_settings.enable_redis_auth_cache`` + # attaches Redis to ``user_api_key_cache``; otherwise DualCache misses in + # memory fall back to a key that outlasts ``user_api_key_cache_ttl``. user_api_key_cache.update_cache_ttl( - default_in_memory_ttl=float(user_api_key_cache_ttl), - default_redis_ttl=None, # user_api_key_cache uses in-memory TTL only; Redis not configured for key lookups + default_in_memory_ttl=ttl, + default_redis_ttl=ttl, ) ### PKCE MULTI-INSTANCE PREREQUISITE CHECK ### # PKCE verifiers are stored in redis_usage_cache when available so they can # be read back by any instance (not just the one that started the auth flow). - # user_api_key_cache is intentionally left in-memory-only to avoid routing - # all API-key lookups through Redis. use_pkce = os.getenv("GENERIC_CLIENT_USE_PKCE", "false").lower() == "true" if use_pkce and redis_usage_cache is None: global _pkce_no_redis_warning_emitted From 33d9190747203486cbc3fa7a8a702e78c5e20527 Mon Sep 17 00:00:00 2001 From: harish-berri Date: Wed, 29 Apr 2026 00:27:05 +0000 Subject: [PATCH 18/28] correct team_id key bug --- litellm/proxy/auth/user_api_key_auth.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 30f78a50813b..4c9861b540bd 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -1485,11 +1485,8 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 # Only cache when the key is a real team_id (non-team keys must not use key=None). if valid_token.team_id is not None and _team_obj is not None: - # Match get_team_object / spend counters: "team_id:{id}". Serialize for Redis - # (json.dumps) — same as _cache_team_object in auth_checks. - team_cache_key = f"team_id:{valid_token.team_id}" await user_api_key_cache.async_set_cache( - key=team_cache_key, + key=valid_token.team_id, value=_team_obj, model_type=LiteLLM_TeamTableCachedObj, ) # save team table in cache - used for tpm/rpm limiting - tpm_rpm_limiter.py From b6e13cb70abbf64fde83dfa3eba9c324e730a5e2 Mon Sep 17 00:00:00 2001 From: harish-berri Date: Wed, 29 Apr 2026 01:20:07 +0000 Subject: [PATCH 19/28] fix linting --- litellm/proxy/common_utils/user_api_key_cache.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/common_utils/user_api_key_cache.py b/litellm/proxy/common_utils/user_api_key_cache.py index 2eebc291bb0f..7c42aa545d09 100644 --- a/litellm/proxy/common_utils/user_api_key_cache.py +++ b/litellm/proxy/common_utils/user_api_key_cache.py @@ -57,7 +57,9 @@ def get_cache( def set_cache(self, key, value, local_only: bool = False, **kwargs): # type: ignore[override] model_type = cast(Optional[Type[BaseModel]], kwargs.pop("model_type", None)) payload = CacheCodec.serialize(value, model_type=model_type) - return super().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, value, local_only: bool = False, **kwargs): # type: ignore[override] model_type = cast(Optional[Type[BaseModel]], kwargs.pop("model_type", None)) From bff26f6c014f2095d4a6447217aab61b44bb2582 Mon Sep 17 00:00:00 2001 From: harish-berri Date: Wed, 29 Apr 2026 01:39:27 +0000 Subject: [PATCH 20/28] fix unused imports --- litellm/proxy/auth/auth_checks.py | 1 - 1 file changed, 1 deletion(-) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 8bd8eb794e02..333684eae562 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -19,7 +19,6 @@ import litellm from litellm._logging import verbose_proxy_logger -from litellm.caching.caching import DualCache from litellm.caching.dual_cache import LimitedSizeOrderedDict from litellm.constants import ( CLI_JWT_EXPIRATION_HOURS, From 54f762eeebf773c8b5783f3b0c24502c60b707d8 Mon Sep 17 00:00:00 2001 From: harish-berri Date: Thu, 30 Apr 2026 19:05:04 +0000 Subject: [PATCH 21/28] remove unrelated debugger tooling --- .vscode/launch.json | 24 ------------------------ 1 file changed, 24 deletions(-) delete mode 100644 .vscode/launch.json diff --git a/.vscode/launch.json b/.vscode/launch.json deleted file mode 100644 index e48948319b02..000000000000 --- a/.vscode/launch.json +++ /dev/null @@ -1,24 +0,0 @@ -{ - "version": "0.2.0", - "configurations": [ - { - "name": "LiteLLM proxy (debug)", - "type": "debugpy", - "request": "launch", - "program": "${workspaceFolder}/litellm/proxy/proxy_cli.py", - "cwd": "${workspaceFolder}", - "args": [ - "--host", - "0.0.0.0", - "--config", - "dev_config.yaml", - "--port", - "4000", - "--num_workers", - "1" - ], - "console": "integratedTerminal", - "justMyCode": false - } - ] -} \ No newline at end of file From 4dc950a24252f48bf3ac7af8c1738c5f320441c6 Mon Sep 17 00:00:00 2001 From: harish-berri Date: Thu, 30 Apr 2026 19:11:29 +0000 Subject: [PATCH 22/28] fix linting --- litellm/proxy/auth/handle_jwt.py | 2 +- litellm/proxy/auth/user_api_key_auth.py | 1 - 2 files changed, 1 insertion(+), 2 deletions(-) diff --git a/litellm/proxy/auth/handle_jwt.py b/litellm/proxy/auth/handle_jwt.py index c55cd7bb05e5..71411bed7fda 100644 --- a/litellm/proxy/auth/handle_jwt.py +++ b/litellm/proxy/auth/handle_jwt.py @@ -1460,7 +1460,7 @@ async def _resolve_single_team_fallback( user_object: Optional[LiteLLM_UserTable], user_id: Optional[str], prisma_client: Optional[PrismaClient], - user_api_key_cache: DualCache, + user_api_key_cache: UserApiKeyCache, parent_otel_span: Optional[Span], proxy_logging_obj: ProxyLogging, team_id_upsert: Optional[bool], diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index a94abbc9d427..d87bd9e7d59f 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -20,7 +20,6 @@ import litellm from litellm._logging import verbose_logger, verbose_proxy_logger from litellm._service_logger import ServiceLogging -from litellm.caching import DualCache from litellm.constants import LITELLM_PROXY_MASTER_KEY_ALIAS from litellm.litellm_core_utils.dd_tracing import tracer from litellm.litellm_core_utils.dot_notation_indexing import get_nested_value From 1f72ae5442b1198c16d607173b9ebee2eb8ae363 Mon Sep 17 00:00:00 2001 From: harish-berri Date: Thu, 30 Apr 2026 22:06:33 +0000 Subject: [PATCH 23/28] fix: add type hint ignore for user_api_key_cache and update docstring for EventDrivenCacheCoordinator --- litellm/proxy/auth/user_api_key_auth.py | 2 +- litellm/proxy/common_utils/cache_coordinator.py | 5 ++++- 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index d87bd9e7d59f..b13e6cebf970 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -345,7 +345,7 @@ async def _load_global_spend() -> Optional[float]: return await _global_spend_coordinator.get_or_load( cache_key=cache_key, - cache=user_api_key_cache, + cache=user_api_key_cache, # pyright: ignore[reportArgumentType] load_fn=_load_global_spend, ) diff --git a/litellm/proxy/common_utils/cache_coordinator.py b/litellm/proxy/common_utils/cache_coordinator.py index 92d9916098d8..fc46ab6d78ed 100644 --- a/litellm/proxy/common_utils/cache_coordinator.py +++ b/litellm/proxy/common_utils/cache_coordinator.py @@ -52,6 +52,9 @@ class EventDrivenCacheCoordinator: - Other requests: wait for the signal, then read from cache. Create one instance per resource (e.g. one for global spend, one for feature flags). + + Args: + log_prefix: Prefix for debug log messages. """ def __init__(self, log_prefix: str = "[CACHE]"): @@ -150,7 +153,7 @@ async def _load_and_cache( elapsed_ms, value, ) - + await cache.async_set_cache(key=cache_key, value=value) if self._log_prefix: verbose_proxy_logger.debug("%s Result cached", self._log_prefix) From 76fb9d3c966516ca0a64a64a05a3687661b7c622 Mon Sep 17 00:00:00 2001 From: harish-berri Date: Thu, 30 Apr 2026 22:13:11 +0000 Subject: [PATCH 24/28] refactor: remove AuthMetrics and related combined_view query metrics This commit deletes the AuthMetrics class and its associated methods, which were responsible for tracking combined_view SQL query metrics. The PrometheusLogger integration has been updated to remove references to these metrics, streamlining the codebase. Additionally, minor whitespace adjustments were made in the cache coordinator for consistency. --- litellm/integrations/prometheus.py | 13 ---- litellm/proxy/auth/auth_checks.py | 2 - litellm/proxy/auth/auth_metrics.py | 60 ------------------- .../proxy/common_utils/cache_coordinator.py | 2 +- litellm/types/integrations/prometheus.py | 7 --- 5 files changed, 1 insertion(+), 83 deletions(-) delete mode 100644 litellm/proxy/auth/auth_metrics.py diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index 1a899cf0e838..d9e57ee7ceec 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -266,19 +266,6 @@ def __init__( # noqa: PLR0915 # LiteLLM Virtual API KEY metrics ######################################## - # Auth DB load diagnostic: count direct combined_view SQL queries. - # Each increment means a virtual-key cache miss that hit the DB. - # Useful for validating that enable_redis_auth_cache is working. - self.litellm_auth_combined_view_queries_total = self._counter_factory( - "litellm_auth_combined_view_queries_total", - "Number of times the combined_view SQL query was issued for virtual-key auth. " - "Each count is a cache miss that hit the database. Use to validate " - "enable_redis_auth_cache is reducing DB load.", - labelnames=self.get_labels_for_metric( - "litellm_auth_combined_view_queries_total" - ), - ) - # Remaining MODEL RPM limit for API Key self.litellm_remaining_api_key_requests_for_model = self._gauge_factory( "litellm_remaining_api_key_requests_for_model", diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 9266a917c5f0..113a8f538c0d 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -59,7 +59,6 @@ SpecialModelNames, UserAPIKeyAuth, ) -from litellm.proxy.auth.auth_metrics import AuthMetrics from litellm.proxy.auth.route_checks import RouteChecks from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler from litellm.proxy.guardrails.tool_name_extraction import ( @@ -2262,7 +2261,6 @@ async def _fetch_key_object_from_db_with_reconnect( Fetch key object from DB and retry once if a DB connection error can be healed. """ try: - AuthMetrics.inc_combined_view_query(hashed_token) return await prisma_client.get_data( token=hashed_token, table_name="combined_view", diff --git a/litellm/proxy/auth/auth_metrics.py b/litellm/proxy/auth/auth_metrics.py deleted file mode 100644 index 7284ba427366..000000000000 --- a/litellm/proxy/auth/auth_metrics.py +++ /dev/null @@ -1,60 +0,0 @@ -""" -Prometheus metric helpers for the auth layer. - -All metrics are thin wrappers around the shared ``PrometheusLogger`` instance so -that every counter follows the same registration path (``_counter_factory``, -label-filter config) as the rest of LiteLLM's metrics. - -Usage:: - - from litellm.proxy.auth.auth_metrics import AuthMetrics - - AuthMetrics.inc_combined_view_query(hashed_token="sk-xxx") -""" - -from litellm._logging import verbose_proxy_logger -from litellm.types.integrations.prometheus import UserAPIKeyLabelValues - - -class AuthMetrics: - """Static helpers for incrementing auth-layer Prometheus counters.""" - - @staticmethod - def _get_prom(): - """Return the active PrometheusLogger, or None if Prometheus is not configured.""" - try: - from litellm.router_utils.cooldown_callbacks import ( - _get_prometheus_logger_from_callbacks, - ) - - return _get_prometheus_logger_from_callbacks() - except Exception: - return None - - @staticmethod - def inc_combined_view_query(hashed_token: str) -> None: - """ - Increment ``litellm_auth_combined_view_queries_total``. - - Called once per virtual-key DB lookup (combined_view query). Each - increment represents a cache miss that hit the database — use this to - validate that ``enable_redis_auth_cache`` is reducing DB load. - """ - try: - prom = AuthMetrics._get_prom() - if prom is not None: - # Counter labelnames include ``hashed_api_key`` plus any - # ``custom_prometheus_metadata_labels`` / ``custom_prometheus_tags`` - # (see ``PrometheusMetricLabels.get_labels``). Use the same - # ``_inc_labeled_counter`` + ``prometheus_label_factory`` path as - # other metrics so label cardinality always matches registration. - prom._inc_labeled_counter( - prom.litellm_auth_combined_view_queries_total, - "litellm_auth_combined_view_queries_total", - UserAPIKeyLabelValues(hashed_api_key=hashed_token), - ) - except Exception as e: - verbose_proxy_logger.debug( - "AuthMetrics.inc_combined_view_query: failed to increment counter: %s", - e, - ) diff --git a/litellm/proxy/common_utils/cache_coordinator.py b/litellm/proxy/common_utils/cache_coordinator.py index fc46ab6d78ed..abb0402d3b90 100644 --- a/litellm/proxy/common_utils/cache_coordinator.py +++ b/litellm/proxy/common_utils/cache_coordinator.py @@ -153,7 +153,7 @@ async def _load_and_cache( elapsed_ms, value, ) - + await cache.async_set_cache(key=cache_key, value=value) if self._log_prefix: verbose_proxy_logger.debug("%s Result cached", self._log_prefix) diff --git a/litellm/types/integrations/prometheus.py b/litellm/types/integrations/prometheus.py index a4573535f725..43a287f29bc1 100644 --- a/litellm/types/integrations/prometheus.py +++ b/litellm/types/integrations/prometheus.py @@ -228,8 +228,6 @@ class UserAPIKeyLabelNames(Enum): "litellm_guardrail_latency_seconds", "litellm_guardrail_errors_total", "litellm_guardrail_requests_total", - # Auth DB diagnostic metrics - "litellm_auth_combined_view_queries_total", # Cache metrics "litellm_cache_hits_metric", "litellm_cache_misses_metric", @@ -309,11 +307,6 @@ class PrometheusMetricLabels: litellm_guardrail_errors_total: List[str] = [] litellm_guardrail_requests_total: List[str] = [] - # Auth DB diagnostic - label by key so you can see which virtual key causes DB hits - litellm_auth_combined_view_queries_total = [ - UserAPIKeyLabelNames.API_KEY_HASH.value, - ] - litellm_proxy_total_requests_metric = [ UserAPIKeyLabelNames.END_USER.value, UserAPIKeyLabelNames.API_KEY_HASH.value, From 3bff192c2f55a0df04d48f0c2bc7512f38eb959d Mon Sep 17 00:00:00 2001 From: harish-berri Date: Thu, 30 Apr 2026 22:35:23 +0000 Subject: [PATCH 25/28] fix linting --- .../proxy/common_utils/user_api_key_cache.py | 36 ++++++++++--------- .../key_management_endpoints.py | 20 +++++------ litellm/proxy/proxy_server.py | 8 ++--- 3 files changed, 34 insertions(+), 30 deletions(-) diff --git a/litellm/proxy/common_utils/user_api_key_cache.py b/litellm/proxy/common_utils/user_api_key_cache.py index 7c42aa545d09..c6b830371ba3 100644 --- a/litellm/proxy/common_utils/user_api_key_cache.py +++ b/litellm/proxy/common_utils/user_api_key_cache.py @@ -39,20 +39,22 @@ class UserApiKeyCache(DualCache): @overload def get_cache( self, - key, - parent_otel_span=None, + key: Any, + parent_otel_span: Any = None, local_only: bool = False, - ) -> Any: ... + *, + model_type: Type[T], + **kwargs: Any, + ) -> Optional[T]: ... @overload def get_cache( self, - key, - parent_otel_span=None, + key: Any, + parent_otel_span: Any = None, local_only: bool = False, - *, - model_type: Type[T], - ) -> Optional[T]: ... + **kwargs: Any, + ) -> Any: ... def set_cache(self, key, value, local_only: bool = False, **kwargs): # type: ignore[override] model_type = cast(Optional[Type[BaseModel]], kwargs.pop("model_type", None)) @@ -86,20 +88,22 @@ async def async_set_cache_pipeline( # type: ignore[override] @overload async def async_get_cache( self, - key, - parent_otel_span=None, + key: Any, + parent_otel_span: Any = None, local_only: bool = False, - ) -> Any: ... + *, + model_type: Type[T], + **kwargs: Any, + ) -> Optional[T]: ... @overload async def async_get_cache( self, - key, - parent_otel_span=None, + key: Any, + parent_otel_span: Any = None, local_only: bool = False, - *, - model_type: Type[T], - ) -> Optional[T]: ... + **kwargs: Any, + ) -> Any: ... def get_cache( # type: ignore[override] self, diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 8129fb0de5e3..ae38b389b465 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -27,7 +27,7 @@ import litellm from litellm._logging import verbose_proxy_logger from litellm._uuid import uuid -from litellm.caching import DualCache +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.constants import ( LENGTH_OF_LITELLM_GENERATED_KEY, LITELLM_PROXY_ADMIN_NAME, @@ -1058,7 +1058,7 @@ async def _check_project_key_limits( project_id: str, data: Union[GenerateKeyRequest, UpdateKeyRequest], prisma_client: PrismaClient, - user_api_key_cache: DualCache, + user_api_key_cache: UserApiKeyCache, ) -> None: """ Validate that key's models and budget respect its project's limits. @@ -1833,7 +1833,7 @@ async def _process_single_key_update( user_api_key_dict: UserAPIKeyAuth, litellm_changed_by: Optional[str], prisma_client: Optional[PrismaClient], - user_api_key_cache: DualCache, + user_api_key_cache: UserApiKeyCache, proxy_logging_obj: Any, llm_router: Optional[Router], user_custom_key_update: Optional[Callable] = None, @@ -3297,7 +3297,7 @@ async def _team_key_deletion_check( user_api_key_dict: UserAPIKeyAuth, key_info: LiteLLM_VerificationToken, prisma_client: PrismaClient, - user_api_key_cache: DualCache, + user_api_key_cache: UserApiKeyCache, ): is_team_key = _is_team_key(data=key_info) @@ -3340,7 +3340,7 @@ async def _team_key_deletion_check( async def can_modify_verification_token( key_info: LiteLLM_VerificationToken, - user_api_key_cache: DualCache, + user_api_key_cache: UserApiKeyCache, user_api_key_dict: UserAPIKeyAuth, prisma_client: PrismaClient, ) -> bool: @@ -3414,7 +3414,7 @@ async def can_modify_verification_token( async def delete_verification_tokens( tokens: List, - user_api_key_cache: DualCache, + user_api_key_cache: UserApiKeyCache, user_api_key_dict: UserAPIKeyAuth, litellm_changed_by: Optional[str] = None, ) -> Tuple[Optional[Dict], List[LiteLLM_VerificationToken]]: @@ -3604,7 +3604,7 @@ async def _persist_deleted_verification_tokens( async def delete_key_aliases( key_aliases: List[str], - user_api_key_cache: DualCache, + user_api_key_cache: UserApiKeyCache, prisma_client: PrismaClient, user_api_key_dict: UserAPIKeyAuth, litellm_changed_by: Optional[str] = None, @@ -3850,7 +3850,7 @@ async def _execute_virtual_key_regeneration( data: Optional[RegenerateKeyRequest], user_api_key_dict: UserAPIKeyAuth, litellm_changed_by: Optional[str], - user_api_key_cache: DualCache, + user_api_key_cache: UserApiKeyCache, proxy_logging_obj: ProxyLogging, ) -> GenerateKeyResponse: """Generate new token, update DB, invalidate cache, and return response.""" @@ -4140,7 +4140,7 @@ async def _check_proxy_or_team_admin_for_key( key_in_db: LiteLLM_VerificationToken, user_api_key_dict: UserAPIKeyAuth, prisma_client: PrismaClient, - user_api_key_cache: DualCache, + user_api_key_cache: UserApiKeyCache, ) -> None: if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value: return @@ -5161,7 +5161,7 @@ async def _check_key_admin_access( user_api_key_dict: UserAPIKeyAuth, hashed_token: str, prisma_client: Any, - user_api_key_cache: DualCache, + user_api_key_cache: UserApiKeyCache, route: str, ) -> None: """ diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 342e1c4ab37b..06fdef295d81 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -1613,7 +1613,7 @@ async def root_redirect(): shared_aiohttp_session: Optional["ClientSession"] = ( None # Global shared session for connection reuse ) -user_api_key_cache = UserApiKeyCache( +user_api_key_cache: UserApiKeyCache = UserApiKeyCache( default_in_memory_ttl=UserAPIKeyCacheTTLEnum.in_memory_cache_ttl.value ) spend_counter_cache = DualCache( @@ -6318,7 +6318,7 @@ def _initialize_jwt_auth( cls, general_settings: dict, prisma_client: Optional[PrismaClient], - user_api_key_cache: DualCache, + user_api_key_cache: UserApiKeyCache, ): """Initialize JWT auth on startup""" if general_settings.get("litellm_jwtauth", None) is not None: @@ -6367,7 +6367,7 @@ def _add_proxy_budget_to_db(cls, litellm_proxy_budget_name: str): async def _warm_global_spend_cache( cls, litellm_proxy_admin_name: str, - user_api_key_cache: DualCache, + user_api_key_cache: UserApiKeyCache, prisma_client: PrismaClient, ) -> None: """Warm global spend cache once at startup to reduce impact of first wave of requests.""" @@ -7007,7 +7007,7 @@ async def _setup_prisma_client( cls, database_url: Optional[str], proxy_logging_obj: ProxyLogging, - user_api_key_cache: DualCache, + user_api_key_cache: UserApiKeyCache, ) -> Optional[PrismaClient]: """ - Sets up prisma client From 6fb1de9763a27be2a1ca41d96177d79fbd870329 Mon Sep 17 00:00:00 2001 From: harish-berri Date: Thu, 30 Apr 2026 22:40:52 +0000 Subject: [PATCH 26/28] refactor: replace DualCache with UserApiKeyCache in multiple modules This commit updates the codebase to replace instances of DualCache with UserApiKeyCache in various files, including utils, expired_ui_session_key_cleanup_manager, and team_member_permission_checks. Additionally, it enhances the UserApiKeyCache class with new methods for cache management, improving type safety and consistency across the application. --- .../expired_ui_session_key_cleanup_manager.py | 4 +- .../proxy/common_utils/user_api_key_cache.py | 104 +++++++++--------- .../team_member_permission_checks.py | 4 +- litellm/proxy/utils.py | 5 +- 4 files changed, 59 insertions(+), 58 deletions(-) diff --git a/litellm/proxy/common_utils/expired_ui_session_key_cleanup_manager.py b/litellm/proxy/common_utils/expired_ui_session_key_cleanup_manager.py index c25d85331285..67a24567461e 100644 --- a/litellm/proxy/common_utils/expired_ui_session_key_cleanup_manager.py +++ b/litellm/proxy/common_utils/expired_ui_session_key_cleanup_manager.py @@ -8,7 +8,7 @@ from typing import Any, Dict, List, Optional from litellm._logging import verbose_proxy_logger -from litellm.caching import DualCache +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.constants import ( EXPIRED_UI_SESSION_KEY_CLEANUP_JOB_NAME, LITELLM_EXPIRED_UI_SESSION_KEY_CLEANUP_BATCH_SIZE, @@ -31,7 +31,7 @@ class ExpiredUISessionKeyCleanupManager: def __init__( self, prisma_client: PrismaClient, - user_api_key_cache: DualCache, + user_api_key_cache: UserApiKeyCache, pod_lock_manager=None, ): self.prisma_client = prisma_client diff --git a/litellm/proxy/common_utils/user_api_key_cache.py b/litellm/proxy/common_utils/user_api_key_cache.py index c6b830371ba3..914be364579a 100644 --- a/litellm/proxy/common_utils/user_api_key_cache.py +++ b/litellm/proxy/common_utils/user_api_key_cache.py @@ -31,10 +31,10 @@ class UserApiKeyCache(DualCache): ``async_set_cache_pipeline`` applies the same untyped Codec pass as omitting ``model_type`` on ``async_set_cache`` (so ``BaseModel`` rows are dumped before Redis). - """ - # Overloads: `model_type` must be a real parameter (not only via **kwargs) so - # the untyped branch cannot match calls that pass `model_type=...`. + ``get_cache`` / ``async_get_cache`` overloads and implementations must be contiguous + (no other methods in between) so mypy resolves ``@overload`` + implementation correctly. + """ @overload def get_cache( @@ -56,55 +56,6 @@ def get_cache( **kwargs: Any, ) -> Any: ... - def set_cache(self, key, value, local_only: bool = False, **kwargs): # type: ignore[override] - model_type = cast(Optional[Type[BaseModel]], kwargs.pop("model_type", None)) - payload = CacheCodec.serialize(value, model_type=model_type) - return super().set_cache( - key=key, value=payload, local_only=local_only, **kwargs - ) - - async def async_set_cache(self, key, value, local_only: bool = False, **kwargs): # type: ignore[override] - model_type = cast(Optional[Type[BaseModel]], kwargs.pop("model_type", None)) - payload = CacheCodec.serialize(value, model_type=model_type) - return await super().async_set_cache( - key=key, value=payload, local_only=local_only, **kwargs - ) - - async def async_set_cache_pipeline( # type: ignore[override] - self, cache_list: list, local_only: bool = False, **kwargs - ) -> None: - """ - Batch writes with the same Codec boundary as ``async_set_cache`` without - ``model_type``: ``BaseModel`` values become JSON-safe dicts; dicts/scalars unchanged. - """ - normalized = [ - (key, CacheCodec.serialize(value, model_type=None)) - for key, value in cache_list - ] - return await super().async_set_cache_pipeline( - cache_list=normalized, local_only=local_only, **kwargs - ) - - @overload - async def async_get_cache( - self, - key: Any, - parent_otel_span: Any = None, - local_only: bool = False, - *, - model_type: Type[T], - **kwargs: Any, - ) -> Optional[T]: ... - - @overload - async def async_get_cache( - self, - key: Any, - parent_otel_span: Any = None, - local_only: bool = False, - **kwargs: Any, - ) -> Any: ... - def get_cache( # type: ignore[override] self, key, @@ -133,6 +84,26 @@ def get_cache( # type: ignore[override] return None return decoded + @overload + async def async_get_cache( + self, + key: Any, + parent_otel_span: Any = None, + local_only: bool = False, + *, + model_type: Type[T], + **kwargs: Any, + ) -> Optional[T]: ... + + @overload + async def async_get_cache( + self, + key: Any, + parent_otel_span: Any = None, + local_only: bool = False, + **kwargs: Any, + ) -> Any: ... + async def async_get_cache( # type: ignore[override] self, key, @@ -160,3 +131,32 @@ async def async_get_cache( # type: ignore[override] ) return None return decoded + + def set_cache(self, key, value, local_only: bool = False, **kwargs): # type: ignore[override] + model_type = cast(Optional[Type[BaseModel]], kwargs.pop("model_type", None)) + payload = CacheCodec.serialize(value, model_type=model_type) + return super().set_cache( + key=key, value=payload, local_only=local_only, **kwargs + ) + + async def async_set_cache(self, key, value, local_only: bool = False, **kwargs): # type: ignore[override] + model_type = cast(Optional[Type[BaseModel]], kwargs.pop("model_type", None)) + payload = CacheCodec.serialize(value, model_type=model_type) + return await super().async_set_cache( + key=key, value=payload, local_only=local_only, **kwargs + ) + + async def async_set_cache_pipeline( # type: ignore[override] + self, cache_list: list, local_only: bool = False, **kwargs + ) -> None: + """ + Batch writes with the same Codec boundary as ``async_set_cache`` without + ``model_type``: ``BaseModel`` values become JSON-safe dicts; dicts/scalars unchanged. + """ + normalized = [ + (key, CacheCodec.serialize(value, model_type=None)) + for key, value in cache_list + ] + return await super().async_set_cache_pipeline( + cache_list=normalized, local_only=local_only, **kwargs + ) diff --git a/litellm/proxy/management_helpers/team_member_permission_checks.py b/litellm/proxy/management_helpers/team_member_permission_checks.py index e035168ca004..50339210a6ec 100644 --- a/litellm/proxy/management_helpers/team_member_permission_checks.py +++ b/litellm/proxy/management_helpers/team_member_permission_checks.py @@ -1,6 +1,5 @@ from typing import List, Optional -from litellm.caching import DualCache from litellm.proxy._types import ( KeyManagementRoutes, LiteLLM_TeamTableCachedObj, @@ -12,6 +11,7 @@ ProxyException, UserAPIKeyAuth, ) +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.proxy.auth.auth_checks import get_team_object from litellm.proxy.auth.route_checks import RouteChecks from litellm.proxy.utils import PrismaClient @@ -65,7 +65,7 @@ async def can_team_member_execute_key_management_endpoint( user_api_key_dict: UserAPIKeyAuth, route: KeyManagementRoutes, prisma_client: PrismaClient, - user_api_key_cache: DualCache, + user_api_key_cache: UserApiKeyCache, existing_key_row: LiteLLM_VerificationToken, ): """ diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index d2dfa177515e..8c5fce84099a 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -101,6 +101,7 @@ UserAPIKeyAuth, ) from litellm.proxy.auth.route_checks import RouteChecks +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.proxy.db.create_views import ( create_missing_views, should_create_missing_views, @@ -340,7 +341,7 @@ class ProxyLogging: def __init__( self, - user_api_key_cache: DualCache, + user_api_key_cache: UserApiKeyCache, premium_user: bool = False, ): ## INITIALIZE LITELLM CALLBACKS ## @@ -5715,7 +5716,7 @@ async def get_available_models_for_user( include_model_access_groups: bool = False, only_model_access_groups: bool = False, return_wildcard_routes: bool = False, - user_api_key_cache: Optional["DualCache"] = None, + user_api_key_cache: Optional["UserApiKeyCache"] = None, ) -> List[str]: """ Get the list of models available to a user based on their API key and team permissions. From 8671ec636b9d855bcad630c29a087aa8e6769307 Mon Sep 17 00:00:00 2001 From: harish-berri Date: Fri, 1 May 2026 00:29:13 +0000 Subject: [PATCH 27/28] fix import error --- litellm/proxy/management_endpoints/ui_sso.py | 1 + .../common_utils/test_user_api_key_cache.py | 106 ++++++++++++++++++ 2 files changed, 107 insertions(+) diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index 7b4daf14a923..c1458dc37f49 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -39,6 +39,7 @@ from fastapi.responses import RedirectResponse import litellm +from litellm.caching.dual_cache import DualCache from litellm._logging import verbose_proxy_logger from litellm._uuid import uuid from litellm.constants import ( diff --git a/tests/test_litellm/proxy/common_utils/test_user_api_key_cache.py b/tests/test_litellm/proxy/common_utils/test_user_api_key_cache.py index cff331d9b192..8667348d2235 100644 --- a/tests/test_litellm/proxy/common_utils/test_user_api_key_cache.py +++ b/tests/test_litellm/proxy/common_utils/test_user_api_key_cache.py @@ -3,9 +3,23 @@ import pytest +from litellm.caching.in_memory_cache import InMemoryCache from litellm.caching.redis_cache import RedisCache from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache +from litellm.proxy.proxy_server import UserAPIKeyCacheTTLEnum + + +class CapturingInMemoryCache(InMemoryCache): + """Records ``ttl`` passed into ``set_cache`` (what DualCache injects).""" + + def __init__(self) -> None: + super().__init__() + self.last_ttl: Any = None + + def set_cache(self, key, value, **kwargs): # type: ignore[override] + self.last_ttl = kwargs.get("ttl") + super().set_cache(key, value, **kwargs) class FakeRedisCache(RedisCache): @@ -18,14 +32,18 @@ class FakeRedisCache(RedisCache): This fake: - raises TypeError if the value is not a dict - raises TypeError if the dict is not JSON-serializable + + Records the ``ttl`` kwarg DualCache forwards on each Redis write for tests. """ def __init__(self): # noqa: super().__init__ skipped intentionally self._store: dict[str, str] = {} + self.last_ttl: Any = None def set_cache(self, key: str, value: Any, **kwargs): # type: ignore[override] if not isinstance(value, dict): raise TypeError("FakeRedisCache only accepts dict payloads") + self.last_ttl = kwargs.get("ttl") self._store[key] = json.dumps(value) return True @@ -38,6 +56,7 @@ def get_cache(self, key: str, **kwargs): # type: ignore[override] async def async_set_cache(self, key: str, value: Any, **kwargs): # type: ignore[override] if not isinstance(value, dict): raise TypeError("FakeRedisCache only accepts dict payloads") + self.last_ttl = kwargs.get("ttl") self._store[key] = json.dumps(value) return True @@ -60,6 +79,93 @@ def _make_key_obj(token: str = "tok") -> UserAPIKeyAuth: class TestUserApiKeyCache: + @pytest.mark.asyncio + async def test_async_set_in_memory_gets_enum_default_when_user_api_key_cache_ttl_omitted( + self, + ): + """ + If ``general_settings.user_api_key_cache_ttl`` is absent, the proxy never + calls ``update_cache_ttl``; ``user_api_key_cache`` keeps + ``default_in_memory_ttl=UserAPIKeyCacheTTLEnum.in_memory_cache_ttl``. + DualCache must forward that as the in-memory ``ttl`` kwarg on each set. + """ + mem = CapturingInMemoryCache() + cache = UserApiKeyCache( + in_memory_cache=mem, + redis_cache=FakeRedisCache(), + default_in_memory_ttl=UserAPIKeyCacheTTLEnum.in_memory_cache_ttl.value, + ) + await cache.async_set_cache( + "k", + _make_key_obj("t"), + model_type=UserAPIKeyAuth, + ) + expected = UserAPIKeyCacheTTLEnum.in_memory_cache_ttl.value + assert mem.last_ttl == expected + + def test_sync_set_in_memory_gets_enum_default_when_user_api_key_cache_ttl_omitted( + self, + ): + mem = CapturingInMemoryCache() + cache = UserApiKeyCache( + in_memory_cache=mem, + redis_cache=FakeRedisCache(), + default_in_memory_ttl=UserAPIKeyCacheTTLEnum.in_memory_cache_ttl.value, + ) + cache.set_cache("sk", _make_key_obj("s"), model_type=UserAPIKeyAuth) + assert mem.last_ttl == UserAPIKeyCacheTTLEnum.in_memory_cache_ttl.value + + @pytest.mark.asyncio + async def test_async_set_forwards_default_in_memory_ttl_to_redis_layer(self): + """ + DualCache injects missing ``ttl`` from ``default_in_memory_ttl`` into kwargs + before calling ``redis_cache.async_set_cache`` — Redis should receive the same + TTL as memory (matches proxy defaults: enum 60s). + """ + fake = FakeRedisCache() + cache = UserApiKeyCache( + redis_cache=fake, + default_in_memory_ttl=60, + ) + + await cache.async_set_cache( + key="ttl-key", + value=_make_key_obj("ttl-tok"), + model_type=UserAPIKeyAuth, + ) + + assert fake.last_ttl == 60 + + @pytest.mark.asyncio + async def test_async_set_explicit_ttl_override_reaches_redis(self): + fake = FakeRedisCache() + cache = UserApiKeyCache( + redis_cache=fake, + default_in_memory_ttl=60, + ) + + await cache.async_set_cache( + key="k", + value=_make_key_obj("x"), + model_type=UserAPIKeyAuth, + ttl=900, + ) + + assert fake.last_ttl == 900 + + def test_sync_set_forwards_default_in_memory_ttl_to_redis_layer(self): + fake = FakeRedisCache() + cache = UserApiKeyCache( + redis_cache=fake, + default_in_memory_ttl=45, + ) + cache.set_cache( + "sk", + _make_key_obj("sync"), + model_type=UserAPIKeyAuth, + ) + assert fake.last_ttl == 45 + @pytest.mark.asyncio async def test_async_set_typed_stores_serialized_payload_in_memory_and_redis(self): cache = UserApiKeyCache(redis_cache=FakeRedisCache()) From 896e1024643e20f609e618df4141b3368aca3a0d Mon Sep 17 00:00:00 2001 From: harish-berri Date: Fri, 1 May 2026 00:32:10 +0000 Subject: [PATCH 28/28] fix linting --- litellm/proxy/management_endpoints/ui_sso.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index c1458dc37f49..9dfc67370fea 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -76,8 +76,10 @@ ) from litellm.proxy.auth.auth_checks import ExperimentalUIJWTToken, get_user_object from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache -from litellm.proxy.auth.auth_utils import _has_user_setup_sso -from litellm.proxy.auth.auth_utils import _get_request_ip_address, _has_user_setup_sso +from litellm.proxy.auth.auth_utils import ( + _get_request_ip_address, + _has_user_setup_sso, +) from litellm.proxy.auth.handle_jwt import JWTHandler from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_utils.admin_ui_utils import (