From 63d39b41fd4cfa3b9aa0c47f3ea04756e4054e98 Mon Sep 17 00:00:00 2001 From: quietjules <349048734+quietjules@users.noreply.github.com> Date: Tue, 10 Nov 2026 20:21:39 -0800 Subject: [PATCH] fix(proxy): keep a request's cache opt-in when its key has cache metadata (#50999) * fix(proxy): merge a key's cache controls into the request's cache settings instead of replacing them * fix(proxy): keep a key with no-cache out of opt-in caching A key whose cache metadata sets no-cache now drops the caller's use-cache, so on a default_off proxy its responses stay out of the shared cache as before the merge change. default_on proxies ignore use-cache, so they are unchanged * fix(proxy): keep base handling of non-dict key cache metadata Key cache metadata that is not a dict, such as false or a string, wipes the request's cache settings as before the merge change, so those keys stay out of opt-in caching * fix(proxy): keep a key with use-cache false out of opt-in caching * fix(proxy): only keep request cache settings on an allowed opt-in --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- litellm/models/verification_token.py | 2 +- litellm/proxy/litellm_pre_call_utils.py | 41 ++++- .../unit/proxy/test_litellm_pre_call_utils.py | 155 ++++++++++++++++++ 3 files changed, 190 insertions(+), 8 deletions(-) diff --git a/litellm/models/verification_token.py b/litellm/models/verification_token.py index 89048c56a9f..4b804ebde1c 100644 --- a/litellm/models/verification_token.py +++ b/litellm/models/verification_token.py @@ -35,7 +35,7 @@ class LiteLLM_VerificationToken(LiteLLMPydanticObjectBase): tpd_limit: int | None = None budget_duration: str | None = None budget_reset_at: datetime | None = None - allowed_cache_controls: list | None = [] + allowed_cache_controls: list[object] | None = [] allowed_routes: list | None = [] key_type: str | None = None permissions: dict = {} diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 447b41f9cfd..0b2ba93af71 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -2217,15 +2217,41 @@ class LiteLLMProxyRequestSetup: return data @staticmethod - def add_key_level_controls(key_metadata: dict | None, data: dict, _metadata_variable_name: str): + def _request_cache_kept_for_key( + request_cache: Mapping[str, object], + key_cache: Mapping[str, object], + allowed_cache_controls: Sequence[object] | None, + ) -> Mapping[str, object]: + key_opts_out: Final = ( + key_cache.get("no-cache", False) is not False or key_cache.get("use-cache", True) is not True + ) + request_is_allowed: Final = not allowed_cache_controls or all( + control in allowed_cache_controls for control in request_cache + ) + if key_opts_out or request_cache.get("use-cache") is not True or not request_is_allowed: + return MappingProxyType({}) + return {k: v for k, v in request_cache.items() if not (k == "s-max-age" and "s-maxage" in key_cache)} + + @staticmethod + def add_key_level_controls( + key_metadata: dict | None, + data: dict, + _metadata_variable_name: str, + allowed_cache_controls: Sequence[object] | None = None, + ): if key_metadata is None: return data - if "cache" in key_metadata: - data["cache"] = {} - if isinstance(key_metadata["cache"], dict): - for k, v in key_metadata["cache"].items(): - if k in SupportedCacheControls: - data["cache"][k] = v + if "cache" in key_metadata and not isinstance(key_metadata["cache"], dict): + data["cache"] = {} # rebind-ok: data is an out-param + elif "cache" in key_metadata: + request_cache: Final = _DECODED_JSON.validate_python(data["cache"]) if "cache" in data else None + key_cache: Final = _request_fields(_DECODED_JSON.validate_python(key_metadata["cache"])) + data["cache"] = { # rebind-ok: data is an out-param + **LiteLLMProxyRequestSetup._request_cache_kept_for_key( + _request_fields(request_cache), key_cache, allowed_cache_controls + ), + **{k: v for k, v in key_cache.items() if k in SupportedCacheControls}, + } ## KEY-LEVEL SPEND LOGS / TAGS if "tags" in key_metadata and key_metadata["tags"] is not None: @@ -2967,6 +2993,7 @@ async def add_litellm_data_to_request( key_metadata=key_metadata, data=data, _metadata_variable_name=_metadata_variable_name, + allowed_cache_controls=user_api_key_dict.allowed_cache_controls, ) data[_metadata_variable_name] = LiteLLMProxyRequestSetup.add_team_and_project_level_controls( user_api_key_dict=user_api_key_dict, diff --git a/tests/unit/proxy/test_litellm_pre_call_utils.py b/tests/unit/proxy/test_litellm_pre_call_utils.py index f4a836b4b4b..f5757a0ff18 100644 --- a/tests/unit/proxy/test_litellm_pre_call_utils.py +++ b/tests/unit/proxy/test_litellm_pre_call_utils.py @@ -17,6 +17,7 @@ from pydantic import ValidationError as PydanticValidationError from starlette.datastructures import Headers import litellm +from litellm.caching.caching import Cache, CacheMode from litellm.constants import ( CALLER_SESSION_ID_METADATA_KEY, COMPLEXITY_ROUTER_CHILD_SESSION_ID_METADATA_KEY, @@ -64,6 +65,7 @@ from litellm.proxy.litellm_pre_call_utils import ( resolve_tenant_otel_destinations, ) from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload +from litellm.types.caching import LiteLLMCacheType from litellm.types.litellm_params import AGENTIC_LOOP_KWARG_NAMES from litellm.types.utils import CredentialItem from tests.unit.proxy.router_entry_point_stubs import route_entry_points_to_stubs @@ -8379,6 +8381,159 @@ async def test_key_level_callback_vars_survive_the_strip(): assert updated["dd_site"] == "us5.datadoghq.com" +@pytest.mark.asyncio +@pytest.mark.parametrize( + "key_cache, request_cache, expected_cache, uses_cache_when_default_off", + [ + pytest.param( + {"ttl": 600}, + {"use-cache": True, "namespace": "team-a"}, + {"use-cache": True, "namespace": "team-a", "ttl": 600}, + True, + id="key-control-added-to-request-opt-in", + ), + pytest.param( + {"no-store": True}, + {"use-cache": True, "no-store": False}, + {"use-cache": True, "no-store": True}, + True, + id="key-control-wins-over-request", + ), + pytest.param( + {"use-cache": True, "namespace": "key-ns"}, + None, + {}, + False, + id="unsupported-key-controls-ignored", + ), + pytest.param(None, {"use-cache": True}, {}, False, id="non-dict-key-cache-still-wipes-request"), + pytest.param( + {"no-cache": True}, + {"use-cache": True, "ttl": 60}, + {"no-cache": True}, + False, + id="key-no-cache-keeps-the-key-out-of-opt-in-caching", + ), + pytest.param( + {"use-cache": False}, + {"use-cache": True, "ttl": 60}, + {}, + False, + id="key-use-cache-false-keeps-the-key-out-of-opt-in-caching", + ), + pytest.param( + {"use-cache": True}, + {"use-cache": True}, + {"use-cache": True}, + True, + id="key-use-cache-true-leaves-the-request-opt-in", + ), + pytest.param( + {"no-cache": "true"}, + {"use-cache": True}, + {"no-cache": "true"}, + False, + id="key-non-boolean-no-cache-keeps-the-key-out-of-opt-in-caching", + ), + pytest.param( + {"no-cache": False}, + {"use-cache": True}, + {"use-cache": True, "no-cache": False}, + True, + id="key-no-cache-false-leaves-the-request-opt-in", + ), + pytest.param( + {"ttl": 600}, + {"no-cache": True, "namespace": "team-a"}, + {"ttl": 600}, + False, + id="request-without-opt-in-is-still-replaced", + ), + pytest.param( + {"ttl": 600}, + {"use-cache": "true", "no-store": True}, + {"ttl": 600}, + False, + id="request-non-boolean-use-cache-is-still-replaced", + ), + pytest.param( + {"s-maxage": 1}, + {"use-cache": True, "s-max-age": 100000}, + {"use-cache": True, "s-maxage": 1}, + True, + id="key-s-maxage-overrides-the-request-s-max-age-alias", + ), + pytest.param( + {"ttl": 600}, + {"use-cache": True, "s-max-age": 5}, + {"use-cache": True, "s-max-age": 5, "ttl": 600}, + True, + id="request-s-max-age-kept-when-the-key-has-no-s-maxage", + ), + ], +) +async def test_key_cache_controls_merge_into_request_cache_controls( + key_cache: dict[str, object] | None, + request_cache: dict[str, object] | None, + expected_cache: dict[str, object], + uses_cache_when_default_off: bool, +): + key: Final = UserAPIKeyAuth(api_key="hashed-key", metadata={"cache": key_cache}) + data: Final = { + "model": "gpt-4.1-mini", + "messages": [{"role": "user", "content": "hello"}], + **({} if request_cache is None else {"cache": request_cache}), + } + + updated: Final = await add_litellm_data_to_request( + data=data, + request=_callback_credential_request_mock(), + user_api_key_dict=key, + proxy_config=MagicMock(), + general_settings={}, + version="test-version", + ) + + assert updated["cache"] == expected_cache + opt_in_cache: Final = Cache(type=LiteLLMCacheType.LOCAL, mode=CacheMode.default_off) + assert opt_in_cache.should_use_cache(**updated) is uses_cache_when_default_off + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "allowed_cache_controls, request_cache, expected_cache", + [ + pytest.param(["no-store"], {"use-cache": True}, {"ttl": 600}, id="opt-in-not-allowed-is-replaced"), + pytest.param( + ["use-cache"], {"use-cache": True, "no-store": True}, {"ttl": 600}, id="any-disallowed-control-is-replaced" + ), + pytest.param(["use-cache"], {"use-cache": True}, {"use-cache": True, "ttl": 600}, id="allowed-opt-in-is-kept"), + pytest.param([], {"use-cache": True}, {"use-cache": True, "ttl": 600}, id="empty-list-allows-the-opt-in"), + ], +) +async def test_key_cache_metadata_only_keeps_request_controls_the_key_allows( + allowed_cache_controls: list[str], + request_cache: dict[str, object], + expected_cache: dict[str, object], +): + key: Final = UserAPIKeyAuth( + api_key="hashed-key", + metadata={"cache": {"ttl": 600}}, + allowed_cache_controls=allowed_cache_controls, + ) + + updated: Final = await add_litellm_data_to_request( + data={"model": "gpt-4.1-mini", "messages": [{"role": "user", "content": "hello"}], "cache": request_cache}, + request=_callback_credential_request_mock(), + user_api_key_dict=key, + proxy_config=MagicMock(), + general_settings={}, + version="test-version", + ) + + assert updated["cache"] == expected_cache + + class TestPromotedTraceControlFields: """LIT-5137: caller metadata trace fields must reach litellm_metadata."""