From 9f9a02b646f27a7c225715bb2c0a5ecc0bf8fb4c Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 16 Oct 2026 20:58:26 -0700 Subject: [PATCH] fix(mcp): log upstream exchange when an MCP health check fails (#44125) Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: jesus --- litellm/experimental_mcp_client/client.py | 14 +++- .../mcp_server/mcp_server_manager.py | 14 +++- .../_experimental/mcp_server/upstream.py | 3 + .../mcp_server/test_mcp_server_manager.py | 67 +++++++++++++++++++ 4 files changed, 95 insertions(+), 3 deletions(-) diff --git a/litellm/experimental_mcp_client/client.py b/litellm/experimental_mcp_client/client.py index e2057e4b627..ed6d89397a0 100644 --- a/litellm/experimental_mcp_client/client.py +++ b/litellm/experimental_mcp_client/client.py @@ -498,6 +498,7 @@ class MCPClient: elicitation_callback: "ElicitationCallback | None" = None, logging_callback: Callable | None = None, protocol_version: MCPUpstreamProtocol = "auto", + upstream_transport: httpx2.AsyncBaseTransport | None = None, ): self.protocol_version: MCPUpstreamProtocol = TypeAdapter[MCPUpstreamProtocol]( MCPUpstreamProtocol @@ -524,6 +525,7 @@ class MCPClient: self._sampling_callback: Callable | None = sampling_callback self._elicitation_callback: ElicitationCallback | None = elicitation_callback self._logging_callback: Callable | None = logging_callback + self._upstream_transport: httpx2.AsyncBaseTransport | None = upstream_transport self._session_callbacks_override: _SessionCallbackMap | None = None # handle the basic auth value if provided if auth_value: @@ -610,7 +612,11 @@ class MCPClient: return stdio_client(server_params), None if self.transport_type == MCPTransport.sse: headers = self._get_auth_headers() - httpx_client_factory = self._create_httpx_client_factory() + httpx_client_factory = ( + self._create_httpx_client_factory() + if self._upstream_transport is None + else self._create_httpx_client_factory(transport=self._upstream_transport) + ) return ( sse_client( url=self.server_url, @@ -622,7 +628,11 @@ class MCPClient: ) # HTTP transport (default) headers = self._get_auth_headers() - httpx_client_factory = self._create_httpx_client_factory() + httpx_client_factory = ( + self._create_httpx_client_factory() + if self._upstream_transport is None + else self._create_httpx_client_factory(transport=self._upstream_transport) + ) verbose_logger.debug("litellm headers for streamable_http_client: %s", headers) http_client = httpx_client_factory( headers=headers, diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index b724a61cd94..02799b90566 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -1970,6 +1970,7 @@ class MCPServerManager: per_user_oauth_token_store: InvalidatableOAuthTokenStore | None = None, per_user_token_cache: MCPPerUserTokenCache | None = None, discovery_clock: Callable[[], float] = time.monotonic, + upstream_transport: httpx2.AsyncBaseTransport | None = None, ): self._per_user_oauth_token_store = per_user_oauth_token_store or LazyPerUserOAuthTokenStore( self.get_mcp_server_by_id @@ -1979,6 +1980,7 @@ class MCPServerManager: oauth_token_store=self._per_user_oauth_token_store, token_exchanger=build_token_exchanger(), ) + self._upstream_transport: httpx2.AsyncBaseTransport | None = upstream_transport discovery_ttl: Final = _mcp_discovery_cache_ttl() self._prompt_discovery_cache = _DiscoveryCache[ListPromptsResult]( discovery_ttl, discovery_clock, TypeAdapter(ListPromptsResult) @@ -4348,6 +4350,7 @@ class MCPServerManager: if resolved_server.allow_elicitation else None ), + upstream_transport=self._upstream_transport, ) _create_mcp_client = create_mcp_client @@ -7466,7 +7469,10 @@ class MCPServerManager: return "ok" # Add timeout wrapper to prevent hanging - await asyncio.wait_for(client.run_with_session(_noop), timeout=MCP_HEALTH_CHECK_TIMEOUT) + await asyncio.wait_for( + client.run_with_session(_noop, quiet_on_error=True), + timeout=MCP_HEALTH_CHECK_TIMEOUT, + ) self._remember_upstream_initialize_instructions(server, client) status = "healthy" except asyncio.TimeoutError: @@ -7478,6 +7484,12 @@ class MCPServerManager: except Exception as e: health_check_error = str(e) status = "unhealthy" + verbose_logger.warning( + "MCP health check failed for %s: %s%s", + server.name, + type(e).__name__, + _upstream_failure_suffix(e), + ) return LiteLLM_MCPServerTable( server_id=server.server_id, diff --git a/litellm/proxy/_experimental/mcp_server/upstream.py b/litellm/proxy/_experimental/mcp_server/upstream.py index b19f904d81d..3b7cf904b6a 100644 --- a/litellm/proxy/_experimental/mcp_server/upstream.py +++ b/litellm/proxy/_experimental/mcp_server/upstream.py @@ -207,6 +207,7 @@ async def prepare_upstream_client( protocol_version: MCPUpstreamProtocol, sampling_callback: SamplingCallback | None = None, elicitation_callback: ElicitationCallback | None = None, + upstream_transport: httpx2.AsyncBaseTransport | None = None, ) -> MCPClient: transport: Final[MCPTransportType] = server.transport or MCPTransport.sse server_spec: Final = None if transport == MCPTransport.stdio else to_server_spec_fail_closed(server) @@ -263,6 +264,7 @@ async def prepare_upstream_client( server_url=server.url or "", transport_type=transport, protocol_version=protocol_version, + upstream_transport=upstream_transport, auth_type=server.auth_type, timeout=timeout, extra_headers=resolved_headers, @@ -304,6 +306,7 @@ async def prepare_upstream_client( server_url=server.url or "", transport_type=transport, protocol_version=protocol_version, + upstream_transport=upstream_transport, auth_type=server.auth_type, auth_value=auth_value, auth_header_name=auth_header_name, diff --git a/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py index 40901d1dd96..691da35b8c4 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -15634,6 +15634,73 @@ def _mcp_upstream(respond): yield +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("responder", "expects_exchange", "expected_fragments"), + ( + pytest.param( + AsyncMock(return_value=httpx2.Response(500, json={"error": "health-check-boom"})), + True, + ( + "MCP health check failed for", + "upstream exchange:", + "-> HTTP 500", + '"method":"initialize"', + "health-check-boom", + ), + id="upstream-http-500", + ), + pytest.param( + AsyncMock(side_effect=httpx2.ConnectError("connection failed")), + False, + ("MCP health check failed for",), + id="connect-error", + ), + ), +) +async def test_health_check_logs_upstream_exchange_only_for_http_failures( + caplog: pytest.LogCaptureFixture, + responder: AsyncMock, + expects_exchange: bool, + expected_fragments: tuple[str, ...], +) -> None: + manager: Final = MCPServerManager(upstream_transport=httpx2.MockTransport(responder)) + token: Final = "health-check-secret-0123456789" + server: Final = MCPServer( + server_id="health-check-upstream", + name="health_check_upstream", + url="https://upstream.example/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.bearer_token, + authentication_token=token, + ) + manager.registry[server.server_id] = server + + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + result: Final = await manager.health_check_server(server.server_id) + + assert result.status == "unhealthy" + warning_messages: Final = tuple( + record.getMessage() + for record in caplog.records + if record.name == "LiteLLM" + and record.levelno == logging.WARNING + and "MCP health check failed for" in record.getMessage() + ) + assert len(warning_messages) == 1 + warning: Final = warning_messages[0] + assert all(fragment in warning for fragment in expected_fragments) + assert ("upstream exchange:" in warning) is expects_exchange + assert not any( + record.name == "LiteLLM" + and record.levelno == logging.WARNING + and "run_with_session failed" in record.getMessage() + for record in caplog.records + ) + assert token not in caplog.text + assert responder.await_args_list[0].args[0].headers["Authorization"] == f"Bearer {token}" + + class _DiscoveryUpstream: def __init__(self) -> None: self.requests: tuple[tuple[str, str], ...] = ()