fix(proxy): run the semantic cache check on MCP sampling, merge A2A cache controls with the agent's own, and drop routes that never reach the cache

litellm/proxy/_experimental/mcp_server/sampling_handler.py

@@ -37,6 +37,8 @@ from fastapi import HTTPException
 from pydantic import TypeAdapter
 
 from litellm._logging import verbose_logger
+from litellm.proxy.common_utils.semantic_cache_gate import gated_cache_controls
+from litellm.proxy.utils import ProxyLogging
 
 # Guard imports that require the mcp package
 try:
@@ -1125,14 +1127,22 @@ async def _run_guardrails_and_call_llm(
 
     import litellm
 
+    cache_controls: Final = await gated_cache_controls(
+        data=completion_kwargs,
+        user_api_key_dict=user_api_key_auth,
+        call_type="acompletion",
+        gates=ProxyLogging.callback_capabilities().semantic_cache_gates,
+        cache=litellm.cache,
+    )
+    call_kwargs: Final = completion_kwargs if cache_controls is None else {**completion_kwargs, "cache": cache_controls}
     try:
         from litellm.proxy.proxy_server import llm_router
 
         if llm_router is not None:
-            return await _AcompletionCall(fn=llm_router.acompletion).fn(**completion_kwargs)
-        return await _AcompletionCall(fn=litellm.acompletion).fn(**completion_kwargs)
+            return await _AcompletionCall(fn=llm_router.acompletion).fn(**call_kwargs)
+        return await _AcompletionCall(fn=litellm.acompletion).fn(**call_kwargs)
     except ImportError:
-        return await _AcompletionCall(fn=litellm.acompletion).fn(**completion_kwargs)
+        return await _AcompletionCall(fn=litellm.acompletion).fn(**call_kwargs)
 
 
 async def handle_sampling_create_message(

litellm/proxy/agent_endpoints/a2a_endpoints.py

@@ -180,9 +180,12 @@ async def _resolve_backend_auth(
     )
 
 
-def _request_cache_controls(data: Mapping[str, object]) -> Mapping[str, object]:
-    cache_controls: Final = data.get("cache")
-    return {"cache": cache_controls} if isinstance(cache_controls, Mapping) else {}
+def _request_cache_controls(agent_params: Mapping[str, object], data: Mapping[str, object]) -> Mapping[str, object]:
+    request_cache: Final = data.get("cache")
+    if not isinstance(request_cache, Mapping):
+        return {}
+    agent_cache: Final = agent_params.get("cache")
+    return {"cache": {**(agent_cache if isinstance(agent_cache, Mapping) else {}), **request_cache}}
 
 
 def _forwarding_headers(
@@ -843,7 +846,10 @@ async def invoke_agent_a2a(
             ),
             backend_auth_header=backend_auth.header,
         )
-        forwarded_litellm_params: Final = {**backend_auth.litellm_params, **_request_cache_controls(data)}
+        forwarded_litellm_params: Final = {
+            **backend_auth.litellm_params,
+            **_request_cache_controls(backend_auth.litellm_params, data),
+        }
 
         # Merge agent-level guardrails into data so post_call_success_hook and
         # _handle_stream_message both pick them up.  A2A agents use model

litellm/proxy/common_utils/semantic_cache_gate.py

@@ -13,8 +13,6 @@ SEMANTIC_CACHE_ROUTES: Final = frozenset(
         "acompletion",
         "atext_completion",
         "aembedding",
-        "atranscription",
-        "arerank",
         "aresponses",
         "_aresponses_websocket",
         "anthropic_messages",

tests/unit/proxy/_experimental/mcp_server/test_mcp_sampling_completion_flow.py

@@ -7,6 +7,7 @@ through the proxy router / guardrails, and the end-to-end
 `handle_sampling_create_message` success and error-propagation behaviour.
 """
 
+from collections.abc import Mapping
 from types import SimpleNamespace
 from unittest.mock import AsyncMock, MagicMock, patch
 
@@ -14,6 +15,12 @@ import pytest
 
 from mcp.types import CreateMessageResult, ErrorData
 
+import litellm
+from litellm.caching.caching import Cache
+from litellm.integrations.custom_logger import CustomLogger
+from litellm.proxy._types import UserAPIKeyAuth
+from litellm.types.caching import LiteLLMCacheType
+
 from litellm.proxy._experimental.mcp_server.sampling_handler import (
     _build_completion_kwargs,
     _run_guardrails_and_call_llm,
@@ -119,6 +126,35 @@ class TestRunGuardrailsAndCallLlm:
         assert result == "router-response"
         router.acompletion.assert_awaited_once()
 
+    async def test_should_skip_the_semantic_cache_for_a_key_the_semantic_cache_check_rejects(self):
+        class _RejectEveryTeam(CustomLogger):
+            async def async_should_semantic_cache(
+                self, user_api_key_dict: UserAPIKeyAuth, data: Mapping[str, object], call_type: str
+            ) -> bool:
+                return False
+
+        async def _cache_controls_the_router_received(**kwargs):
+            return kwargs.get("cache")
+
+        semantic_cache = Cache(
+            type=LiteLLMCacheType.VALKEY_SEMANTIC, host="localhost", port="6379", similarity_threshold=0.8
+        )
+        with (
+            patch("litellm.proxy.proxy_server.proxy_logging_obj", None),
+            patch(
+                "litellm.proxy.proxy_server.llm_router",
+                SimpleNamespace(acompletion=_cache_controls_the_router_received),
+            ),
+            patch.object(litellm, "callbacks", [_RejectEveryTeam()]),
+            patch.object(litellm, "cache", semantic_cache),
+        ):
+            received = await _run_guardrails_and_call_llm(
+                completion_kwargs={"model": "gpt-4o", "messages": [{"role": "user", "content": "hi"}]},
+                user_api_key_auth=UserAPIKeyAuth(api_key="sk-test", team_id="team-other"),
+            )
+
+        assert received == {"no-cache": True, "no-store": True}
+
     async def test_should_propagate_guardrail_rejection(self):
         plo = MagicMock()
         plo.pre_call_hook = AsyncMock(side_effect=ValueError("blocked by guardrail"))

tests/unit/proxy/agent_endpoints/test_a2a_endpoints.py

@@ -643,6 +643,8 @@ class _RejectEveryTeam(CustomLogger):
 @pytest.mark.asyncio
 @pytest.mark.parametrize("method", ["message/send", "message/stream"])
 async def test_message_methods_hand_the_semantic_cache_check_result_to_the_agent_call(method: str):
+    agent = _make_agent_mock()
+    agent.litellm_params = {"cache": {"ttl": 60}}
     mock_request = _make_request_mock(method, _HELLO_MESSAGE_PARAMS)
     semantic_cache = Cache(
         type=LiteLLMCacheType.VALKEY_SEMANTIC, host="localhost", port="6379", similarity_threshold=0.8
@@ -650,10 +652,14 @@ async def test_message_methods_hand_the_semantic_cache_check_result_to_the_agent
 
     with patch.object(litellm, "callbacks", [_RejectEveryTeam()]), patch.object(litellm, "cache", semantic_cache):
         captured = await _invoke_message_method(
-            method, mock_request, UserAPIKeyAuth(api_key="sk-test", team_id="team-other"), _add_proxy_data
+            method,
+            mock_request,
+            UserAPIKeyAuth(api_key="sk-test", team_id="team-other"),
+            _add_proxy_data,
+            agent=agent,
         )
 
-    assert (captured.litellm_params or {}).get("cache") == {"no-cache": True, "no-store": True}
+    assert (captured.litellm_params or {}).get("cache") == {"ttl": 60, "no-cache": True, "no-store": True}
 
 
 @pytest.mark.asyncio

tests/unit/proxy/common_utils/test_semantic_cache_gate.py

@@ -112,7 +112,7 @@ async def test_the_gate_is_never_consulted_without_a_semantic_cache(cache: Cache
 
 
 @pytest.mark.asyncio
-@pytest.mark.parametrize("route", ["acreate_batch", "avector_store_search", "aspeech"])
+@pytest.mark.parametrize("route", ["acreate_batch", "avector_store_search", "aspeech", "atranscription", "arerank"])
 async def test_the_gate_is_never_consulted_for_a_route_that_cannot_reach_the_cache(route: str):
     gate = _TeamAllowlist(frozenset())