fix(mcp): surface connection test failures safely

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)

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,

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"<html>secret-page</html>", 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

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: