diff --git a/litellm/experimental_mcp_client/client.py b/litellm/experimental_mcp_client/client.py index 3503468c735..cabcc6c03ba 100644 --- a/litellm/experimental_mcp_client/client.py +++ b/litellm/experimental_mcp_client/client.py @@ -18,6 +18,7 @@ from mcp import ClientSession, McpError, ReadResourceResult, Resource, StdioServ from mcp.client.sse import sse_client from mcp.client.stdio import stdio_client from mcp.shared.message import SessionMessage +from mcp.shared.session import RequestResponder from typing_extensions import Unpack _TransportStreams: TypeAlias = tuple[ @@ -56,10 +57,13 @@ def missing_streamable_http_client_error() -> ImportError: from mcp.types import CallToolRequestParams as MCPCallToolRequestParams from mcp.types import CallToolResult as MCPCallToolResult from mcp.types import ( + ClientResult, GetPromptRequestParams, GetPromptResult, Prompt, ResourceTemplate, + ServerNotification, + ServerRequest, TextContent, ) from mcp.types import Tool as MCPTool @@ -442,6 +446,18 @@ class MCPClient: in_flight_error: BaseException | None = None try: read_stream, write_stream = transport[0], transport[1] + stream_error: Final[asyncio.Future[Exception]] = asyncio.get_running_loop().create_future() + + async def receive_message( + message: RequestResponder[ServerRequest, ClientResult] | ServerNotification | Exception, + ) -> None: + if not isinstance(message, ValueError): + return + if not stream_error.done(): + stream_error.set_result(message) + # The SDK closes pending requests when its message handler raises. + raise RuntimeError("MCP response stream failed") + # Build session kwargs with optional callbacks session_kwargs: Final[dict[str, Any]] = {} if self._sampling_callback is not None: @@ -456,6 +472,7 @@ class MCPClient: read_stream, write_stream, read_timeout_seconds=timedelta(seconds=self.timeout), + message_handler=receive_message if self.transport_type == MCPTransport.http else None, **session_kwargs, ) session: Final = await session_ctx.__aenter__() @@ -467,6 +484,10 @@ class MCPClient: if isinstance(ins, str) and ins.strip(): self._last_initialize_instructions = ins.strip() return await operation(session) + except McpError: + if stream_error.done(): + raise stream_error.result() + raise finally: try: await session_ctx.__aexit__(None, None, None) diff --git a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py index 329dddbdf05..c0ffeaeac62 100644 --- a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py @@ -3,12 +3,15 @@ import importlib from collections.abc import Awaitable, Callable, Mapping, Sequence from dataclasses import dataclass from datetime import datetime +from traceback import walk_tb from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal +from uuid import uuid4 import anyio import httpx from fastapi import APIRouter, Depends, HTTPException, Query, Request, status +from pydantic import ValidationError from starlette.datastructures import Headers from litellm._logging import verbose_logger @@ -30,6 +33,8 @@ from litellm.proxy._experimental.mcp_server.faults.list_outcomes import ( list_fault_http_status, outcome_wire_value, ) +from litellm.proxy._experimental.mcp_server.faults.traversal import iter_exception_tree +from litellm.proxy._experimental.mcp_server.oauth_utils import _redact_mcp_resource_url from litellm.proxy._experimental.mcp_server.ui_session_utils import ( acting_user_auth, build_effective_auth_contexts, @@ -78,11 +83,38 @@ _MCP_GUARDRAIL_REJECTIONS: Final = ( def _connection_error_message(exc: BaseException, url: str | None, timeout_seconds: float) -> str: + reference: Final = uuid4().hex + verbose_logger.error( + "MCP connection test failed (reference=%s): %s", + reference, + tuple( + ( + type(cause).__name__, + tuple( + (frame.f_code.co_filename, lineno, frame.f_code.co_name) + for frame, lineno in walk_tb(cause.__traceback__) + ), + ) + for cause in iter_exception_tree(exc) + ), + ) + return next( + ( + message + for cause in iter_exception_tree(exc) + if (message := _known_connection_error_message(cause, url, timeout_seconds)) is not None + ), + "An unexpected error occurred while testing the MCP connection. " + f"Retry; if it persists, share reference {reference} with your gateway administrator.", + ) + + +def _known_connection_error_message(exc: BaseException, url: str | None, timeout_seconds: float) -> str | None: if isinstance(exc, MCPServerURLCredentialsError): return str(exc.detail) if isinstance(exc, TimeoutError): return ( - f"Failed to connect to MCP server: no response from {url or 'the server'} " + f"Failed to connect to MCP server: no response from {_redact_mcp_resource_url(url) or 'the server'} " f"within {timeout_seconds:.0f}s. Check that the LiteLLM proxy can reach this URL " "from its network (DNS, egress rules, firewalls) and that the server answers MCP requests." ) @@ -99,10 +131,32 @@ def _connection_error_message(exc: BaseException, url: str | None, timeout_secon return "Failed to connect to MCP server: the connection timed out." if isinstance(exc, httpx.HTTPStatusError): return f"Failed to connect to MCP server: it returned HTTP {exc.response.status_code}." - return "Failed to connect to MCP server. Check proxy logs for details." + if isinstance(exc, ValueError) and str(exc).startswith("Unexpected content type:"): + return ( + "Failed to connect to MCP server: the endpoint returned an unsupported content type. " + "Check that the URL is an MCP endpoint, not a web page, and matches the selected transport." + ) + if isinstance(exc, ValidationError) and exc.title in ("JSONRPCMessage", "InitializeResult", "ListToolsResult"): + return ( + "Failed to connect to MCP server: the endpoint returned invalid JSON or an invalid MCP response. " + "Check the MCP endpoint URL and the server's protocol implementation." + ) + if MCP_AVAILABLE and isinstance(exc, McpError): + if exc.error.code == 32600 and exc.error.message == "Session terminated": + return ( + "Failed to connect to MCP server: the MCP session was terminated. " + "Check that the URL points to an MCP endpoint and matches the selected transport, " + "then retry to start a new session." + ) + return ( + f"Failed to connect to MCP server: the MCP request failed (JSON-RPC code {exc.error.code}). " + "Check that the endpoint supports MCP initialization and tool listing, and check the upstream server logs." + ) + return None if MCP_AVAILABLE: + from mcp.shared.exceptions import McpError from mcp.types import Tool as MCPTool from litellm.experimental_mcp_client.client import MCPClient @@ -1342,7 +1396,6 @@ if MCP_AVAILABLE: except (KeyboardInterrupt, SystemExit, asyncio.CancelledError): raise except BaseException as e: - verbose_logger.error("Error in MCP operation: %s", e, exc_info=True) return { "status": "error", "error": True, diff --git a/tests/test_litellm/experimental_mcp_client/test_mcp_client.py b/tests/test_litellm/experimental_mcp_client/test_mcp_client.py index 4db131da62c..49268b4d22d 100644 --- a/tests/test_litellm/experimental_mcp_client/test_mcp_client.py +++ b/tests/test_litellm/experimental_mcp_client/test_mcp_client.py @@ -1,9 +1,11 @@ import asyncio import base64 +import json import os import sys from importlib import metadata from pathlib import Path +from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import anyio @@ -11,6 +13,8 @@ import httpx import pytest from litellm.proxy._experimental.mcp_server.outbound_credentials.httpx_auth import StaticHeaderAuth from mcp import McpError +from mcp.client.streamable_http import streamable_http_client +from pydantic import ValidationError from mcp.shared.message import SessionMessage from mcp.types import ( LATEST_PROTOCOL_VERSION, @@ -1224,14 +1228,14 @@ def test_without_a_configured_slot_the_existing_precedence_is_unchanged(): _REDIRECT_CASES = [ - ("https://upstream.example.com/mcp", "https://upstream.example.com/other"), # same origin + ("https://upstream.example.com/mcp", "https://upstream.example.com/other"), # same origin ("https://upstream.example.com/mcp", "https://upstream.example.com:443/other"), # explicit default port - ("https://upstream.example.com/mcp", "https://attacker.example.com/collect"), # different host - ("https://upstream.example.com/mcp", "http://upstream.example.com/collect"), # scheme downgrade - ("https://upstream.example.com/mcp", "https://upstream.example.com:8443/other"), # different port - ("https://upstream.example.com/mcp", "https://sub.upstream.example.com/x"), # different host - ("http://upstream.example.com/mcp", "https://upstream.example.com/other"), # http -> https upgrade - ("http://upstream.example.com/mcp", "http://upstream.example.com/other"), # same origin, plain http + ("https://upstream.example.com/mcp", "https://attacker.example.com/collect"), # different host + ("https://upstream.example.com/mcp", "http://upstream.example.com/collect"), # scheme downgrade + ("https://upstream.example.com/mcp", "https://upstream.example.com:8443/other"), # different port + ("https://upstream.example.com/mcp", "https://sub.upstream.example.com/x"), # different host + ("http://upstream.example.com/mcp", "https://upstream.example.com/other"), # http -> https upgrade + ("http://upstream.example.com/mcp", "http://upstream.example.com/other"), # same origin, plain http ] @@ -1283,3 +1287,109 @@ def test_a_differently_cased_injected_header_cannot_shadow_the_slot() -> None: headers = client._get_auth_headers() assert [v for k, v in headers.items() if k.lower() == "esb-oauth"] == ["Bearer minted-token"] assert headers["X-Trace"] == "keep" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("content_type", "body", "expected_type"), + [ + ("text/html", b"secret-page", ValueError), + ("application/json", b"secret-invalid-json", ValidationError), + ("application/json", b'{"secret":"invalid-rpc"}', ValidationError), + ("application/json", b'{"jsonrpc":"2.0","id":0,"result":{"secret":"invalid-schema"}}', ValidationError), + ], +) +async def test_invalid_http_response_surfaces_without_waiting_for_timeout( + content_type: str, body: bytes, expected_type: type[Exception] +) -> None: + from litellm.proxy._experimental.mcp_server.rest_endpoints import _connection_error_message + + def respond(request: httpx.Request) -> httpx.Response: + return httpx.Response(200, headers={"Content-Type": content_type}, content=body) + + async with httpx.AsyncClient(transport=httpx.MockTransport(respond)) as http_client: + client: Final = MCPClient(server_url="https://example.com/mcp", timeout=30) + with pytest.raises(expected_type) as caught: + await asyncio.wait_for( + client._execute_session_operation( + streamable_http_client(client.server_url, http_client=http_client), + lambda session: session.list_tools(), + ), + timeout=3, + ) + + message: Final = _connection_error_message(caught.value, client.server_url, 30) + assert "unsupported content type" in message or "invalid MCP response" in message + assert "secret" not in message + assert "timed out" not in message + + +@pytest.mark.asyncio +@pytest.mark.parametrize("status_code", [200, 401, 503]) +async def test_http_response_handler_preserves_success_and_http_errors(status_code: int) -> None: + def respond(request: httpx.Request) -> httpx.Response: + if request.method == "DELETE": + return httpx.Response(200) + payload: Final = json.loads(request.content) + if "id" not in payload: + return httpx.Response(202) + result: Final = ( + { + "protocolVersion": LATEST_PROTOCOL_VERSION, + "capabilities": {}, + "serverInfo": {"name": "test", "version": "1"}, + } + if payload["method"] == "initialize" + else {"tools": []} + ) + return httpx.Response(status_code, json={"jsonrpc": "2.0", "id": payload["id"], "result": result}) + + async with httpx.AsyncClient(transport=httpx.MockTransport(respond)) as http_client: + client: Final = MCPClient(server_url="https://example.com/mcp", timeout=30) + operation: Final = client._execute_session_operation( + streamable_http_client(client.server_url, http_client=http_client), lambda session: session.list_tools() + ) + if status_code == 200: + result: Final = await asyncio.wait_for(operation, timeout=3) + assert result.tools == [] + else: + with pytest.raises(httpx.HTTPStatusError) as caught: + await asyncio.wait_for(operation, timeout=3) + assert caught.value.response.status_code == status_code + + +@pytest.mark.asyncio +async def test_invalid_tool_list_schema_is_identified_as_an_upstream_response() -> None: + from litellm.proxy._experimental.mcp_server.rest_endpoints import _connection_error_message + + def respond(request: httpx.Request) -> httpx.Response: + if request.method == "DELETE": + return httpx.Response(200) + payload: Final = json.loads(request.content) + if "id" not in payload: + return httpx.Response(202) + result: Final = ( + { + "protocolVersion": LATEST_PROTOCOL_VERSION, + "capabilities": {}, + "serverInfo": {"name": "test", "version": "1"}, + } + if payload["method"] == "initialize" + else {"tools": "secret-invalid-tools"} + ) + return httpx.Response(200, json={"jsonrpc": "2.0", "id": payload["id"], "result": result}) + + async with httpx.AsyncClient(transport=httpx.MockTransport(respond)) as http_client: + client: Final = MCPClient(server_url="https://example.com/mcp", timeout=30) + with pytest.raises(ValidationError) as caught: + await asyncio.wait_for( + client._execute_session_operation( + streamable_http_client(client.server_url, http_client=http_client), + lambda session: session.list_tools(), + ), + timeout=3, + ) + + message: Final = _connection_error_message(caught.value, client.server_url, 30) + assert "invalid MCP response" in message + assert "secret" not in message diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py index bd692776c82..73734131130 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py @@ -3,7 +3,7 @@ import inspect import json import sys from datetime import datetime -from typing import Any, Dict, Optional +from typing import Any, Dict, Final, Optional from unittest.mock import AsyncMock, MagicMock if sys.version_info < (3, 11): # BaseExceptionGroup is a builtin only from 3.11 @@ -113,7 +113,7 @@ class TestExecuteWithMcpClient: assert "stack_trace" not in result @pytest.mark.asyncio - async def test_timeout_caps_hanging_operation_and_names_url(self, monkeypatch): + async def test_timeout_caps_hanging_operation_and_names_origin(self, monkeypatch): async def fake_create_client(*args, **kwargs): return object() @@ -138,7 +138,7 @@ class TestExecuteWithMcpClient: ) assert result["error"] is True - assert "https://mcp.example.com/mcp/" in result["message"] + assert "https://mcp.example.com" in result["message"] @pytest.mark.asyncio async def test_timeout_covers_client_creation(self, monkeypatch): @@ -166,15 +166,15 @@ class TestExecuteWithMcpClient: ) assert result["error"] is True - assert "https://mcp.example.com/mcp/" in result["message"] + assert "https://mcp.example.com" in result["message"] def test_timeout_defaults_to_tool_listing_timeout(self): default = inspect.signature(rest_endpoints._execute_with_mcp_client).parameters["timeout_seconds"].default assert default == MCP_TOOL_LISTING_TIMEOUT - def test_connection_error_message_timeout_names_url_and_budget(self): + def test_connection_error_message_timeout_names_origin_and_budget(self): message = rest_endpoints._connection_error_message(TimeoutError(), "https://api.example.com/mcp/", 30.0) - assert "https://api.example.com/mcp/" in message + assert "https://api.example.com" in message assert "30s" in message def test_connection_error_message_hides_arbitrary_http_exception_detail(self): @@ -592,7 +592,7 @@ class TestExecuteWithMcpClient: assert result["status"] == "error" assert result["error"] is True - assert "Failed to connect to MCP server" in result["message"] + assert "reference" in result["message"] # Error message must not leak raw exception details assert "cancel scope" not in result["message"] @@ -3430,7 +3430,139 @@ class TestConnectionErrorMessage: def test_unknown_error_falls_back_to_generic(self): message = rest_endpoints._connection_error_message(RuntimeError("weird"), "https://example.com", 30.0) assert "weird" not in message - assert "proxy logs" in message.lower() + assert "reference" in message.lower() + + def test_sdk_session_terminated_explains_endpoint_and_retry(self) -> None: + from mcp.shared.exceptions import McpError + from mcp.types import ErrorData + + message: Final = rest_endpoints._connection_error_message( + McpError(ErrorData(code=32600, message="Session terminated")), "https://example.com/mcp", 30.0 + ) + + assert "session was terminated" in message + assert "MCP endpoint" in message + assert "transport" in message + assert "retry" in message + assert "404" not in message + + @pytest.mark.parametrize("code", [-32700, -32601, -32602, -32603, -32000, 32600, 408]) + def test_rpc_errors_include_code_without_echoing_upstream_data(self, code: int) -> None: + from mcp.shared.exceptions import McpError + from mcp.types import ErrorData + + message: Final = rest_endpoints._connection_error_message( + McpError(ErrorData(code=code, message="secret-message", data={"token": "secret-data"})), + "https://example.com/secret-path?token=secret-query", + 30.0, + ) + + assert f"JSON-RPC code {code}" in message + assert "secret" not in message + assert "timed out" not in message + assert "session was terminated" not in message + + @pytest.mark.parametrize("status_code", [401, 403, 404, 405, 429, 503]) + def test_wrapped_http_failures_preserve_status(self, status_code: int) -> None: + response: Final = httpx.Response(status_code, text="secret-body") + upstream: Final = httpx.HTTPStatusError( + "secret-exception", + request=httpx.Request("POST", "https://example.com/?token=secret-query"), + response=response, + ) + wrapped: Final = BaseExceptionGroup( + "secret-group", [asyncio.CancelledError(), BaseExceptionGroup("nested", [upstream])] + ) + + message: Final = rest_endpoints._connection_error_message(wrapped, "https://example.com", 30.0) + + assert f"HTTP {status_code}" in message + assert "secret" not in message + + def test_explicit_cause_is_classified_before_incidental_context(self) -> None: + wrapped: Final = RuntimeError("secret-wrapper") + wrapped.__cause__ = httpx.ConnectError("secret-cause") + wrapped.__context__ = TimeoutError("secret-context") + + message: Final = rest_endpoints._connection_error_message(wrapped, "https://example.com", 30.0) + + assert "unreachable" in message + assert "secret" not in message + + def test_timeout_url_redacts_credentials_path_query_and_fragment(self) -> None: + message: Final = rest_endpoints._connection_error_message( + TimeoutError("secret-error"), + "https://secret-user:secret-pass@example.com:8443/secret-path?token=secret-query#secret-fragment", + 30.0, + ) + + assert "https://example.com:8443" in message + assert "30s" in message + assert "secret" not in message + + def test_unknown_failure_reference_matches_safe_diagnostics(self, caplog: pytest.LogCaptureFixture) -> None: + import re + + try: + raise RuntimeError("secret-exception-body") + except RuntimeError as exc: + message: Final = rest_endpoints._connection_error_message( + exc, "https://secret-user:secret-password@example.com/secret-path?token=secret-query", 30.0 + ) + + reference: Final = re.search(r"reference ([a-f0-9]{32})", message) + assert reference is not None + diagnostics: Final = tuple( + record for record in caplog.records if "MCP connection test failed" in record.message + ) + assert len(diagnostics) == 1 + assert reference.group(1) in diagnostics[0].message + assert "RuntimeError" in diagnostics[0].message + assert "test_unknown_failure_reference_matches_safe_diagnostics" in diagnostics[0].message + assert diagnostics[0].exc_info is None + assert "secret" not in message + diagnostics[0].message + + @pytest.mark.parametrize("exc", [ValueError("secret-config"), HTTPException(500, "secret-detail")]) + def test_unrelated_errors_are_not_misreported_as_invalid_mcp(self, exc: Exception) -> None: + message: Final = rest_endpoints._connection_error_message(exc, "https://example.com", 30.0) + + assert "reference" in message + assert "invalid MCP response" not in message + assert "secret" not in message + + def test_configuration_validation_error_uses_unknown_fallback(self) -> None: + from pydantic import ValidationError + + with pytest.raises(ValidationError) as caught: + NewMCPServerRequest.model_validate({"server_name": "example", "transport": "secret-invalid-transport"}) + + message: Final = rest_endpoints._connection_error_message(caught.value, "https://example.com", 30.0) + assert "reference" in message + assert "invalid MCP response" not in message + assert "secret" not in message + + @pytest.mark.asyncio + async def test_connection_test_preserves_cancellation(self) -> None: + async def cancelled_operation(client: rest_endpoints.MCPClient) -> dict[str, object]: + raise asyncio.CancelledError + + payload: Final = NewMCPServerRequest(server_name="cancelled", url="https://example.com", auth_type=MCPAuth.none) + with pytest.raises(asyncio.CancelledError): + await rest_endpoints._execute_with_mcp_client(payload, cancelled_operation) + + @pytest.mark.asyncio + async def test_unknown_failure_preserves_response_contract(self) -> None: + async def failing_operation(client: rest_endpoints.MCPClient) -> dict[str, object]: + raise RuntimeError("secret-operation") + + payload: Final = NewMCPServerRequest(server_name="unknown", url="https://example.com", auth_type=MCPAuth.none) + result: Final = await rest_endpoints._execute_with_mcp_client(payload, failing_operation) + + assert result["error"] is True + assert result["status"] == "error" + assert "reference" in result["message"] + assert "secret" not in result["message"] + assert "stack_trace" not in result class TestGetServerAuthHeaderGroupDefault: