From 422790b188c65efbb4bf50dcc4f078986529dafb Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 15 Jul 2026 12:17:58 +0000 Subject: [PATCH] fix(mcp): stop tools/list from prefetching upstream instructions and blocking on listing-timeout cleanup --- .../mcp_server/mcp_server_manager.py | 44 ++++- .../proxy/_experimental/mcp_server/server.py | 9 +- .../mcp_server/test_mcp_server.py | 162 ++++++++++++++++++ .../mcp_server/test_mcp_server_manager.py | 43 +++++ 4 files changed, 248 insertions(+), 10 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index e6e265abb61..28ebd67df9a 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -17,7 +17,6 @@ from contextlib import asynccontextmanager from typing import Any, AsyncIterator, Callable, Literal, Optional, Union, cast from urllib.parse import urlparse -import anyio import httpx from fastapi import HTTPException from httpx import HTTPStatusError @@ -508,6 +507,25 @@ def _extract_upstream_auth_failure( return None +def _detach_task(task: "asyncio.Future[Any]") -> None: + """Cancel a still-running task and consume its outcome in the background. + + Awaiting the task from the request path would block on the MCP SDK's + session/transport teardown, which can hang far past the intended deadline + against an unresponsive upstream. Cancelling and retrieving the eventual + result via a done callback lets that cleanup finish detached from the + caller without emitting an "exception was never retrieved" warning. + """ + task.cancel() + + def _consume(finished: "asyncio.Future[Any]") -> None: + if finished.cancelled(): + return + finished.exception() + + task.add_done_callback(_consume) + + def _warn_on_server_name_fields( *, server_id: str, @@ -3353,8 +3371,14 @@ class MCPServerManager: """ Fetch tools from MCP client with timeout and error handling. - Uses anyio.fail_after() instead of asyncio.wait_for() to avoid conflicts - with the MCP SDK's anyio TaskGroup. See GitHub issue #20715 for details. + The listing runs in its own task so the deadline returns promptly even + when the MCP SDK's session/transport teardown is slow: on timeout the + task is cancelled and its cleanup detached from the request path rather + than awaited, which could otherwise take far longer than + MCP_TOOL_LISTING_TIMEOUT against an unresponsive upstream. Running the + operation in a dedicated task also keeps cancellation inside the SDK's + own anyio TaskGroup instead of firing it into this coroutine's scope, + avoiding the asyncio.wait_for conflict from GitHub issue #20715. An upstream HTTP 401 is converted into :class:`MCPUpstreamAuthError` instead of being swallowed to an empty tool list, regardless of the @@ -3374,14 +3398,16 @@ class MCPServerManager: Returns: List of tools from the server """ - try: - with anyio.fail_after(MCP_TOOL_LISTING_TIMEOUT): - tools = await client.list_tools(raise_on_error=True) - verbose_logger.debug(f"Tools from {server_name}: {tools}") - return tools - except TimeoutError: + listing_task = asyncio.ensure_future(client.list_tools(raise_on_error=True)) + done, _pending = await asyncio.wait({listing_task}, timeout=MCP_TOOL_LISTING_TIMEOUT) + if listing_task not in done: verbose_logger.warning(f"Timeout while listing tools from {server_name}") + _detach_task(listing_task) return [] + try: + tools = listing_task.result() + verbose_logger.debug(f"Tools from {server_name}: {tools}") + return tools except asyncio.CancelledError: verbose_logger.warning(f"Task cancelled while listing tools from {server_name}") return [] diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 68a61b85175..5530e259991 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -1713,13 +1713,19 @@ if MCP_AVAILABLE: mcp_servers: Optional[List[str]], client_ip: Optional[str], scoped_server_endpoint: bool = False, + prefetch_upstream_instructions: bool = True, ) -> AsyncIterator[None]: allowed = await _get_allowed_mcp_servers( user_api_key_auth=user_api_key_auth, mcp_servers=mcp_servers, client_ip=client_ip, ) - if allowed: + if allowed and prefetch_upstream_instructions: + # Only the `initialize` response carries these instructions, so + # opening upstream sessions to fetch them for any other method + # (e.g. `tools/list`) just adds latency to that request. The merge + # below still surfaces anything already cached by a prior probe. + # # return_exceptions=True: a per-server probe failure (incl. CancelledError # bubbled from anyio task group teardown on connection refused) must not # cancel sibling probes or 500 the gateway initialize request. @@ -4193,6 +4199,7 @@ if MCP_AVAILABLE: mcp_servers, _client_ip, scoped_server_endpoint=scoped_server_endpoint, + prefetch_upstream_instructions=is_initialize, ): await target_manager.handle_request(scope, receive, local_send) if use_stateful and session_id and scope.get("method") == "DELETE": diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py index a983ac3ff48..dc84f51c2ae 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py @@ -5121,6 +5121,168 @@ class TestGatewayCreateInitializationOptions: assert getattr(opts, "instructions", None) is None +class TestGatewayInstructionsPrefetchGating: + """Regression for #33374: the upstream initialize-instructions prefetch must + only run for the `initialize` request, never for `tools/list` (or any other + method), so a slow upstream server does not add its initialize latency to + every listing. Instructions already cached from a prior probe still merge.""" + + def _server(self) -> MCPServer: + return MCPServer( + server_id="slow-1", + name="slow", + alias="slow", + transport=MCPTransport.http, + url="https://example.com/mcp", + ) + + @pytest.mark.asyncio + async def test_scope_skips_prefetch_when_not_initialize(self): + try: + from litellm.proxy._experimental.mcp_server.server import ( + _gateway_initialize_instructions_request_scope, + global_mcp_server_manager, + server, + ) + except ImportError: + pytest.skip("MCP server not available") + + srv = self._server() + global_mcp_server_manager._upstream_initialize_instructions_by_server_id["slow-1"] = "cached hi" + try: + with ( + patch( + "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers", + new_callable=AsyncMock, + return_value=[srv], + ), + patch.object( + global_mcp_server_manager, + "_ensure_upstream_initialize_instructions_cached", + new_callable=AsyncMock, + ) as prefetch, + ): + async with _gateway_initialize_instructions_request_scope( + user_api_key_auth=None, + mcp_servers=["slow"], + client_ip=None, + prefetch_upstream_instructions=False, + ): + assert server.create_initialization_options().instructions == "cached hi" + prefetch.assert_not_awaited() + finally: + global_mcp_server_manager._upstream_initialize_instructions_by_server_id.pop("slow-1", None) + + @pytest.mark.asyncio + async def test_scope_prefetches_when_initialize(self): + try: + from litellm.proxy._experimental.mcp_server.server import ( + _gateway_initialize_instructions_request_scope, + global_mcp_server_manager, + ) + except ImportError: + pytest.skip("MCP server not available") + + srv = self._server() + with ( + patch( + "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers", + new_callable=AsyncMock, + return_value=[srv], + ), + patch.object( + global_mcp_server_manager, + "_ensure_upstream_initialize_instructions_cached", + new_callable=AsyncMock, + ) as prefetch, + ): + async with _gateway_initialize_instructions_request_scope( + user_api_key_auth=None, + mcp_servers=["slow"], + client_ip=None, + prefetch_upstream_instructions=True, + ): + pass + prefetch.assert_awaited_once() + + async def _run_streamable(self, body: bytes): + from litellm.proxy._experimental.mcp_server import server as mcp_server + from litellm.proxy._experimental.mcp_server.server import ( + global_mcp_server_manager, + handle_streamable_http_mcp, + ) + + srv = self._server() + + async def receive(): + return {"type": "http.request", "body": body, "more_body": False} + + scope = { + "type": "http", + "method": "POST", + "path": "/mcp", + "headers": [(b"content-type", b"application/json"), (b"authorization", b"Bearer sk-test")], + } + + with ( + patch( + "litellm.proxy._experimental.mcp_server.server.extract_mcp_auth_context", + new_callable=AsyncMock, + return_value=(UserAPIKeyAuth(api_key="sk-test"), None, None, None, None, None), + ), + patch("litellm.proxy._experimental.mcp_server.server.set_auth_context"), + patch( + "litellm.proxy._experimental.mcp_server.server._raise_preemptive_401_for_unauthenticated_servers", + new_callable=AsyncMock, + ), + patch( + "litellm.proxy._experimental.mcp_server.server._check_passthrough_upstream_auth", + new_callable=AsyncMock, + ), + patch("litellm.proxy._experimental.mcp_server.server._SESSION_MANAGERS_INITIALIZED", True), + patch( + "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers", + new_callable=AsyncMock, + return_value=[srv], + ), + patch.object( + global_mcp_server_manager, + "_ensure_upstream_initialize_instructions_cached", + new_callable=AsyncMock, + ) as prefetch, + patch.object(mcp_server.session_manager_stateful, "handle_request", new_callable=AsyncMock), + patch.object(mcp_server.session_manager_stateful, "_server_instances", {}), + patch.object(mcp_server.session_manager_stateless, "handle_request", new_callable=AsyncMock), + patch.object(mcp_server.session_manager_stateless, "_server_instances", {}), + ): + await handle_streamable_http_mcp(scope, receive, AsyncMock()) + return prefetch + + @pytest.mark.asyncio + async def test_streamable_tools_list_does_not_prefetch(self): + import json as _json + + try: + body = _json.dumps({"jsonrpc": "2.0", "id": 1, "method": "tools/list"}).encode() + prefetch = await self._run_streamable(body) + except ImportError: + pytest.skip("MCP server not available") + prefetch.assert_not_awaited() + + @pytest.mark.asyncio + async def test_streamable_initialize_prefetches(self): + import json as _json + + try: + body = _json.dumps( + {"jsonrpc": "2.0", "id": 1, "method": "initialize", "params": {}} + ).encode() + prefetch = await self._run_streamable(body) + except ImportError: + pytest.skip("MCP server not available") + prefetch.assert_awaited() + + @pytest.mark.asyncio async def test_list_tools_with_legacy_db_m2m_server_resolves_oauth2_flow(): """ diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py index adcfff6fe9d..1858297668d 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -7042,6 +7042,49 @@ class TestMCPToolsListAuthSurfacing: assert [t.name for t in result] == ["good-do_thing"] +class TestFetchToolsTimeoutDoesNotBlockOnCleanup: + """Regression for #33374: the tool-listing deadline must return promptly + even when the MCP SDK's session/transport teardown is slow. The listing + runs in its own task and, on timeout, that task is cancelled and its + cleanup detached from the request path rather than awaited.""" + + @pytest.mark.asyncio + async def test_timeout_returns_at_deadline_without_awaiting_cleanup(self, monkeypatch): + import time + + from litellm.proxy._experimental.mcp_server import mcp_server_manager as mgr_mod + + cleanup_started = asyncio.Event() + cleanup_finished = asyncio.Event() + + class HangingClient: + async def list_tools(self, raise_on_error: bool = False): + try: + await asyncio.sleep(10) + except asyncio.CancelledError: + cleanup_started.set() + await asyncio.sleep(0.5) + cleanup_finished.set() + raise + return [] + + monkeypatch.setattr(mgr_mod, "MCP_TOOL_LISTING_TIMEOUT", 0.05) + manager = MCPServerManager() + + start = time.monotonic() + result = await manager._fetch_tools_with_timeout(HangingClient(), "slow-server") + elapsed = time.monotonic() - start + + assert result == [] + assert elapsed < 0.4, "listing must return at the deadline, not after upstream cleanup" + assert not cleanup_finished.is_set(), "caller must not block on the slow teardown" + + # The cancelled listing keeps unwinding detached from the request path; + # it eventually runs (and finishes) its slow teardown in the background. + await asyncio.wait_for(cleanup_finished.wait(), timeout=2) + assert cleanup_started.is_set() + + def test_should_strip_caller_authorization_for_token_exchange(): """OBO: the inbound bearer is the subject token (exchanged), never forwarded upstream raw.""" server = MCPServer(