From d0a9e2181c1063aecb60a80889522804c6602a48 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 21 Oct 2026 04:08:31 -0700 Subject: [PATCH] fix(caching): forward bridged cache controls only to a semantic cache, and only a caller's opt-outs on A2A The generateContent and A2A completion bridges forwarded the request's cache controls for every cache type, so an exact-match cache started honoring them. Both now forward only when the proxy's cache is semantic. The A2A bridge merged every request control over the agent's own cache setting, so a caller could name another agent's namespace and read its entries. It now forwards only a no-cache or no-store the request sets. The generateContent cache test waits on the pending cache writes instead of sleeping, and the pre-call gate test runs the real request helpers. --- litellm/google_genai/adapters/handler.py | 2 +- .../proxy/agent_endpoints/a2a_endpoints.py | 15 +++-- .../google_genai/test_google_genai_adapter.py | 55 ++++++++++++++++--- .../agent_endpoints/test_a2a_endpoints.py | 40 +++++++++++++- .../proxy/test_common_request_processing.py | 36 +++++------- 5 files changed, 113 insertions(+), 35 deletions(-) diff --git a/litellm/google_genai/adapters/handler.py b/litellm/google_genai/adapters/handler.py index e386008f5d3..f0c8fbab106 100644 --- a/litellm/google_genai/adapters/handler.py +++ b/litellm/google_genai/adapters/handler.py @@ -69,7 +69,7 @@ class GenerateContentToCompletionHandler: completion_kwargs["proxy_server_request"] = extra_kwargs["proxy_server_request"] if "litellm_metadata" in extra_kwargs: completion_kwargs["litellm_metadata"] = extra_kwargs["litellm_metadata"] - if "cache" in extra_kwargs: + if "cache" in extra_kwargs and litellm.cache is not None and litellm.cache.is_semantic_cache(): completion_kwargs["cache"] = extra_kwargs["cache"] if stream: diff --git a/litellm/proxy/agent_endpoints/a2a_endpoints.py b/litellm/proxy/agent_endpoints/a2a_endpoints.py index 6a0f4462160..20cc9844920 100644 --- a/litellm/proxy/agent_endpoints/a2a_endpoints.py +++ b/litellm/proxy/agent_endpoints/a2a_endpoints.py @@ -23,6 +23,7 @@ from pydantic import TypeAdapter, ValidationError import litellm from litellm._logging import verbose_proxy_logger +from litellm.caching.caching import Cache from litellm.litellm_core_utils.url_utils import SSRFError, validate_url from litellm.llms.a2a.common_utils import resolve_a2a_hop_auth_header from litellm.proxy._types import UserAPIKeyAuth @@ -40,6 +41,7 @@ from litellm.proxy.agent_endpoints.databricks_oauth import ( ) from litellm.proxy.agent_endpoints.utils import merge_agent_headers from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.common_utils.semantic_cache_gate import SKIP_CACHE_CONTROLS from litellm.proxy.common_utils.sse_keepalive import ( SSE_COMMENT_PING, coerce_keepalive_interval, @@ -180,12 +182,17 @@ async def _resolve_backend_auth( ) -def _request_cache_controls(agent_params: Mapping[str, object], data: Mapping[str, object]) -> Mapping[str, object]: +def _request_cache_controls( + agent_params: Mapping[str, object], data: Mapping[str, object], cache: Cache | None +) -> Mapping[str, object]: request_cache: Final = data.get("cache") - if not isinstance(request_cache, Mapping): + if cache is None or not cache.is_semantic_cache() or not isinstance(request_cache, Mapping): + return {} + bypass: Final = {key: value for key, value in SKIP_CACHE_CONTROLS.items() if request_cache.get(key) is value} + if not bypass: return {} agent_cache: Final = agent_params.get("cache") - return {"cache": {**(agent_cache if isinstance(agent_cache, Mapping) else {}), **request_cache}} + return {"cache": {**(agent_cache if isinstance(agent_cache, Mapping) else {}), **bypass}} def _forwarding_headers( @@ -848,7 +855,7 @@ async def invoke_agent_a2a( ) forwarded_litellm_params: Final = { **backend_auth.litellm_params, - **_request_cache_controls(backend_auth.litellm_params, data), + **_request_cache_controls(backend_auth.litellm_params, data, litellm.cache), } # Merge agent-level guardrails into data so post_call_success_hook and diff --git a/tests/unit/google_genai/test_google_genai_adapter.py b/tests/unit/google_genai/test_google_genai_adapter.py index dbcc80a5d25..4add065b36f 100644 --- a/tests/unit/google_genai/test_google_genai_adapter.py +++ b/tests/unit/google_genai/test_google_genai_adapter.py @@ -5,8 +5,11 @@ Test to verify the Google GenAI generate_content adapter functionality import asyncio import json import unittest +from collections.abc import Mapping +from typing import Final import pytest +import respx from litellm.google_genai.main import agenerate_content @@ -1645,15 +1648,26 @@ async def test_generate_content_sends_response_schema_and_tool_parameters_to_the @pytest.mark.asyncio -async def test_generate_content_bridge_honors_the_request_cache_controls(respx_mock, monkeypatch): +@pytest.mark.parametrize( + ("first_cache_controls", "second_cache_controls"), + [({"no-store": True}, None), (None, {"no-cache": True})], + ids=["no-store", "no-cache"], +) +async def test_generate_content_bridge_leaves_an_exact_match_cache_as_it_was( + respx_mock: respx.MockRouter, + monkeypatch: pytest.MonkeyPatch, + first_cache_controls: Mapping[str, bool] | None, + second_cache_controls: Mapping[str, bool] | None, +) -> None: import httpx from litellm.caching.caching import Cache + from litellm.caching.caching_handler import _PENDING_CACHE_WRITES from litellm.types.caching import LiteLLMCacheType monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") monkeypatch.setattr(litellm, "cache", Cache(type=LiteLLMCacheType.LOCAL)) - route = respx_mock.post("https://api.openai.com/v1/chat/completions").mock( + route: Final = respx_mock.post("https://api.openai.com/v1/chat/completions").mock( return_value=httpx.Response( 200, json={ @@ -1666,16 +1680,41 @@ async def test_generate_content_bridge_honors_the_request_cache_controls(respx_m }, ) ) - request = { + request: Final = { "model": "openai/gpt-4.1", "contents": [{"role": "user", "parts": [{"text": "What is the capital of France?"}]}], - "cache": {"no-cache": True, "no-store": True}, "api_key": "sk-test", } - await agenerate_content(**request) - await asyncio.sleep(0.5) - await agenerate_content(**request) + await agenerate_content(**request, **({} if first_cache_controls is None else {"cache": first_cache_controls})) + await asyncio.gather(*_PENDING_CACHE_WRITES) + await agenerate_content(**request, **({} if second_cache_controls is None else {"cache": second_cache_controls})) - assert route.call_count == 2 + assert route.call_count == 1 + + +@pytest.mark.parametrize( + ("semantic", "expected_cache"), + [(True, {"no-cache": True, "no-store": True}), (False, None)], + ids=["semantic", "exact-match"], +) +def test_generate_content_bridge_forwards_the_request_cache_controls_only_to_a_semantic_cache( + monkeypatch: pytest.MonkeyPatch, semantic: bool, expected_cache: Mapping[str, bool] | None +) -> None: + from litellm.caching.caching import Cache + from litellm.google_genai.adapters.handler import GenerateContentToCompletionHandler + from litellm.types.caching import LiteLLMCacheType + + cache_type: Final = LiteLLMCacheType.VALKEY_SEMANTIC if semantic else LiteLLMCacheType.LOCAL + monkeypatch.setattr( + litellm, "cache", Cache(type=cache_type, host="localhost", port="6379", similarity_threshold=0.8) + ) + + completion_kwargs: Final = GenerateContentToCompletionHandler._prepare_completion_kwargs( + model="openai/gpt-4.1", + contents=[{"role": "user", "parts": [{"text": "What is the capital of France?"}]}], + extra_kwargs={"cache": {"no-cache": True, "no-store": True}}, + ) + + assert completion_kwargs.get("cache") == expected_cache diff --git a/tests/unit/proxy/agent_endpoints/test_a2a_endpoints.py b/tests/unit/proxy/agent_endpoints/test_a2a_endpoints.py index 011df4f8daf..c7c501149f6 100644 --- a/tests/unit/proxy/agent_endpoints/test_a2a_endpoints.py +++ b/tests/unit/proxy/agent_endpoints/test_a2a_endpoints.py @@ -386,7 +386,12 @@ def _make_agent_mock(url: str = "http://backend-agent:10001") -> MagicMock: return agent -def _make_request_mock(method: str, params: Mapping[str, object], request_id: object = "req-1") -> MagicMock: +def _make_request_mock( + method: str, + params: Mapping[str, object], + request_id: object = "req-1", + cache: Mapping[str, object] | None = None, +) -> MagicMock: req = MagicMock() req.headers = {} req.json = AsyncMock( @@ -395,6 +400,7 @@ def _make_request_mock(method: str, params: Mapping[str, object], request_id: ob "id": request_id, "method": method, "params": params, + **({} if cache is None else {"cache": cache}), } ) return req @@ -662,6 +668,38 @@ async def test_message_methods_hand_the_semantic_cache_check_result_to_the_agent assert (captured.litellm_params or {}).get("cache") == {"ttl": 60, "no-cache": True, "no-store": True} +@pytest.mark.asyncio +@pytest.mark.parametrize("method", ["message/send", "message/stream"]) +@pytest.mark.parametrize( + ("cache_type", "expected_agent_cache"), + [ + (LiteLLMCacheType.VALKEY_SEMANTIC, {"namespace": "team-a", "no-store": True}), + (LiteLLMCacheType.LOCAL, {"namespace": "team-a"}), + ], + ids=["semantic", "exact-match"], +) +async def test_message_methods_forward_only_the_callers_semantic_cache_opt_outs_to_the_agent_call( + method: str, cache_type: LiteLLMCacheType, expected_agent_cache: Mapping[str, object] +) -> None: + agent: Final = _make_agent_mock() + agent.litellm_params = {"cache": {"namespace": "team-a"}} + mock_request: Final = _make_request_mock( + method, _HELLO_MESSAGE_PARAMS, cache={"namespace": "team-b", "use-cache": True, "ttl": 1, "no-store": True} + ) + cache: Final = Cache(type=cache_type, host="localhost", port="6379", similarity_threshold=0.8) + + with patch.object(litellm, "cache", cache): + captured: Final = await _invoke_message_method( + method, + mock_request, + UserAPIKeyAuth(api_key="sk-test", team_id="team-a"), + _add_proxy_data, + agent=agent, + ) + + assert (captured.litellm_params or {}).get("cache") == expected_agent_cache + + @pytest.mark.asyncio @pytest.mark.parametrize("method", ["message/send", "message/stream"]) async def test_message_methods_forward_caller_identity_headers(method: str): diff --git a/tests/unit/proxy/test_common_request_processing.py b/tests/unit/proxy/test_common_request_processing.py index 700c4958a69..8e52a263bc5 100644 --- a/tests/unit/proxy/test_common_request_processing.py +++ b/tests/unit/proxy/test_common_request_processing.py @@ -525,6 +525,7 @@ class TestProxyBaseLLMRequestProcessing: expected_cache_controls: dict[str, object], ) -> None: from litellm.caching.caching import Cache + from litellm.proxy.proxy_server import ProxyConfig as RealProxyConfig from litellm.types.caching import LiteLLMCacheType class _Gate(CustomLogger): @@ -533,37 +534,30 @@ class TestProxyBaseLLMRequestProcessing: ) -> bool: return gate_allows - async def passthrough_pre_call_hook( - user_api_key_dict: ProxyUserAPIKeyAuth, - data: dict[str, object], - call_type: str, - skip_guardrails: bool = False, - ) -> dict[str, object]: - return data - monkeypatch.setattr(litellm, "callbacks", [_Gate()]) monkeypatch.setattr( litellm, "cache", Cache(type=LiteLLMCacheType.VALKEY_SEMANTIC, host="localhost", port="6379", similarity_threshold=0.8), ) - body: Final = {"model": "gpt-4o-mini", "messages": [{"role": "user", "content": "hi"}], "cache": {"ttl": 60}} - monkeypatch.setattr( - litellm.proxy.common_request_processing, - "add_litellm_data_to_request", - AsyncMock(return_value={**body, "metadata": {}, "proxy_server_request": {"body": body}}), + request: Final = Request( + { + "type": "http", + "method": "POST", + "path": "/v1/chat/completions", + "headers": [(b"content-type", b"application/json")], + "query_string": b"", + } ) - mock_request: Final = MagicMock(spec=Request) - mock_request.headers = {} - mock_proxy_logging_obj: Final = MagicMock(spec=ProxyLogging) - mock_proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=passthrough_pre_call_hook) - returned_data, _ = await ProxyBaseLLMRequestProcessing(data={}).common_processing_pre_call_logic( - request=mock_request, + returned_data, _ = await ProxyBaseLLMRequestProcessing( + data={"model": "gpt-4o-mini", "messages": [{"role": "user", "content": "hi"}], "cache": {"ttl": 60}} + ).common_processing_pre_call_logic( + request=request, general_settings={}, user_api_key_dict=ProxyUserAPIKeyAuth(api_key="sk-hashed", team_id="team-other"), - proxy_logging_obj=mock_proxy_logging_obj, - proxy_config=MagicMock(spec=ProxyConfig), + proxy_logging_obj=ProxyLogging(user_api_key_cache=DualCache()), + proxy_config=RealProxyConfig(), route_type="acompletion", )