diff --git a/litellm/caching/caching.py b/litellm/caching/caching.py index b58d1be4971..eadf8987fca 100644 --- a/litellm/caching/caching.py +++ b/litellm/caching/caching.py @@ -410,7 +410,7 @@ class Cache: ) return "".join(f"{field}: {value}" for field, value in scope_values if value is not None) - def get_cache_key(self, **kwargs: object) -> str: + def get_cache_key(self, **kwargs: object) -> str | None: """ Get the cache key for the given arguments. @@ -421,11 +421,18 @@ class Cache: str: The cache key generated from the arguments, or None if no cache key could be generated. """ cache_key = "" + authenticated_caller, authenticated_namespace = self._get_authenticated_cache_namespace(**kwargs) + if authenticated_caller and authenticated_namespace is None: + # An authenticated request without a stable server-side identity must + # never fall back to a shared or client-selected cache key. + return None # verbose_logger.debug("\nGetting Cache key. Kwargs: %s", kwargs) preset_cache_key: Final = self.get_preset_cache_key_from_kwargs(**kwargs) if preset_cache_key is not None: verbose_logger.debug("\nReturning preset cache key: %s", preset_cache_key) + if authenticated_caller: + return self._add_namespace_to_cache_key(preset_cache_key, **kwargs) return preset_cache_key combined_kwargs: Final = ModelParamHelper.get_all_llm_api_params() @@ -459,6 +466,8 @@ class Cache: ) hashed_cache_key = Cache._get_hashed_cache_key(cache_key + reasoning_summary_material) hashed_cache_key = self._add_namespace_to_cache_key(hashed_cache_key, **kwargs) + if hashed_cache_key is None: + return None verbose_logger.debug( "\nCreated cache key: %s (source material length: %d)", hashed_cache_key, @@ -587,7 +596,7 @@ class Cache: self, hash_hex: str, **kwargs: object, # kwargs-ok: cache request metadata is extensible - ) -> str: + ) -> str | None: """ If a redis namespace is provided, add it to the cache key @@ -622,29 +631,9 @@ class Cache: ) if isinstance(source, Mapping) ) - authenticated_namespace: str | None = None # rebind-ok: populated once while scanning trusted auth metadata - try: - from litellm.proxy._types import UserAPIKeyAuth - except ImportError: - - class UserAPIKeyAuth: # pragma: no cover - proxy types are available in proxy deployments - pass - - for metadata in metadata_sources: - auth_object: object | None = metadata.get("user_api_key_auth") - if isinstance(auth_object, UserAPIKeyAuth): - identity_fields = { - "api_key_hash": getattr(auth_object, "token", None), - "team_id": getattr(auth_object, "team_id", None), - "project_id": getattr(auth_object, "project_id", None), - "org_id": getattr(auth_object, "org_id", None), - "user_id": getattr(auth_object, "user_id", None), - "end_user_id": getattr(auth_object, "end_user_id", None), - "user_role": getattr(auth_object, "user_role", None), - } - identity = json.dumps(identity_fields, sort_keys=True, default=str, separators=(",", ":")) - authenticated_namespace = "caller:" + hashlib.sha256(identity.encode("utf-8")).hexdigest()[:32] - break + authenticated_caller, authenticated_namespace = self._get_authenticated_cache_namespace(**kwargs) + if authenticated_caller and authenticated_namespace is None: + return None requested_namespace_value: Final = dynamic_cache_control.get("namespace") requested_namespace: Final[str | None] = ( requested_namespace_value if isinstance(requested_namespace_value, str) else None @@ -662,10 +651,71 @@ class Cache: else requested_namespace or metadata_namespace or self.namespace ) if namespace: - hash_hex = f"{namespace}:{hash_hex}" # rebind-ok: namespace prefix is the final cache key representation + prefix = f"{namespace}:" + if not hash_hex.startswith(prefix): + hash_hex = f"{prefix}{hash_hex}" # rebind-ok: namespace prefix is the final cache key representation verbose_logger.debug("Final hashed key: %s", hash_hex) return hash_hex + def _get_authenticated_cache_namespace(self, **kwargs: object) -> tuple[bool, str | None]: + """Return whether trusted auth metadata exists and its stable cache namespace. + + ``UserAPIKeyAuth.token`` is the proxy's server-side hashed token identity. + API-key-only objects without it cannot be safely separated, so callers + must bypass cache reads and writes rather than use a shared fallback. + """ + litellm_params_value: Final = kwargs.get("litellm_params") + litellm_params: Final[Mapping[str, object]] = ( + cast(Mapping[str, object], litellm_params_value) if isinstance(litellm_params_value, Mapping) else {} + ) + metadata_sources: Final[tuple[Mapping[str, object], ...]] = tuple( + cast(Mapping[str, object], source) + for source in ( + kwargs.get("metadata"), + kwargs.get("litellm_metadata"), + litellm_params.get("metadata"), + litellm_params.get("litellm_metadata"), + ) + if isinstance(source, Mapping) + ) + try: + from litellm.proxy._types import UserAPIKeyAuth + except ImportError: + + class UserAPIKeyAuth: # pragma: no cover - proxy types are available in proxy deployments + pass + + for metadata in metadata_sources: + auth_object: object | None = metadata.get("user_api_key_auth") + if not isinstance(auth_object, UserAPIKeyAuth): + continue + token = getattr(auth_object, "token", None) + if not isinstance(token, str) or not token: + return True, None + identity_fields = { + "api_key_hash": token, + "team_id": getattr(auth_object, "team_id", None), + "project_id": getattr(auth_object, "project_id", None), + "org_id": getattr(auth_object, "org_id", None), + "user_id": getattr(auth_object, "user_id", None), + "end_user_id": getattr(auth_object, "end_user_id", None), + "user_role": getattr(auth_object, "user_role", None), + } + identity = json.dumps(identity_fields, sort_keys=True, default=str, separators=(",", ":")) + return True, "caller:" + hashlib.sha256(identity.encode("utf-8")).hexdigest()[:32] + return False, None + + def _prepare_explicit_cache_key(self, explicit_key: object, **kwargs: object) -> str | None: + """Validate and scope a caller-provided cache key before backend access.""" + if not isinstance(explicit_key, str): + return None + authenticated_caller, authenticated_namespace = self._get_authenticated_cache_namespace(**kwargs) + if not authenticated_caller: + return explicit_key + if authenticated_namespace is None: + return None + return self._add_namespace_to_cache_key(explicit_key, **kwargs) + def generate_streaming_content(self, content): chunk_size: Final = 5 # Adjust the chunk size as needed for i in range(0, len(content), chunk_size): @@ -759,7 +809,7 @@ class Cache: return with response_cache_phase("get"): if "cache_key" in kwargs: - cache_key = kwargs["cache_key"] + cache_key = self._prepare_explicit_cache_key(kwargs["cache_key"], **kwargs) else: cache_key = self.get_cache_key(**kwargs) if cache_key is not None: @@ -792,7 +842,7 @@ class Cache: with response_cache_phase("get"): if "cache_key" in kwargs: - cache_key = kwargs["cache_key"] + cache_key = self._prepare_explicit_cache_key(kwargs["cache_key"], **kwargs) else: cache_key = self.get_cache_key(**kwargs) if cache_key is not None: @@ -813,31 +863,27 @@ class Cache: """ Common implementation across sync + async add_cache functions """ - try: - if "cache_key" in kwargs: - cache_key = kwargs["cache_key"] - else: - cache_key = self.get_cache_key(**kwargs) - if cache_key is not None: - if isinstance(result, BaseModel): - result = result.model_dump_json(exclude_unset=True) + if "cache_key" in kwargs: + cache_key = self._prepare_explicit_cache_key(kwargs["cache_key"], **kwargs) + else: + cache_key = self.get_cache_key(**kwargs) + if cache_key is not None: + if isinstance(result, BaseModel): + result = result.model_dump_json(exclude_unset=True) - ## DEFAULT TTL ## - if self.ttl is not None: - kwargs["ttl"] = self.ttl - ## Get Cache-Controls ## - _cache_kwargs: Final = kwargs.get("cache", None) - if isinstance(_cache_kwargs, dict): - for k, v in _cache_kwargs.items(): - if k == "ttl": - kwargs["ttl"] = v + ## DEFAULT TTL ## + if self.ttl is not None: + kwargs["ttl"] = self.ttl + ## Get Cache-Controls ## + _cache_kwargs: Final = kwargs.get("cache", None) + if isinstance(_cache_kwargs, dict): + for k, v in _cache_kwargs.items(): + if k == "ttl": + kwargs["ttl"] = v - cached_data: Final = {"timestamp": time.time(), "response": result} - return cache_key, cached_data, kwargs - else: - raise Exception("cache key is None") - except Exception as e: - raise e + cached_data: Final = {"timestamp": time.time(), "response": result} + return cache_key, cached_data, kwargs + return None def add_cache(self, result, **kwargs): """ @@ -854,7 +900,10 @@ class Cache: if self.should_use_cache(**kwargs) is not True: return with response_cache_phase("set"): - cache_key, cached_data, kwargs = self._add_cache_logic(result=result, **kwargs) + cache_entry = self._add_cache_logic(result=result, **kwargs) + if cache_entry is None: + return + cache_key, cached_data, kwargs = cache_entry self.cache.set_cache(cache_key, cached_data, **kwargs) except Exception as e: self._log_add_cache_failure(e) @@ -878,7 +927,10 @@ class Cache: # high traffic - fill in results in memory and then flush await self.batch_cache_write(result, **kwargs) else: - cache_key, cached_data, kwargs = self._add_cache_logic(result=result, **kwargs) + cache_entry = self._add_cache_logic(result=result, **kwargs) + if cache_entry is None: + return + cache_key, cached_data, kwargs = cache_entry if dynamic_cache_object is not None: await dynamic_cache_object.async_set_cache(cache_key, cached_data, **kwargs) else: @@ -989,7 +1041,7 @@ class Cache: input: str, kwargs: dict, idx_in_result_data: int = 0, - ) -> tuple[str, dict, dict]: + ) -> tuple[str, dict, dict] | None: preset_cache_key: Final = self.get_cache_key(**{**kwargs, "input": input}) kwargs["cache_key"] = preset_cache_key embedding_response: Final = result.data[idx_in_result_data] @@ -1013,11 +1065,10 @@ class Cache: prompt_tokens_details=prompt_tokens_details, ) - cache_key, cached_data, kwargs = self._add_cache_logic( + return self._add_cache_logic( result=embedding_dict, **kwargs, ) - return cache_key, cached_data, kwargs async def async_add_cache_pipeline(self, result, dynamic_cache_object: BaseCache | None = None, **kwargs): """ @@ -1046,16 +1097,16 @@ class Cache: cache_list: Final = [] if isinstance(kwargs["input"], list): for idx, i in enumerate(kwargs["input"]): - ( - cache_key, - cached_data, - kwargs, - ) = self.add_embedding_response_to_cache(result, i, kwargs, idx) + cache_entry = self.add_embedding_response_to_cache(result, i, kwargs, idx) + if cache_entry is None: + return + cache_key, cached_data, kwargs = cache_entry cache_list.append((cache_key, cached_data)) elif isinstance(kwargs["input"], str): - cache_key, cached_data, kwargs = self.add_embedding_response_to_cache( - result, kwargs["input"], kwargs - ) + cache_entry = self.add_embedding_response_to_cache(result, kwargs["input"], kwargs) + if cache_entry is None: + return + cache_key, cached_data, kwargs = cache_entry cache_list.append((cache_key, cached_data)) if dynamic_cache_object is not None: @@ -1087,7 +1138,10 @@ class Cache: return False async def batch_cache_write(self, result, **kwargs): - cache_key, cached_data, kwargs = self._add_cache_logic(result=result, **kwargs) + cache_entry = self._add_cache_logic(result=result, **kwargs) + if cache_entry is None: + return + cache_key, cached_data, kwargs = cache_entry await self.cache.batch_cache_write(cache_key, cached_data, **kwargs) async def ping(self): diff --git a/litellm/proxy/hooks/batch_redis_get.py b/litellm/proxy/hooks/batch_redis_get.py index 94eb0bb52d1..076825e1d6b 100644 --- a/litellm/proxy/hooks/batch_redis_get.py +++ b/litellm/proxy/hooks/batch_redis_get.py @@ -93,10 +93,12 @@ class _PROXY_BatchRedisRequests(CustomLogger): - return redis cache request """ try: # never block execution + if litellm.cache is None: + return None cache_key: str | None = None if "cache_key" in kwargs: - cache_key = kwargs["cache_key"] - elif litellm.cache is not None: + cache_key = litellm.cache._prepare_explicit_cache_key(kwargs["cache_key"], **kwargs) + else: cache_key = litellm.cache.get_cache_key( *args, **kwargs ) # returns ":" - we pass redis_namespace in async_pre_call_hook. Done to avoid rewriting the async_set_cache logic diff --git a/tests/unit/caching/test_batch_redis_get_cache.py b/tests/unit/caching/test_batch_redis_get_cache.py new file mode 100644 index 00000000000..226c66ade20 --- /dev/null +++ b/tests/unit/caching/test_batch_redis_get_cache.py @@ -0,0 +1,52 @@ +import time +from typing import Final + +import pytest + +import litellm +from litellm.caching.caching import Cache +from litellm.caching.in_memory_cache import InMemoryCache +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.hooks.batch_redis_get import PROXY_BatchRedisRequests +from litellm.types.caching import LiteLLMCacheType + + +class _BackendSpy: + def __init__(self) -> None: + self.keys: list[str] = [] + + async def async_get_cache(self, key: str, *args, **kwargs): + self.keys.append(key) + return {"timestamp": time.time(), "response": {"key": key}} + + +def _caller(token: str | None) -> UserAPIKeyAuth: + return UserAPIKeyAuth(api_key=f"raw-{token}").model_copy(update={"token": token}) + + +@pytest.mark.asyncio +async def test_batch_redis_hook_scopes_explicit_keys_and_fails_closed(monkeypatch: pytest.MonkeyPatch) -> None: + cache: Final = Cache(type=LiteLLMCacheType.LOCAL) + backend = _BackendSpy() + cache.cache = backend + monkeypatch.setattr(litellm, "cache", cache) + hook = PROXY_BatchRedisRequests() + hook.in_memory_cache = InMemoryCache() + + request_a = {"metadata": {"user_api_key_auth": _caller("hashed-a")}, "cache_key": "shared"} + request_b = {"metadata": {"user_api_key_auth": _caller("hashed-b")}, "cache_key": "shared"} + + first_a = await hook.async_get_cache(**request_a) + second_a = await hook.async_get_cache(**request_a) + first_b = await hook.async_get_cache(**request_b) + missing = await hook.async_get_cache( + metadata={"user_api_key_auth": _caller(None)}, + cache_key="shared", + ) + + assert first_a == second_a + assert first_a != first_b + assert missing is None + assert len(backend.keys) == 2 + assert backend.keys[0] != backend.keys[1] + assert all(key.startswith("caller:") for key in backend.keys) diff --git a/tests/unit/caching/test_caching.py b/tests/unit/caching/test_caching.py index 43590ac7af7..00dd73f09a9 100644 --- a/tests/unit/caching/test_caching.py +++ b/tests/unit/caching/test_caching.py @@ -3,6 +3,7 @@ import hashlib import json import logging import re +import time import traceback import uuid from datetime import timedelta @@ -53,6 +54,137 @@ def test_authenticated_proxy_callers_get_distinct_exact_cache_keys() -> None: assert key_a != key_b +def test_authenticated_callers_without_server_token_fail_closed_for_all_key_forms() -> None: + cache: Final = Cache(type=LiteLLMCacheType.LOCAL, namespace="qa") + caller_a: Final = UserAPIKeyAuth(api_key="raw-a", team_id="team-a", user_id="user-a").model_copy( + update={"token": None} + ) + caller_b: Final = UserAPIKeyAuth(api_key="raw-b", team_id="team-a", user_id="user-a").model_copy( + update={"token": None} + ) + + for caller in (caller_a, caller_b): + metadata = {"user_api_key_auth": caller, "redis_namespace": "client-selected"} + assert cache.get_cache_key(model="gpt-4.1-mini", messages=messages, metadata=metadata) is None + assert ( + cache.get_cache_key( + model="gpt-4.1-mini", + messages=messages, + metadata=metadata, + litellm_params={"preset_cache_key": "shared"}, + ) + is None + ) + + +def test_explicit_cache_keys_are_idempotently_scoped_to_valid_callers() -> None: + cache: Final = Cache(type=LiteLLMCacheType.LOCAL, namespace="qa") + caller_a: Final = UserAPIKeyAuth(api_key="raw-a", token="hashed-a", team_id="team-a") + caller_b: Final = UserAPIKeyAuth(api_key="raw-b", token="hashed-b", team_id="team-a") + metadata_a = {"user_api_key_auth": caller_a, "redis_namespace": "client-selected"} + metadata_b = {"user_api_key_auth": caller_b, "redis_namespace": "client-selected"} + + key_a = cache.get_cache_key(model="gpt-4.1-mini", messages=messages, metadata=metadata_a) + assert key_a is not None + explicit_a = cache._prepare_explicit_cache_key("shared", metadata=metadata_a) + assert explicit_a is not None + assert explicit_a != cache._prepare_explicit_cache_key("shared", metadata=metadata_b) + assert cache._prepare_explicit_cache_key(explicit_a, metadata=metadata_a) == explicit_a + + +class _CacheBackendSpy: + def __init__(self) -> None: + self.get_calls = 0 + self.set_calls = 0 + + def get_cache(self, *args, **kwargs): + self.get_calls += 1 + return {"timestamp": time.time(), "response": {"value": "cached"}} + + def set_cache(self, *args, **kwargs): + self.set_calls += 1 + + async def async_get_cache(self, *args, **kwargs): + self.get_calls += 1 + return {"timestamp": time.time(), "response": {"value": "cached"}} + + async def async_set_cache(self, *args, **kwargs): + self.set_calls += 1 + + +def _missing_token_cache_request() -> dict[str, object]: + return { + "model": "gpt-4.1-mini", + "messages": messages, + "metadata": { + "user_api_key_auth": UserAPIKeyAuth(api_key="raw-a", team_id="team-a").model_copy( + update={"token": None} + ), + }, + } + + +def test_missing_token_skips_sync_backend_access_for_explicit_key(caplog) -> None: + cache: Final = Cache(type=LiteLLMCacheType.LOCAL) + backend = _CacheBackendSpy() + cache.cache = backend + request = _missing_token_cache_request() + + assert cache.get_cache(**request, cache_key="shared") is None + cache.add_cache("response", **request, cache_key="shared") + + assert backend.get_calls == 0 + assert backend.set_calls == 0 + assert "exception in add_cache" not in caplog.text + + +@pytest.mark.asyncio +async def test_missing_token_skips_async_backend_access_for_explicit_key() -> None: + cache: Final = Cache(type=LiteLLMCacheType.LOCAL) + backend = _CacheBackendSpy() + cache.cache = backend + request = _missing_token_cache_request() + + assert await cache.async_get_cache(**request, cache_key="shared") is None + await cache.async_add_cache("response", **request, cache_key="shared") + + assert backend.get_calls == 0 + assert backend.set_calls == 0 + + +def test_valid_token_explicit_key_reaches_sync_backend() -> None: + cache: Final = Cache(type=LiteLLMCacheType.LOCAL) + backend = _CacheBackendSpy() + cache.cache = backend + request = { + "model": "gpt-4.1-mini", + "messages": messages, + "metadata": {"user_api_key_auth": UserAPIKeyAuth(api_key="raw-a", token="hashed-a")}, + } + + cache.add_cache("response", **request, cache_key="shared") + assert cache.get_cache(**request, cache_key="shared") == {"value": "cached"} + assert backend.set_calls == 1 + assert backend.get_calls == 1 + + +@pytest.mark.asyncio +async def test_valid_token_explicit_key_reaches_async_backend() -> None: + cache: Final = Cache(type=LiteLLMCacheType.LOCAL) + backend = _CacheBackendSpy() + cache.cache = backend + request = { + "model": "gpt-4.1-mini", + "messages": messages, + "metadata": {"user_api_key_auth": UserAPIKeyAuth(api_key="raw-a", token="hashed-a")}, + } + + await cache.async_add_cache("response", **request, cache_key="shared") + assert await cache.async_get_cache(**request, cache_key="shared") == {"value": "cached"} + assert backend.set_calls == 1 + assert backend.get_calls == 1 + + @pytest.fixture def preserve_litellm_set_verbose(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr(litellm, "set_verbose", litellm.set_verbose)