diff --git a/litellm/experimental_mcp_client/client.py b/litellm/experimental_mcp_client/client.py index 3503468c735..8c7f4557992 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 @@ -146,8 +150,8 @@ _SDK_READ_TIMEOUT_CODE: Final = int(httpx.codes.REQUEST_TIMEOUT) otherwise carries JSON-RPC error codes.""" -def _as_read_timeout(exc: BaseException) -> TimeoutError | None: - """The session read timeout elapsing, re-expressed as a ``TimeoutError``, or ``None``. +def as_mcp_read_timeout(exc: BaseException) -> TimeoutError | None: + """Normalize an MCP SDK read timeout for client and gateway diagnostics, or return ``None``. The SDK reports its own elapsed read timeout as ``McpError`` carrying an HTTP status code in a field that otherwise holds JSON-RPC error codes, and it relays an upstream's JSON-RPC error @@ -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, httpx.RequestError, OSError)): + 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, **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) @@ -501,11 +522,10 @@ class MCPClient: transport_ctx, http_client = self._create_transport_context() return await self._execute_session_operation(transport_ctx, operation) except Exception as e: - read_timeout: Final = _as_read_timeout(e) + read_timeout: Final = as_mcp_read_timeout(e) if read_timeout is not None: verbose_logger.warning( - "MCP client timed out after %ss waiting for %s to answer; the server accepted the " - "request and ended its response stream without a JSON-RPC reply", + "MCP client timed out after %ss waiting for a valid MCP response from %s", self.timeout, self.server_url or "stdio", ) diff --git a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py index 329dddbdf05..5b74caee3a2 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,39 @@ _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'} " + "Failed to connect to MCP server: no valid MCP response received from " + f"{_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,13 +132,45 @@ 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, (httpx.NetworkError, httpx.RemoteProtocolError, ConnectionError)): + return ( + "Failed to connect to MCP server: the connection was interrupted. " + "Check the server and network connection, then retry." + ) + 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 == -32000 and exc.error.message == "Connection closed": + return ( + "Failed to connect to MCP server: the connection was closed before the request completed. " + "Check that the server stays running and returns a complete MCP response, then retry." + ) + 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 + from litellm.experimental_mcp_client.client import MCPClient, as_mcp_read_timeout from litellm.llms.litellm_proxy.skills.skill_search import ( DEFAULT_SKILL_SEARCH_TOP_K, ) @@ -1342,11 +1407,18 @@ 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) + effective_timeout: Final = ( + min(request.timeout if request.timeout is not None else MCP_CLIENT_TIMEOUT, timeout_seconds) + if any( + isinstance(cause, McpError) and as_mcp_read_timeout(cause) is not None + for cause in iter_exception_tree(e) + ) + else timeout_seconds + ) return { "status": "error", "error": True, - "message": _connection_error_message(e, request.url, timeout_seconds), + "message": _connection_error_message(e, request.url, effective_timeout), } async def _preview_openapi_tools(spec_path: str) -> dict: 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..b07e5876e8b 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,12 @@ import asyncio import base64 +import json import os import sys +from collections.abc import AsyncIterator from importlib import metadata from pathlib import Path +from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import anyio @@ -11,15 +14,19 @@ 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, + CallToolResult, ErrorData, Implementation, InitializeResult, JSONRPCError, JSONRPCMessage, JSONRPCResponse, + LoggingMessageNotificationParams, ServerCapabilities, ) @@ -29,8 +36,9 @@ import litellm.experimental_mcp_client.client as mcp_client_module from litellm.experimental_mcp_client.client import ( MCP_STREAMABLE_HTTP_REQUIREMENT, MCPClient, - _as_read_timeout, _first_non_cancelled_cause, + _TransportContext, + as_mcp_read_timeout, missing_streamable_http_client_error, strip_auth_scheme, ) @@ -859,25 +867,25 @@ def _raise_mcp_error_while_handling_a_timeout(code: int, message: str) -> McpErr return raised -def test_as_read_timeout_separates_the_sdk_timeout_from_a_relayed_upstream_error(): +def test_as_mcp_read_timeout_separates_the_sdk_timeout_from_a_relayed_upstream_error(): """Neither signal alone is enough. The code alone cannot separate the SDK's own timeout from an upstream JSON-RPC error that happens to use 408, and the context chain alone cannot separate it from any other relayed error that surfaces while a timeout is being handled, so both must hold. """ timeout_code = int(httpx.codes.REQUEST_TIMEOUT) - translated = _as_read_timeout(_raise_mcp_error_while_handling_a_timeout(timeout_code, "Timed out while waiting")) + translated = as_mcp_read_timeout(_raise_mcp_error_while_handling_a_timeout(timeout_code, "Timed out while waiting")) assert isinstance(translated, TimeoutError) assert str(translated) == "Timed out while waiting" relayed_408 = McpError(ErrorData(code=timeout_code, message="upstream said 408")) - assert _as_read_timeout(relayed_408) is None, "an upstream 408 with no elapsed timeout is not our timeout" + assert as_mcp_read_timeout(relayed_408) is None, "an upstream 408 with no elapsed timeout is not our timeout" relayed_other = _raise_mcp_error_while_handling_a_timeout(-32603, "upstream internal error") - assert _as_read_timeout(relayed_other) is None, "a non-timeout code is not our timeout, whatever the chain" + assert as_mcp_read_timeout(relayed_other) is None, "a non-timeout code is not our timeout, whatever the chain" - assert _as_read_timeout(McpError(ErrorData(code=-32603, message="boom"))) is None - assert _as_read_timeout(RuntimeError("not an McpError")) is None + assert as_mcp_read_timeout(McpError(ErrorData(code=-32603, message="boom"))) is None + assert as_mcp_read_timeout(RuntimeError("not an McpError")) is None @pytest.mark.asyncio @@ -1224,14 +1232,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 +1291,398 @@ 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"", 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_http_response_handler_preserves_notifications_and_tool_listing() -> None: + notification: Final = { + "jsonrpc": "2.0", + "method": "notifications/message", + "params": {"level": "info", "data": "Listing tools"}, + } + logging_callback: Final = AsyncMock() + + 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) + if payload["method"] == "initialize": + return httpx.Response( + 200, + json={ + "jsonrpc": "2.0", + "id": payload["id"], + "result": { + "protocolVersion": LATEST_PROTOCOL_VERSION, + "capabilities": {"logging": {}, "tools": {}}, + "serverInfo": {"name": "test", "version": "1"}, + }, + }, + ) + response: Final = { + "jsonrpc": "2.0", + "id": payload["id"], + "result": {"tools": [{"name": "search", "inputSchema": {"type": "object"}}]}, + } + return httpx.Response( + 200, + headers={"Content-Type": "text/event-stream"}, + content="".join(f"event: message\ndata: {json.dumps(message)}\n\n" for message in (notification, response)), + ) + + async with httpx.AsyncClient(transport=httpx.MockTransport(respond)) as http_client: + client: Final = MCPClient(server_url="https://example.com/mcp", timeout=30, logging_callback=logging_callback) + result: Final = 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, + ) + + assert [tool.name for tool in result.tools] == ["search"] + logging_callback.assert_awaited_once_with(LoggingMessageNotificationParams(level="info", data="Listing tools")) + + +@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 + + +class _DiagnosticSSEStream(httpx.AsyncByteStream): + def __init__(self, messages: asyncio.Queue[bytes | Exception | None]) -> None: + self.messages = messages + + async def __aiter__(self) -> AsyncIterator[bytes]: + yield b"event: endpoint\ndata: /messages\n\n" + while True: + message: Final = await self.messages.get() + if message is None: + return + if isinstance(message, Exception): + raise message + yield b"event: message\ndata: " + message + b"\n\n" + + +_DIAGNOSTIC_STDIO_SERVER: Final = """ +import json, sys +mode, failure_method = sys.argv[1:] +for line in sys.stdin: + request = json.loads(line) + if "method" not in request or "id" not in request: + continue + if request["method"] == failure_method: + if mode == "bad-json": + print("secret-invalid-json", flush=True) + continue + if mode == "closed": + sys.exit(0) + if mode == "silent": + print(json.dumps({"jsonrpc": "2.0", "method": "notifications/message", "params": {"level": "info", "data": "Waiting"}}), flush=True) + continue + if request["method"] == "initialize": + result = {"protocolVersion": request["params"]["protocolVersion"], "capabilities": {"tools": {}, "logging": {}}, "serverInfo": {"name": "diagnostic", "version": "1"}} + elif request["method"] == "tools/list": + print(json.dumps({"jsonrpc": "2.0", "method": "notifications/message", "params": {"level": "info", "data": "Listing tools"}}), flush=True) + print(json.dumps({"jsonrpc": "2.0", "id": "unmatched", "result": {}}), flush=True) + print(json.dumps({"jsonrpc": "2.0", "id": "server-ping", "method": "ping"}), flush=True) + result = {"tools": [{"name": "ping", "inputSchema": {"type": "object"}}]} + else: + result = {"content": [{"type": "text", "text": "pong"}], "isError": False} + print(json.dumps({"jsonrpc": "2.0", "id": request["id"], "result": result}), flush=True) +""" + + +def _diagnostic_transport(transport: MCPTransport, mode: str, failure_method: str) -> _TransportContext: + from mcp import StdioServerParameters + from mcp.client.sse import sse_client + from mcp.client.stdio import stdio_client + + if transport == MCPTransport.stdio: + return stdio_client( + StdioServerParameters( + command=sys.executable, args=["-u", "-c", _DIAGNOSTIC_STDIO_SERVER, mode, failure_method] + ) + ) + messages: Final[asyncio.Queue[bytes | Exception | None]] = asyncio.Queue() + + async def respond(request: httpx.Request) -> httpx.Response: + if request.method == "GET": + return httpx.Response( + 200, headers={"Content-Type": "text/event-stream"}, stream=_DiagnosticSSEStream(messages) + ) + payload: Final = json.loads(request.content) + if "method" not in payload or "id" not in payload: + return httpx.Response(202) + if payload["method"] == failure_method and mode != "ok": + if mode == "bad-json": + await messages.put(b"secret-invalid-json") + elif mode == "io-error": + await messages.put(httpx.ReadError("secret-read-error")) + elif mode == "closed": + await messages.put(None) + elif mode == "silent": + await messages.put( + b'{"jsonrpc":"2.0","method":"notifications/message","params":{"level":"info","data":"Waiting"}}' + ) + return httpx.Response(202) + if payload["method"] == "tools/list": + for message in ( + { + "jsonrpc": "2.0", + "method": "notifications/message", + "params": {"level": "info", "data": "Listing tools"}, + }, + {"jsonrpc": "2.0", "id": "unmatched", "result": {}}, + {"jsonrpc": "2.0", "id": "server-ping", "method": "ping"}, + ): + await messages.put(json.dumps(message).encode()) + result: Final = ( + { + "protocolVersion": LATEST_PROTOCOL_VERSION, + "capabilities": {"tools": {}, "logging": {}}, + "serverInfo": {"name": "diagnostic", "version": "1"}, + } + if payload["method"] == "initialize" + else {"tools": [{"name": "ping", "inputSchema": {"type": "object"}}]} + if payload["method"] == "tools/list" + else {"content": [{"type": "text", "text": "pong"}], "isError": False} + ) + await messages.put(json.dumps({"jsonrpc": "2.0", "id": payload["id"], "result": result}).encode()) + return httpx.Response(202) + + def factory( + headers: dict[str, str] | None = None, + timeout: httpx.Timeout | None = None, + auth: httpx.Auth | None = None, + ) -> httpx.AsyncClient: + return httpx.AsyncClient(transport=httpx.MockTransport(respond), headers=headers, timeout=timeout, auth=auth) + + return sse_client("https://example.com/sse", httpx_client_factory=factory) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("transport", [MCPTransport.sse, MCPTransport.stdio]) +@pytest.mark.parametrize("failure_method", ["initialize", "tools/list"]) +async def test_transport_parsing_failure_is_preserved(transport: MCPTransport, failure_method: str) -> None: + client: Final = MCPClient(server_url="https://example.com/sse", transport_type=transport, timeout=0.2) + with pytest.raises(ValidationError): + await asyncio.wait_for( + client._execute_session_operation( + _diagnostic_transport(transport, "bad-json", failure_method), lambda session: session.list_tools() + ), + timeout=3, + ) + + +@pytest.mark.asyncio +async def test_sse_read_failure_is_preserved() -> None: + client: Final = MCPClient(server_url="https://example.com/sse", transport_type=MCPTransport.sse, timeout=0.2) + with pytest.raises(httpx.ReadError, match="secret-read-error"): + await asyncio.wait_for( + client._execute_session_operation( + _diagnostic_transport(MCPTransport.sse, "io-error", "tools/list"), lambda session: session.list_tools() + ), + timeout=3, + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("transport", [MCPTransport.sse, MCPTransport.stdio]) +@pytest.mark.parametrize("mode", ["ok", "closed", "silent"]) +async def test_transport_completion_and_normal_messages(transport: MCPTransport, mode: str) -> None: + from mcp import ClientSession + from litellm.proxy._experimental.mcp_server.rest_endpoints import _connection_error_message + + logging_callback: Final = AsyncMock() + client: Final = MCPClient( + server_url="https://example.com/sse", transport_type=transport, timeout=0.2, logging_callback=logging_callback + ) + + async def operation(session: ClientSession) -> CallToolResult: + tools: Final = await session.list_tools() + assert [tool.name for tool in tools.tools] == ["ping"] + return await session.call_tool("ping", {}) + + pending: Final = client._execute_session_operation(_diagnostic_transport(transport, mode, "tools/list"), operation) + if mode == "ok": + result: Final = await asyncio.wait_for(pending, timeout=3) + assert result.isError is False + assert result.content[0].text == "pong" + logging_callback.assert_awaited_once_with(LoggingMessageNotificationParams(level="info", data="Listing tools")) + else: + with pytest.raises(McpError) as caught: + await asyncio.wait_for(pending, timeout=3) + if mode == "closed": + assert "connection was closed" in _connection_error_message(caught.value, client.server_url, 0.2) + else: + assert isinstance(as_mcp_read_timeout(caught.value), TimeoutError) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("transport", [MCPTransport.sse, MCPTransport.stdio]) +async def test_transport_cancellation_cleans_up_a_pending_request(transport: MCPTransport) -> None: + ready: Final = asyncio.Event() + + async def on_log(message: LoggingMessageNotificationParams) -> None: + if message.data == "Waiting": + ready.set() + + client: Final = MCPClient( + server_url="https://example.com/sse", transport_type=transport, timeout=30, logging_callback=on_log + ) + task: Final = asyncio.create_task( + client._execute_session_operation( + _diagnostic_transport(transport, "silent", "tools/list"), lambda session: session.list_tools() + ) + ) + try: + await asyncio.wait_for(ready.wait(), timeout=3) + finally: + task.cancel() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(task, timeout=3) + + +class _InterruptedHTTPBody(httpx.AsyncByteStream): + async def __aiter__(self) -> AsyncIterator[bytes]: + yield b'{"jsonrpc":' + raise httpx.RemoteProtocolError("secret-incomplete-response") + + +@pytest.mark.asyncio +async def test_interrupted_http_response_preserves_the_transport_failure() -> None: + def respond(request: httpx.Request) -> httpx.Response: + return httpx.Response(200, headers={"Content-Type": "application/json"}, stream=_InterruptedHTTPBody()) + + 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(httpx.RemoteProtocolError, match="secret-incomplete-response"): + 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, + ) + + +@pytest.mark.asyncio +async def test_empty_http_event_stream_uses_the_existing_request_deadline() -> None: + def respond(request: httpx.Request) -> httpx.Response: + return httpx.Response(200, headers={"Content-Type": "text/event-stream"}, content=b"") + + async with httpx.AsyncClient(transport=httpx.MockTransport(respond)) as http_client: + client: Final = MCPClient(server_url="https://example.com/mcp", timeout=0.2) + with pytest.raises(McpError) 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, + ) + assert isinstance(as_mcp_read_timeout(caught.value), TimeoutError) 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..3487f634251 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"] @@ -3427,10 +3427,191 @@ class TestConnectionErrorMessage: message = rest_endpoints._connection_error_message(exc, "https://example.com", 30.0) assert "503" in message + @pytest.mark.parametrize( + "error_type", [httpx.ReadError, httpx.WriteError, httpx.RemoteProtocolError, ConnectionResetError] + ) + def test_interrupted_connection_message_is_safe(self, error_type: type[Exception]) -> None: + message: Final = rest_endpoints._connection_error_message( + error_type("secret-transport-detail"), "https://example.com/?token=secret-query", 30 + ) + assert "connection was interrupted" in message + assert "secret" not in message + + def test_closed_connection_explains_incomplete_request(self) -> None: + from mcp import McpError + from mcp.types import ErrorData + + message: Final = rest_endpoints._connection_error_message( + McpError(ErrorData(code=-32000, message="Connection closed", data="secret-data")), None, 30 + ) + assert "connection was closed before the request completed" in message + assert "secret" not in message + + def test_timeout_does_not_claim_the_server_sent_nothing(self) -> None: + message: Final = rest_endpoints._connection_error_message(TimeoutError(), None, 30) + assert "no valid MCP response received" in message + + @pytest.mark.asyncio + @pytest.mark.parametrize("sdk_timeout", [True, False]) + @pytest.mark.parametrize("read_timeout", [0, 1]) + async def test_timeout_message_uses_the_deadline_that_expired(self, sdk_timeout: bool, read_timeout: int) -> None: + from mcp import McpError + from mcp.types import ErrorData + + async def operation(client: rest_endpoints.MCPClient) -> dict[str, object]: + try: + raise TimeoutError("secret-timeout") + except TimeoutError as elapsed: + if not sdk_timeout: + raise + try: + raise McpError(ErrorData(code=408, message="secret-sdk-timeout")) from elapsed + except McpError as sdk_error: + raise TimeoutError() from sdk_error + + payload: Final = NewMCPServerRequest( + server_name="timeout", url="https://example.com", auth_type=MCPAuth.none, timeout=read_timeout + ) + result: Final = await rest_endpoints._execute_with_mcp_client(payload, operation, timeout_seconds=30) + assert (f"within {read_timeout}s" if sdk_timeout else "within 30s") in result["message"] + assert "secret" not in result["message"] + 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: