From 34bbeaaa381292b1c4aa125757a768783c63556b Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 21 Oct 2026 01:54:16 -0700 Subject: [PATCH] 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 --- .../mcp_server/sampling_handler.py | 16 +++++++-- .../proxy/agent_endpoints/a2a_endpoints.py | 14 +++++--- .../proxy/common_utils/semantic_cache_gate.py | 2 -- .../test_mcp_sampling_completion_flow.py | 36 +++++++++++++++++++ .../agent_endpoints/test_a2a_endpoints.py | 10 ++++-- .../common_utils/test_semantic_cache_gate.py | 2 +- 6 files changed, 68 insertions(+), 12 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/sampling_handler.py b/litellm/proxy/_experimental/mcp_server/sampling_handler.py index e1469f2fff7..64f5530fa5c 100644 --- a/litellm/proxy/_experimental/mcp_server/sampling_handler.py +++ b/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( diff --git a/litellm/proxy/agent_endpoints/a2a_endpoints.py b/litellm/proxy/agent_endpoints/a2a_endpoints.py index f883d49532a..6a0f4462160 100644 --- a/litellm/proxy/agent_endpoints/a2a_endpoints.py +++ b/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 diff --git a/litellm/proxy/common_utils/semantic_cache_gate.py b/litellm/proxy/common_utils/semantic_cache_gate.py index b080c166f2d..8b45ce58031 100644 --- a/litellm/proxy/common_utils/semantic_cache_gate.py +++ b/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", diff --git a/tests/unit/proxy/_experimental/mcp_server/test_mcp_sampling_completion_flow.py b/tests/unit/proxy/_experimental/mcp_server/test_mcp_sampling_completion_flow.py index d17b407a1be..30b9617c42d 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_mcp_sampling_completion_flow.py +++ b/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")) diff --git a/tests/unit/proxy/agent_endpoints/test_a2a_endpoints.py b/tests/unit/proxy/agent_endpoints/test_a2a_endpoints.py index 223b247a779..011df4f8daf 100644 --- a/tests/unit/proxy/agent_endpoints/test_a2a_endpoints.py +++ b/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 diff --git a/tests/unit/proxy/common_utils/test_semantic_cache_gate.py b/tests/unit/proxy/common_utils/test_semantic_cache_gate.py index 9c95d28973d..eef95c5e589 100644 --- a/tests/unit/proxy/common_utils/test_semantic_cache_gate.py +++ b/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())