From 640223719f5f648a2fb6f11acf01cfb5843e2baa Mon Sep 17 00:00:00 2001 From: shrey Date: Tue, 3 Nov 2026 20:42:10 -0800 Subject: [PATCH] feat(mcp): add an opt-in per-server setting to skip MCP tool output validation (#49961) --- litellm/mcp_client/client.py | 18 +- .../proxy/mcp_server/mcp_server_manager.py | 2 +- litellm/proxy/mcp_server/upstream.py | 3 + .../types/mcp_server/mcp_server_manager.py | 6 + tests/_support/mcp_count_tool_upstream.py | 42 ++++ tests/unit/mcp_client/test_mcp_client.py | 54 +++++ tests/unit/proxy/mcp_server/test_upstream.py | 192 ++++++++++++++++++ tests/unit/proxy/test__types.py | 33 +++ 8 files changed, 347 insertions(+), 3 deletions(-) create mode 100644 tests/_support/mcp_count_tool_upstream.py create mode 100644 tests/unit/proxy/mcp_server/test_upstream.py diff --git a/litellm/mcp_client/client.py b/litellm/mcp_client/client.py index 107b7155303..0461f1f06ff 100644 --- a/litellm/mcp_client/client.py +++ b/litellm/mcp_client/client.py @@ -478,6 +478,13 @@ def _notify_cleanup(callback: Callable[[], None] | None) -> None: callback() +class _UnvalidatedClientSession(ClientSession): + """Skips the SDK's output check, which lists the upstream's tools on a session with no schema cached.""" + + async def validate_tool_result(self, name: str, result: MCPCallToolResult) -> None: + return None + + class MCPClient: """ MCP Client supporting: @@ -505,12 +512,15 @@ class MCPClient: elicitation_callback: "ElicitationCallback | None" = None, logging_callback: Callable | None = None, protocol_version: MCPUpstreamProtocol = "auto", + skip_tool_output_validation: bool = False, upstream_transport: httpx2.AsyncBaseTransport | None = None, ): self.protocol_version: MCPUpstreamProtocol = TypeAdapter[MCPUpstreamProtocol]( MCPUpstreamProtocol ).validate_python(protocol_version) validate_mcp_protocol_transport(self.protocol_version, transport_type) + # 2026-07-28 lists tools before every call, so its output check adds no list + self.skip_tool_output_validation: Final = skip_tool_output_validation and self.protocol_version != "2026-07-28" self.server_url: str = server_url self.transport_type: MCPTransport = transport_type self.auth_type: MCPAuthType = auth_type @@ -546,7 +556,7 @@ class MCPClient: "response_hooks": _mcp_response_hooks(), } - def session_settings(self) -> tuple[MCPTransport, MCPUpstreamProtocol, float, bool, bool, bool]: + def session_settings(self) -> tuple[MCPTransport, MCPUpstreamProtocol, float, bool, bool, bool, bool]: callbacks: Final = self.session_callbacks() return ( self.transport_type, @@ -555,6 +565,7 @@ class MCPClient: callbacks["sampling_callback"] is not None, callbacks["elicitation_callback"] is not None, callbacks["logging_callback"] is not None, + self.skip_tool_output_validation, ) def route_session_callbacks(self, callbacks: _SessionCallbackMap) -> None: @@ -776,7 +787,10 @@ class MCPClient: ) # The SDK drops a response stream that ends without a JSON-RPC reply, so nothing else # ever fails the request. - session_ctx: Final = ClientSession( + session_type: Final = ( + _UnvalidatedClientSession if self.skip_tool_output_validation else ClientSession + ) + session_ctx: Final = session_type( read_stream, write_stream, read_timeout_seconds=self.timeout, diff --git a/litellm/proxy/mcp_server/mcp_server_manager.py b/litellm/proxy/mcp_server/mcp_server_manager.py index eacba25257e..48e7370d36d 100644 --- a/litellm/proxy/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/mcp_server/mcp_server_manager.py @@ -367,7 +367,7 @@ _UpstreamSessionKey: TypeAlias = tuple[ str, str, str, - tuple[MCPTransport, MCPUpstreamProtocol, float, bool, bool, bool], + tuple[MCPTransport, MCPUpstreamProtocol, float, bool, bool, bool, bool], ] diff --git a/litellm/proxy/mcp_server/upstream.py b/litellm/proxy/mcp_server/upstream.py index 96a1c2db9ba..3dd245947db 100644 --- a/litellm/proxy/mcp_server/upstream.py +++ b/litellm/proxy/mcp_server/upstream.py @@ -231,6 +231,7 @@ async def prepare_upstream_client( server_url="", transport_type=transport, protocol_version=protocol_version, + skip_tool_output_validation=server.skip_tool_output_validation, auth_type=server.auth_type, auth_value=auth_value, timeout=timeout, @@ -265,6 +266,7 @@ async def prepare_upstream_client( server_url=server.url or "", transport_type=transport, protocol_version=protocol_version, + skip_tool_output_validation=server.skip_tool_output_validation, upstream_transport=upstream_transport, auth_type=server.auth_type, timeout=timeout, @@ -308,6 +310,7 @@ async def prepare_upstream_client( server_url=server.url or "", transport_type=transport, protocol_version=protocol_version, + skip_tool_output_validation=server.skip_tool_output_validation, upstream_transport=upstream_transport, auth_type=server.auth_type, auth_value=auth_value, diff --git a/litellm/types/mcp_server/mcp_server_manager.py b/litellm/types/mcp_server/mcp_server_manager.py index 0d39b0a0616..1d3d5911eb1 100644 --- a/litellm/types/mcp_server/mcp_server_manager.py +++ b/litellm/types/mcp_server/mcp_server_manager.py @@ -23,6 +23,8 @@ from litellm.types.mcp import ( def _validate_mcp_protocol_metadata(value: dict[str, object]) -> dict[str, object]: if "protocol_version" in value: TypeAdapter[MCPUpstreamProtocol](MCPUpstreamProtocol).validate_python(value["protocol_version"]) + if "skip_tool_output_validation" in value: + TypeAdapter[bool](bool).validate_python(value["skip_tool_output_validation"], strict=True) return value @@ -297,6 +299,10 @@ class MCPServer(LiteLLMBaseModel): """ return self.oauth2_flow == "client_credentials" + @property + def skip_tool_output_validation(self) -> bool: + return (self.mcp_info or {}).get("skip_tool_output_validation") is True + @model_validator(mode="after") def resolve_protocol_version(self) -> Self: if "protocol_version" not in self.model_fields_set and self.mcp_info is not None: diff --git a/tests/_support/mcp_count_tool_upstream.py b/tests/_support/mcp_count_tool_upstream.py new file mode 100644 index 00000000000..f1aa61c7dba --- /dev/null +++ b/tests/_support/mcp_count_tool_upstream.py @@ -0,0 +1,42 @@ +import json +import queue +from collections.abc import Callable, Mapping +from typing import Final + +import httpx2 + +COUNT_OUTPUT_SCHEMA: Final = {"type": "object", "properties": {"n": {"type": "integer"}}, "required": ["n"]} + + +def count_tool_upstream( + structured: Mapping[str, object], methods_seen: "queue.SimpleQueue[str]" +) -> Callable[[httpx2.Request], httpx2.Response]: + """An MCP upstream with one tool, `count`, whose declared output schema wants an integer `n`.""" + + def respond(request: httpx2.Request) -> httpx2.Response: + if request.method != "POST": + return httpx2.Response(200) + payload: Final = json.loads(request.content) + methods_seen.put(payload["method"]) + if "id" not in payload: + return httpx2.Response(202) + results: Final[Mapping[str, Mapping[str, object]]] = { + "initialize": { + "protocolVersion": payload["params"].get("protocolVersion"), + "capabilities": {"tools": {}}, + "serverInfo": {"name": "count", "version": "1"}, + }, + "tools/list": { + "tools": [{"name": "count", "inputSchema": {"type": "object"}, "outputSchema": COUNT_OUTPUT_SCHEMA}] + }, + "tools/call": {"content": [{"type": "text", "text": "ok"}], "structuredContent": structured}, + } + return httpx2.Response( + 200, json={"jsonrpc": "2.0", "id": payload["id"], "result": results.get(payload["method"], {})} + ) + + return respond + + +def methods_received(methods_seen: "queue.SimpleQueue[str]") -> tuple[str, ...]: + return tuple(methods_seen.get_nowait() for _ in range(methods_seen.qsize())) diff --git a/tests/unit/mcp_client/test_mcp_client.py b/tests/unit/mcp_client/test_mcp_client.py index 43021d248a9..866a29b70e9 100644 --- a/tests/unit/mcp_client/test_mcp_client.py +++ b/tests/unit/mcp_client/test_mcp_client.py @@ -3,6 +3,7 @@ import base64 import importlib import json import os +import queue import sys from collections.abc import AsyncIterator, Callable from contextlib import AbstractAsyncContextManager, asynccontextmanager @@ -74,6 +75,7 @@ from litellm.proxy.mcp_server.outbound_credentials.httpx_auth import StaticHeade from litellm.types.mcp import MCPAuth, MCPStdioConfig, MCPTransport from litellm.types.mcp_server.mcp_server_manager import MCPServer from tests._support.clock_loops import AutojumpClockLoop, ManualClockLoop +from tests._support.mcp_count_tool_upstream import count_tool_upstream, methods_received _JSONRPC_MESSAGE_ADAPTER: Final = TypeAdapter(JSONRPCMessage) @@ -5250,3 +5252,55 @@ async def test_notification_callback_keeps_stream_errors_separate_from_protocol_ await _receive_session_message(failure, stream_error=stream_error, on_notification=receiver) assert stream_error.result() is first receiver.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_a_skipping_client_never_asks_the_upstream_for_its_tools() -> None: + methods_seen: Final[queue.SimpleQueue[str]] = queue.SimpleQueue() + client: Final = _MockTransportClient( + count_tool_upstream({"n": 1}, methods_seen), + server_url="https://upstream.example/mcp", + transport_type=MCPTransport.http, + skip_tool_output_validation=True, + ) + + result: Final = await client.call_tool(CallToolRequestParams(name="count", arguments={})) + + assert isinstance(result, CallToolResult) + assert result.is_error is False + assert "tools/list" not in methods_received(methods_seen) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("skip_validation", "rejected"), [(False, True), (True, False)], ids=["validating", "skipping"] +) +async def test_invalid_structured_output_is_rejected_unless_validation_is_skipped( + skip_validation: bool, rejected: bool +) -> None: + client: Final = _MockTransportClient( + count_tool_upstream({"n": "not-an-integer"}, queue.SimpleQueue()), + server_url="https://upstream.example/mcp", + transport_type=MCPTransport.http, + skip_tool_output_validation=skip_validation, + ) + + result: Final = await client.call_tool(CallToolRequestParams(name="count", arguments={})) + + assert isinstance(result, CallToolResult) + assert result.is_error is rejected + + +@pytest.mark.asyncio +async def test_a_skipped_validation_passes_the_upstreams_result_through_unchanged() -> None: + client: Final = _MockTransportClient( + count_tool_upstream({"n": "not-an-integer"}, queue.SimpleQueue()), + server_url="https://upstream.example/mcp", + transport_type=MCPTransport.http, + skip_tool_output_validation=True, + ) + + result: Final = await client.call_tool(CallToolRequestParams(name="count", arguments={})) + + assert isinstance(result, CallToolResult) + assert result.structured_content == {"n": "not-an-integer"} diff --git a/tests/unit/proxy/mcp_server/test_upstream.py b/tests/unit/proxy/mcp_server/test_upstream.py new file mode 100644 index 00000000000..7c88c679972 --- /dev/null +++ b/tests/unit/proxy/mcp_server/test_upstream.py @@ -0,0 +1,192 @@ +import queue +from collections.abc import Iterator, Mapping +from typing import Final +from unittest.mock import MagicMock + +import httpx2 +import pytest +from mcp.server.context import ServerRequestContext +from mcp.types import CallToolResult + +from litellm.proxy.mcp_server.mcp_context import active_mcp_request_ctx_var +from litellm.proxy.mcp_server.mcp_server_manager import MCPServerManager +from litellm.proxy.mcp_server.outbound_credentials.resolver import UpstreamCredentialProvider +from litellm.proxy.mcp_server.upstream import prepare_upstream_client +from litellm.types.mcp import MCPAuth, MCPTransport, MCPUpstreamProtocol +from litellm.types.mcp_server.mcp_server_manager import MCPServer +from tests._support.mcp_count_tool_upstream import count_tool_upstream, methods_received + +_HTTP_URL: Final = "https://upstream.example/mcp" +_SERVER_SHAPES: Final = pytest.mark.parametrize( + ("server_fields", "per_request_token", "forwarded_token"), + [ + ({"transport": MCPTransport.http, "url": _HTTP_URL}, None, None), + ( + { + "transport": MCPTransport.http, + "url": _HTTP_URL, + "auth_type": MCPAuth.bearer_token, + "authentication_token": "static-token", + }, + None, + None, + ), + ({"transport": MCPTransport.http, "url": _HTTP_URL, "auth_type": MCPAuth.bearer_token}, "caller-key", None), + ({"transport": MCPTransport.http, "url": _HTTP_URL, "auth_type": MCPAuth.oauth_delegate}, None, "caller-token"), + ], + ids=["no-auth", "static-bearer-token", "per-request-token", "forwarded-caller-token"], +) + + +@pytest.fixture +def _active_mcp_request_context() -> Iterator[None]: + context: Final = ServerRequestContext( + session=MagicMock(), lifespan_context={}, protocol_version="2025-11-25", method="tools/call", request=None + ) + token: Final = active_mcp_request_ctx_var.set(context) + try: + yield + finally: + active_mcp_request_ctx_var.reset(token) + + +def _server(server_fields: Mapping[str, object], skip_setting: bool | None) -> MCPServer: + return MCPServer.model_validate( + { + **server_fields, + "server_id": "skip-validation", + "name": "skip_validation", + "mcp_info": {} if skip_setting is None else {"skip_tool_output_validation": skip_setting}, + } + ) + + +async def _call_count_tool( + manager: MCPServerManager, + server: MCPServer, + *, + per_request_token: str | None = None, + forwarded_token: str | None = None, + gateway_session_id: str | None = None, +) -> CallToolResult: + result: Final = await manager._call_regular_mcp_tool( + mcp_server=server, + original_tool_name="count", + arguments={}, + tasks=[], + mcp_auth_header=per_request_token, + mcp_server_auth_headers=None, + oauth2_headers=None if forwarded_token is None else {"Authorization": f"Bearer {forwarded_token}"}, + raw_headers=None if gateway_session_id is None else {"mcp-session-id": gateway_session_id}, + proxy_logging_obj=None, + ) + assert isinstance(result, CallToolResult) + return result + + +@pytest.mark.asyncio +@_SERVER_SHAPES +@pytest.mark.parametrize( + ("skip_setting", "rejected"), [(None, True), (False, True), (True, False)], ids=["unset", "off", "on"] +) +async def test_the_setting_alone_decides_whether_a_tool_result_is_checked( + server_fields: Mapping[str, object], + per_request_token: str | None, + forwarded_token: str | None, + skip_setting: bool | None, + rejected: bool, +) -> None: + upstream: Final = count_tool_upstream({"n": "not-an-integer"}, queue.SimpleQueue()) + manager: Final = MCPServerManager(upstream_transport=httpx2.MockTransport(upstream)) + + result: Final = await _call_count_tool( + manager, + _server(server_fields, skip_setting), + per_request_token=per_request_token, + forwarded_token=forwarded_token, + ) + + assert result.is_error is rejected + + +@pytest.mark.asyncio +@_SERVER_SHAPES +async def test_a_skipping_server_is_never_asked_for_its_tools( + server_fields: Mapping[str, object], per_request_token: str | None, forwarded_token: str | None +) -> None: + methods_seen: Final[queue.SimpleQueue[str]] = queue.SimpleQueue() + manager: Final = MCPServerManager( + upstream_transport=httpx2.MockTransport(count_tool_upstream({"n": 1}, methods_seen)) + ) + + result: Final = await _call_count_tool( + manager, + _server(server_fields, True), + per_request_token=per_request_token, + forwarded_token=forwarded_token, + ) + + assert result.is_error is False + assert "tools/list" not in methods_received(methods_seen) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("first_setting", "second_setting", "rejected_in_order"), + [(True, False, (False, True)), (False, True, (True, False))], + ids=["turned-off", "turned-on"], +) +async def test_a_changed_setting_applies_to_the_next_call_in_a_live_gateway_session( + _active_mcp_request_context: None, + first_setting: bool, + second_setting: bool, + rejected_in_order: tuple[bool, bool], +) -> None: + upstream: Final = count_tool_upstream({"n": "not-an-integer"}, queue.SimpleQueue()) + manager: Final = MCPServerManager(upstream_transport=httpx2.MockTransport(upstream)) + manager.track_gateway_session("gateway-1") + fields: Final = {"transport": MCPTransport.http, "url": _HTTP_URL} + try: + first: Final = await _call_count_tool(manager, _server(fields, first_setting), gateway_session_id="gateway-1") + second: Final = await _call_count_tool(manager, _server(fields, second_setting), gateway_session_id="gateway-1") + finally: + await manager.close_all_upstream_sessions() + + assert (first.is_error, second.is_error) == rejected_in_order + + +@pytest.mark.asyncio +async def test_a_stdio_upstream_client_carries_the_setting(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("LITELLM_ENABLE_MCP_STDIO", "true") + server: Final = _server({"transport": MCPTransport.stdio, "command": "python", "args": ["server.py"]}, True) + + client: Final = await prepare_upstream_client( + server, provider=UpstreamCredentialProvider(), root_path="/mcp", protocol_version="auto" + ) + + assert client.skip_tool_output_validation is True + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("protocol_version", "expected"), [("auto", True), ("2025-11-25", True), ("2026-07-28", False)] +) +async def test_upstream_client_ignores_the_setting_where_the_call_already_lists_tools( + protocol_version: MCPUpstreamProtocol, expected: bool +) -> None: + server: Final = MCPServer.model_validate( + { + "server_id": "skip-validation", + "name": "skip_validation", + "url": _HTTP_URL, + "transport": MCPTransport.http, + "protocol_version": protocol_version, + "mcp_info": {"skip_tool_output_validation": True}, + } + ) + + client: Final = await prepare_upstream_client( + server, provider=UpstreamCredentialProvider(), root_path="/mcp", protocol_version=protocol_version + ) + + assert client.skip_tool_output_validation is expected diff --git a/tests/unit/proxy/test__types.py b/tests/unit/proxy/test__types.py index 76ca49362b4..40e66f1da65 100644 --- a/tests/unit/proxy/test__types.py +++ b/tests/unit/proxy/test__types.py @@ -520,6 +520,39 @@ def test_mcp_metadata_rejects_unavailable_upstream_protocol(revision): model.model_validate(payload) +@pytest.mark.parametrize("request_model", (NewMCPServerRequest, UpdateMCPServerRequest)) +@pytest.mark.parametrize("value", ["true", "false", 1, 0, None]) +def test_mcp_metadata_rejects_a_non_boolean_skip_tool_output_validation( + request_model: type[NewMCPServerRequest] | type[UpdateMCPServerRequest], value: object +) -> None: + payload: Final = { + "server_id": "test", + "transport": "http", + "url": "https://example.com/mcp", + "mcp_info": {"skip_tool_output_validation": value}, + } + + with pytest.raises(ValidationError): + request_model.model_validate(payload) + + +@pytest.mark.parametrize("request_model", (NewMCPServerRequest, UpdateMCPServerRequest)) +@pytest.mark.parametrize("value", [True, False]) +def test_mcp_metadata_accepts_a_boolean_skip_tool_output_validation( + request_model: type[NewMCPServerRequest] | type[UpdateMCPServerRequest], value: bool +) -> None: + parsed: Final = request_model.model_validate( + { + "server_id": "test", + "transport": "http", + "url": "https://example.com/mcp", + "mcp_info": {"skip_tool_output_validation": value}, + } + ) + + assert parsed.mcp_info == {"skip_tool_output_validation": value} + + MCP_SERVER_REQUESTS = (NewMCPServerRequest, UpdateMCPServerRequest) STDIO_SERVER_FIELDS = {"server_id": "stdio-1", "transport": "stdio", "command": "python", "args": ["server.py"]}