fix(mcp): preserve stream failures across transports

litellm/experimental_mcp_client/client.py

@@ -451,7 +451,7 @@ class MCPClient:
             async def receive_message(
                 message: RequestResponder[ServerRequest, ClientResult] | ServerNotification | Exception,
             ) -> None:
-                if not isinstance(message, ValueError):
+                if not isinstance(message, (ValueError, httpx.RequestError, OSError)):
                     return
                 if not stream_error.done():
                     stream_error.set_result(message)
@@ -472,7 +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,
+                message_handler=receive_message,
                 **session_kwargs,
             )
             session: Final = await session_ctx.__aenter__()
@@ -525,8 +525,7 @@ class MCPClient:
             read_timeout: Final = _as_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",
                 )

litellm/proxy/_experimental/mcp_server/rest_endpoints.py

@@ -114,7 +114,8 @@ def _known_connection_error_message(exc: BaseException, url: str | None, timeout
         return str(exc.detail)
     if isinstance(exc, TimeoutError):
         return (
-            f"Failed to connect to MCP server: no response from {_redact_mcp_resource_url(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."
         )
@@ -131,6 +132,11 @@ def _known_connection_error_message(exc: BaseException, url: str | None, timeout
         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}."
+    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. "
@@ -142,6 +148,11 @@ def _known_connection_error_message(exc: BaseException, url: str | None, timeout
             "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. "
@@ -159,7 +170,7 @@ 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_read_timeout
     from litellm.llms.litellm_proxy.skills.skill_search import (
         DEFAULT_SKILL_SEARCH_TOP_K,
     )
@@ -1396,10 +1407,18 @@ if MCP_AVAILABLE:
         except (KeyboardInterrupt, SystemExit, asyncio.CancelledError):
             raise
         except BaseException as e:
+            effective_timeout: Final = (
+                min(request.timeout or MCP_CLIENT_TIMEOUT, timeout_seconds)
+                if any(
+                    isinstance(cause, McpError) and _as_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:

tests/test_litellm/experimental_mcp_client/test_mcp_client.py

@@ -3,6 +3,7 @@ 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
@@ -18,6 +19,7 @@ from pydantic import ValidationError
 from mcp.shared.message import SessionMessage
 from mcp.types import (
     LATEST_PROTOCOL_VERSION,
+    CallToolResult,
     ErrorData,
     Implementation,
     InitializeResult,
@@ -36,6 +38,7 @@ from litellm.experimental_mcp_client.client import (
     MCPClient,
     _as_read_timeout,
     _first_non_cancelled_cause,
+    _TransportContext,
     missing_streamable_http_client_error,
     strip_auth_scheme,
 )
@@ -1296,6 +1299,7 @@ def test_a_differently_cased_injected_header_cannot_shadow_the_slot() -> None:
     [
         ("text/html", b"<html>secret-page</html>", 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),
     ],
@@ -1446,3 +1450,239 @@ async def test_invalid_tool_list_schema_is_identified_as_an_upstream_response()
     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_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_read_timeout(caught.value), TimeoutError)

tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py

@@ -3427,6 +3427,54 @@ 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])
+    async def test_timeout_message_uses_the_deadline_that_expired(self, sdk_timeout: bool) -> 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=1
+        )
+        result: Final = await rest_endpoints._execute_with_mcp_client(payload, operation, timeout_seconds=30)
+        assert ("within 1s" 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