From 5447b949bfdbf4a51bd3e88b211b665082f87d87 Mon Sep 17 00:00:00 2001 From: yassin Date: Wed, 12 Aug 2026 06:04:38 +0000 Subject: [PATCH] fix(mcp): bound MCP client requests with a session read timeout An upstream streamable-HTTP response stream that ends without a JSON-RPC reply is dropped by the MCP SDK, so the request stays pending forever. Tool discovery then only ended when the outer listing cancel scope killed it, which logged a cancelled list_tools, ignored the server's own timeout, and reported no tools to the client; prompts and resources had no outer guard at all. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/experimental_mcp_client/client.py | 10 ++- .../mcp_server/faults/list_outcomes.py | 3 + .../test_mcp_client.py | 78 +++++++++++++++++++ .../mcp_server/faults/test_list_outcomes.py | 15 ++++ 4 files changed, 105 insertions(+), 1 deletion(-) diff --git a/litellm/experimental_mcp_client/client.py b/litellm/experimental_mcp_client/client.py index d474291f1cb..64eb96d9318 100644 --- a/litellm/experimental_mcp_client/client.py +++ b/litellm/experimental_mcp_client/client.py @@ -6,6 +6,7 @@ import asyncio import base64 import os from collections.abc import Awaitable, Callable, Generator +from datetime import timedelta from typing import Any, Final, TypeVar import httpx @@ -347,7 +348,14 @@ class MCPClient: session_kwargs["elicitation_callback"] = self._elicitation_callback if self._logging_callback is not None: session_kwargs["logging_callback"] = self._logging_callback - session_ctx: Final = ClientSession(read_stream, write_stream, **session_kwargs) + # A streamable-HTTP response stream that ends without a JSON-RPC reply is dropped by the + # SDK, leaving the request pending forever; the read timeout is what bounds it. + session_ctx: Final = ClientSession( + read_stream, + write_stream, + read_timeout_seconds=timedelta(seconds=self.timeout), + **session_kwargs, + ) session: Final = await session_ctx.__aenter__() try: init_result: Final = await session.initialize() diff --git a/litellm/proxy/_experimental/mcp_server/faults/list_outcomes.py b/litellm/proxy/_experimental/mcp_server/faults/list_outcomes.py index 42b2d29cd52..3ecbada03ad 100644 --- a/litellm/proxy/_experimental/mcp_server/faults/list_outcomes.py +++ b/litellm/proxy/_experimental/mcp_server/faults/list_outcomes.py @@ -14,6 +14,7 @@ from collections.abc import Iterator from typing import Final, Literal, NamedTuple, NoReturn, TypeAlias import httpx +from mcp import McpError from mcp.types import Tool as MCPTool from pydantic import BaseModel, ConfigDict from typing_extensions import assert_never @@ -124,6 +125,8 @@ def classify_list_exception(exc: BaseException) -> ServerListFault: return ServerListFault(tag=tag, status_code=exc.status_code) if isinstance(exc, TimeoutError): return ServerListFault(tag="timeout") + if isinstance(exc, McpError) and exc.error.code == httpx.codes.REQUEST_TIMEOUT: + return ServerListFault(tag="timeout") if isinstance(exc, ConnectionError): return ServerListFault(tag="unreachable") auth: Final = upstream_auth_challenge(exc) 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 8e6fa35b452..60aebaab39b 100644 --- a/tests/test_litellm/experimental_mcp_client/test_mcp_client.py +++ b/tests/test_litellm/experimental_mcp_client/test_mcp_client.py @@ -3,8 +3,19 @@ import os import sys from unittest.mock import AsyncMock, MagicMock, patch +import anyio import httpx import pytest +from mcp import McpError +from mcp.shared.message import SessionMessage +from mcp.types import ( + LATEST_PROTOCOL_VERSION, + Implementation, + InitializeResult, + JSONRPCMessage, + JSONRPCResponse, + ServerCapabilities, +) # Add the parent directory to the path so we can import litellm sys.path.insert(0, "../../../") @@ -701,3 +712,70 @@ async def test_run_with_session_quiet_on_error_demotes_warning_to_debug(): assert any("run_with_session failed" in m for m in warning_msgs), ( "the default path must keep the operator-visible warning" ) + + +class _DroppedResponseTransport: + """An upstream that answers ``initialize`` and then never answers anything else. + + This is the shape the MCP SDK leaves behind when a streamable-HTTP response stream ends without a + JSON-RPC reply: it logs "SSE stream ended" at debug and returns, so the pending request is never + resolved and never fails. + """ + + def __init__(self) -> None: + self._to_client_tx, self._to_client_rx = anyio.create_memory_object_stream(10) + self._from_client_tx, self._from_client_rx = anyio.create_memory_object_stream(10) + self._task_group = None + + async def __aenter__(self): + self._task_group = anyio.create_task_group() + await self._task_group.__aenter__() + self._task_group.start_soon(self._serve) + return self._to_client_rx, self._from_client_tx + + async def __aexit__(self, *exc_info): + self._task_group.cancel_scope.cancel() + return await self._task_group.__aexit__(None, None, None) + + async def _serve(self) -> None: + async for session_message in self._from_client_rx: + request = session_message.message.root + if getattr(request, "method", None) != "initialize": + continue + result = InitializeResult( + protocolVersion=LATEST_PROTOCOL_VERSION, + capabilities=ServerCapabilities(), + serverInfo=Implementation(name="dropped-response-upstream", version="1.0.0"), + ) + await self._to_client_tx.send( + SessionMessage( + JSONRPCMessage( + JSONRPCResponse( + jsonrpc="2.0", + id=request.id, + result=result.model_dump(by_alias=True, mode="json", exclude_none=True), + ) + ) + ) + ) + + +class _DroppedResponseClient(MCPClient): + def _create_transport_context(self): + return _DroppedResponseTransport(), None + + +@pytest.mark.asyncio +async def test_list_tools_times_out_instead_of_hanging_when_upstream_drops_the_response(): + """An upstream that accepts the request and never answers must fail the client's own timeout. + + Without a session read timeout the request waits forever, so tool discovery only ends when an outer + cancel scope kills it. That surfaces as a cancelled list_tools with no tools and no truthful fault, + and it wedges the paths that have no outer guard at all (prompts, resources). + """ + client = _DroppedResponseClient(server_url="http://upstream.local/mcp", timeout=0.5) + + with pytest.raises(McpError) as exc_info: + await asyncio.wait_for(client.list_tools(raise_on_error=True), timeout=10) + + assert exc_info.value.error.code == httpx.codes.REQUEST_TIMEOUT diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/faults/test_list_outcomes.py b/tests/test_litellm/proxy/_experimental/mcp_server/faults/test_list_outcomes.py index cb27e992ecb..81624592d84 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/faults/test_list_outcomes.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/faults/test_list_outcomes.py @@ -4,6 +4,8 @@ stay truthful to who failed.""" import httpx import pytest +from mcp import McpError +from mcp.types import INTERNAL_ERROR, ErrorData from litellm.proxy._experimental.mcp_server.exceptions import ( MCPServerListError, @@ -33,6 +35,19 @@ def test_timeout_and_connection_errors_classify_without_status(): assert classify_list_exception(ConnectionError()).tag == "unreachable" +def test_sdk_request_timeout_classifies_as_timeout_not_internal(): + """The MCP SDK reports an unanswered request as McpError(408), the shape an upstream that drops + the response stream produces; classifying it as internal would blame the gateway.""" + exc = McpError(ErrorData(code=httpx.codes.REQUEST_TIMEOUT, message="Timed out while waiting for response")) + assert classify_list_exception(exc).tag == "timeout" + assert list_fault_http_status(classify_list_exception(exc)) == 504 + + +def test_other_sdk_errors_still_classify_as_internal(): + exc = McpError(ErrorData(code=INTERNAL_ERROR, message="boom")) + assert classify_list_exception(exc).tag == "internal" + + def test_embedded_upstream_response_status_wins(): response = httpx.Response(503, request=httpx.Request("POST", "https://mcp.example.com/mcp")) exc = httpx.HTTPStatusError("boom", request=response.request, response=response)