From dbf94902294cf5fa4fe5af6681fcd83ffd743231 Mon Sep 17 00:00:00 2001 From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Date: Wed, 9 Sep 2026 09:44:14 -0700 Subject: [PATCH] fix(mcp): expose shared SDK timeout normalization --- litellm/experimental_mcp_client/client.py | 6 +++--- .../_experimental/mcp_server/rest_endpoints.py | 4 ++-- .../experimental_mcp_client/test_mcp_client.py | 18 +++++++++--------- .../mcp_server/test_rest_endpoints.py | 7 ++++--- 4 files changed, 18 insertions(+), 17 deletions(-) diff --git a/litellm/experimental_mcp_client/client.py b/litellm/experimental_mcp_client/client.py index 5cdc4efdfef..8c7f4557992 100644 --- a/litellm/experimental_mcp_client/client.py +++ b/litellm/experimental_mcp_client/client.py @@ -150,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 @@ -522,7 +522,7 @@ 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 a valid MCP response from %s", diff --git a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py index 31252479445..5b74caee3a2 100644 --- a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py @@ -170,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, _as_read_timeout + 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, ) @@ -1410,7 +1410,7 @@ if MCP_AVAILABLE: 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_read_timeout(cause) is not None + isinstance(cause, McpError) and as_mcp_read_timeout(cause) is not None for cause in iter_exception_tree(e) ) else timeout_seconds 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 a1e781d17b7..b07e5876e8b 100644 --- a/tests/test_litellm/experimental_mcp_client/test_mcp_client.py +++ b/tests/test_litellm/experimental_mcp_client/test_mcp_client.py @@ -36,9 +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, ) @@ -867,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 @@ -1619,7 +1619,7 @@ async def test_transport_completion_and_normal_messages(transport: MCPTransport, 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) + assert isinstance(as_mcp_read_timeout(caught.value), TimeoutError) @pytest.mark.asyncio @@ -1685,4 +1685,4 @@ async def test_empty_http_event_stream_uses_the_existing_request_deadline() -> N ), timeout=3, ) - assert isinstance(_as_read_timeout(caught.value), TimeoutError) + 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 101b02e6ffd..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 @@ -3453,7 +3453,8 @@ class TestConnectionErrorMessage: @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: + @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 @@ -3469,10 +3470,10 @@ class TestConnectionErrorMessage: raise TimeoutError() from sdk_error payload: Final = NewMCPServerRequest( - server_name="timeout", url="https://example.com", auth_type=MCPAuth.none, timeout=1 + 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 ("within 1s" if sdk_timeout else "within 30s") in result["message"] + 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):