From e4df98272e46a5abe3136abbcf5a8104262f88da Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Thu, 15 Oct 2026 15:46:50 -0700 Subject: [PATCH] fix(mcp): resolve /mcp/proxy tool_id from the owning server instead of relisting every server (#46545) Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/experimental_mcp_client/client.py | 5 +- .../mcp_server/mcp_server_manager.py | 13 +++- .../_experimental/mcp_server/operations.py | 78 +++++++++++++++++++ .../_experimental/mcp_server/tool_search.py | 38 +++++---- .../mcp_server/faults/test_list_outcomes.py | 44 +++++++++-- .../mcp_server/test_mcp_proxy_mode.py | 66 ++++++++++++++++ 6 files changed, 222 insertions(+), 22 deletions(-) diff --git a/litellm/experimental_mcp_client/client.py b/litellm/experimental_mcp_client/client.py index bc3ba796ee3..e2057e4b627 100644 --- a/litellm/experimental_mcp_client/client.py +++ b/litellm/experimental_mcp_client/client.py @@ -38,6 +38,9 @@ _TransportContext: TypeAlias = AbstractAsyncContextManager[_TransportStreams] _SessionCallbackName: TypeAlias = Literal["sampling_callback", "elicitation_callback", "logging_callback"] +SESSION_CLEANUP_TIMEOUT_MESSAGE: Final = "MCP session cleanup timed out" + + class _SessionCallbackMap(TypedDict): sampling_callback: ReadOnly[SamplingFnT | None] elicitation_callback: ReadOnly[ElicitationFnT | None] @@ -815,7 +818,7 @@ class MCPClient: raise ( in_flight_error if in_flight_error is not None - else asyncio.CancelledError("MCP session cleanup timed out") + else asyncio.CancelledError(SESSION_CLEANUP_TIMEOUT_MESSAGE) ) return result diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index f31742d9647..4e4741b2d17 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -74,6 +74,7 @@ from litellm.constants import ( ) from litellm.exceptions import BlockedPiiEntityError, GuardrailRaisedException from litellm.experimental_mcp_client.client import ( + SESSION_CLEANUP_TIMEOUT_MESSAGE, MCPClient, PersistentMCPSession, UpstreamSessionClosedError, @@ -1843,6 +1844,13 @@ def _mcp_discovery_cache_ttl() -> float: return ttl +def _listing_task_was_cancelled(error: asyncio.CancelledError) -> bool: + """The MCP client raises its own CancelledError(SESSION_CLEANUP_TIMEOUT_MESSAGE) when a session's + teardown outlives its budget while the listing task keeps running. That one is the server's fault, + anything else is a real cancellation.""" + return error.args != (SESSION_CLEANUP_TIMEOUT_MESSAGE,) + + class MCPServerManager: _STDIO_ENV_TEMPLATE_PATTERN = re.compile(r"^\$\{(X-[^}]+)\}$") @@ -5591,7 +5599,10 @@ class MCPServerManager: verbose_logger.warning("Timeout while listing tools from %s", server_name) raise MCPServerListError(ServerListFault(tag="timeout"), server_name) from e except asyncio.CancelledError as e: - verbose_logger.warning("Task cancelled while listing tools from %s", server_name) + if _listing_task_was_cancelled(e): + verbose_logger.debug("Task cancelled while listing tools from %s", server_name) + raise + verbose_logger.warning("Listing tools from %s was abandoned by the client session: %s", server_name, e) raise MCPServerListError(ServerListFault(tag="internal"), server_name) from e except ConnectionError as e: verbose_logger.warning("Connection error while listing tools from %s: %s", server_name, e) diff --git a/litellm/proxy/_experimental/mcp_server/operations.py b/litellm/proxy/_experimental/mcp_server/operations.py index 1cd14e032a3..5047948350c 100644 --- a/litellm/proxy/_experimental/mcp_server/operations.py +++ b/litellm/proxy/_experimental/mcp_server/operations.py @@ -1032,6 +1032,84 @@ def _aggregate_server_key(server: MCPServer) -> str: return get_server_prefix(server) or "unknown" +_proxy_tool_lookups: Final[set[asyncio.Task[MCPTool | None]]] = set() # mutable-ok: refs until teardown ends + + +def _release_proxy_tool_lookup(lookup: asyncio.Task[MCPTool | None]) -> None: + _proxy_tool_lookups.discard(lookup) + if not lookup.cancelled() and lookup.exception() is not None: + verbose_logger.debug("MCP proxy tool lookup failed: %s", lookup.exception()) + + +@catalog_operation(lambda: global_mcp_server_manager) +async def find_mcp_proxy_tool( + tool_id: str, + user_api_key_auth: UserAPIKeyAuth, + mcp_servers: Sequence[str] | None, + client_ip: str | None, + mcp_auth_header: str | None, + mcp_server_auth_headers: Mapping[str, Mapping[str, str]] | None, + oauth2_headers: Mapping[str, str] | None, + raw_headers: Mapping[str, str] | None, +) -> MCPTool | None: + """Resolve an /mcp/proxy tool_id against the caller's catalog, returning as soon as the + server that owns it has listed instead of waiting for every allowed server.""" + from litellm.proxy._experimental.mcp_server.catalog import get_filtered_server_tools + from litellm.proxy._experimental.mcp_server.tool_search import mcp_proxy_tool_id + + allowed_mcp_servers: Final = await _get_allowed_mcp_servers( + user_api_key_auth=user_api_key_auth, + mcp_servers=mcp_servers, + client_ip=client_ip, + ) + if mcp_servers and not allowed_mcp_servers: + await raise_denied_scoped_mcp_access( + requested_names=mcp_servers, + user_api_key_auth=user_api_key_auth, + client_ip=client_ip, + ) + admission: Final = await _admit_mcp_servers(allowed_mcp_servers, user_api_key_auth) + if not admission.admitted_servers and admission.rejected_servers: + raise admission.rejected_servers[0][1] + prefetched_oauth_creds: Final = ( + await _prefetch_oauth_creds_for_user(user_api_key_auth) + if any(server.auth_type == MCPAuth.oauth2 for server in allowed_mcp_servers) + else {} + ) + context: Final = OperationContext( + _caller=user_api_key_auth, + mcp_auth_header=mcp_auth_header, + mcp_servers=tuple(mcp_servers) if mcp_servers is not None else None, + mcp_server_auth_headers=mcp_server_auth_headers, + oauth2_headers=oauth2_headers, + raw_headers=raw_headers, + client_ip=client_ip, + mcp_proxy_mode=True, + ) + + async def _owned_tool(server: MCPServer) -> MCPTool | None: + page, _ = await get_filtered_server_tools( + server, + context=context, + allowed_mcp_servers=allowed_mcp_servers, + prefetched_oauth_creds=prefetched_oauth_creds, + ) + return next((tool for tool in page.tools if mcp_proxy_tool_id(tool) == tool_id), None) + + lookups: Final = tuple(asyncio.create_task(_owned_tool(server)) for server in admission.admitted_servers) + for lookup in lookups: + _proxy_tool_lookups.add(lookup) + lookup.add_done_callback(_release_proxy_tool_lookup) + try: + for lookup in asyncio.as_completed(lookups): + if (tool := await lookup) is not None: + return tool + return None + finally: + for lookup in lookups: + lookup.cancel() + + @catalog_operation(lambda: global_mcp_server_manager) async def _get_tools_from_mcp_servers( user_api_key_auth: UserAPIKeyAuth | None, diff --git a/litellm/proxy/_experimental/mcp_server/tool_search.py b/litellm/proxy/_experimental/mcp_server/tool_search.py index 5eacc404913..18642a9162c 100644 --- a/litellm/proxy/_experimental/mcp_server/tool_search.py +++ b/litellm/proxy/_experimental/mcp_server/tool_search.py @@ -606,21 +606,20 @@ async def handle_mcp_proxy_tool( from litellm.proxy import proxy_server from litellm.proxy._experimental.mcp_server.operations import ( _list_mcp_tools, + find_mcp_proxy_tool, ) - listing: Final = await _list_mcp_tools( - user_api_key_auth=user_api_key_dict, - mcp_servers=mcp_servers, - client_ip=client_ip, - mcp_auth_header=mcp_auth_header, - mcp_server_auth_headers=mcp_server_auth_headers, - oauth2_headers=oauth2_headers, - raw_headers=raw_headers, - mcp_proxy_mode=True, - ) - tools_by_id: Final = {mcp_proxy_tool_id(tool): tool for tool in listing.tools} - if name == MCP_PROXY_SEARCH_TOOL_NAME: + listing: Final = await _list_mcp_tools( + user_api_key_auth=user_api_key_dict, + mcp_servers=mcp_servers, + client_ip=client_ip, + mcp_auth_header=mcp_auth_header, + mcp_server_auth_headers=mcp_server_auth_headers, + oauth2_headers=oauth2_headers, + raw_headers=raw_headers, + mcp_proxy_mode=True, + ) llm_router: Final = proxy_server.llm_router proxy_logging_obj: Final = proxy_server.proxy_logging_obj settings: Final = mcp_tool_search_settings() @@ -646,7 +645,20 @@ async def handle_mcp_proxy_tool( return _text_tool_result(json.dumps(tuple(_proxy_search_result(hit) for hit in results)), is_error=False) tool_id: Final = arguments.get("tool_id") - tool: Final = tools_by_id.get(tool_id) if isinstance(tool_id, str) else None + tool: Final = ( + await find_mcp_proxy_tool( + tool_id, + user_api_key_auth=user_api_key_dict, + mcp_servers=mcp_servers, + client_ip=client_ip, + mcp_auth_header=mcp_auth_header, + mcp_server_auth_headers=mcp_server_auth_headers, + oauth2_headers=oauth2_headers, + raw_headers=raw_headers, + ) + if isinstance(tool_id, str) + else None + ) if tool is None: return _text_tool_result("Unknown or unauthorized tool_id", is_error=True) diff --git a/tests/unit/proxy/_experimental/mcp_server/faults/test_list_outcomes.py b/tests/unit/proxy/_experimental/mcp_server/faults/test_list_outcomes.py index fe2254fbd46..ee33c41afb6 100644 --- a/tests/unit/proxy/_experimental/mcp_server/faults/test_list_outcomes.py +++ b/tests/unit/proxy/_experimental/mcp_server/faults/test_list_outcomes.py @@ -131,21 +131,51 @@ def test_single_upstream_http_status_is_truthful(tag, status_code, expected): @pytest.mark.asyncio -async def test_cancelled_fetch_is_a_classified_fault_not_a_healthy_empty_server(): - """A cancelled per-server fetch must not masquerade as ok(tool_count=0): cancellation was already - suppressed before the outcome plumbing existed, so it stays suppressed, but as an internal fault - the outcome reporting can see.""" +@pytest.mark.parametrize("cancel_message", [None, "client disconnected"]) +async def test_cancelled_fetch_propagates_instead_of_masquerading_as_a_healthy_empty_server(cancel_message: str | None): + """A cancelled per-server fetch must not masquerade as ok(tool_count=0), and it is not an upstream + fault either: the caller cancelled it (a sibling server already answered, or the request went + away), so the cancellation propagates instead of being logged as a listing failure.""" import asyncio - from unittest.mock import AsyncMock, MagicMock + from unittest.mock import MagicMock from litellm.proxy._experimental.mcp_server.mcp_server_manager import MCPServerManager manager = MCPServerManager() client = MagicMock() - client.list_tools = AsyncMock(side_effect=asyncio.CancelledError()) + started = asyncio.Event() + + async def list_tools(raise_on_error: bool = False) -> list[object]: + started.set() + await asyncio.Event().wait() + return [] + + client.list_tools = list_tools + listing = asyncio.create_task(manager._fetch_tools_with_timeout(client, "cancelled_srv")) + await started.wait() + listing.cancel(cancel_message) + + with pytest.raises(asyncio.CancelledError): + await listing + + +@pytest.mark.asyncio +async def test_client_side_cancelled_error_is_a_classified_fault_not_a_cancellation(): + """The MCP client raises its own CancelledError when a session's teardown outlives its cleanup + budget while the listing task keeps running. That is one server's fault, so it classifies as an + internal listing failure instead of cancelling the aggregate listing around it.""" + import asyncio + from unittest.mock import AsyncMock, MagicMock + + from litellm.experimental_mcp_client.client import SESSION_CLEANUP_TIMEOUT_MESSAGE + from litellm.proxy._experimental.mcp_server.mcp_server_manager import MCPServerManager + + manager = MCPServerManager() + client = MagicMock() + client.list_tools = AsyncMock(side_effect=asyncio.CancelledError(SESSION_CLEANUP_TIMEOUT_MESSAGE)) with pytest.raises(MCPServerListError) as exc_info: - await manager._fetch_tools_with_timeout(client, "cancelled_srv") + await manager._fetch_tools_with_timeout(client, "slow_teardown_srv") assert exc_info.value.fault.tag == "internal" diff --git a/tests/unit/proxy/_experimental/mcp_server/test_mcp_proxy_mode.py b/tests/unit/proxy/_experimental/mcp_server/test_mcp_proxy_mode.py index 6f1bbad5040..7dadfc6493a 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_mcp_proxy_mode.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_mcp_proxy_mode.py @@ -1,6 +1,7 @@ from litellm.proxy._experimental.mcp_server import operations as mcp_operations import asyncio import json +import logging from datetime import datetime from unittest.mock import AsyncMock, patch @@ -185,3 +186,68 @@ async def test_proxy_call_tool_on_a_never_listed_tool_hands_the_pre_hook_no_list assert pre_call_tool_check.await_args.kwargs["name"] == "echo" assert pre_call_tool_check.await_args.kwargs["tool"] is None assert listed is None + + +@pytest.mark.asyncio +@pytest.mark.parametrize("meta_tool", ("get_tool_schema", "call_tool")) +async def test_proxy_schema_and_call_resolve_tool_id_without_waiting_on_unrelated_servers( + meta_tool: str, caplog: pytest.LogCaptureFixture +) -> None: + """get_tool_schema and call_tool only need the server that owns tool_id. A sibling server whose + tools/list never answers, and whose session is slow to tear down once cancelled, must not hold up a + tool on a server that already listed, and cancelling that sibling is not an error.""" + manager = mcp_operations.global_mcp_server_manager + fast = MCPServer(server_id="proxy-fast", name="proxy-fast", transport=MCPTransport.http, url="http://fast") + slow = MCPServer(server_id="proxy-slow", name="proxy-slow", transport=MCPTransport.http, url="http://slow") + auth = UserAPIKeyAuth(api_key="sk-proxy-fast", user_id="proxy-caller") + slow_release = asyncio.Event() + slow_teardown = asyncio.Event() + served_as = with_mcp_proxy_identity(MCPTool(name="proxy-fast-echo", inputSchema={}), fast.server_id) + + async def fetch_tools(client: object, server_name: str) -> list[MCPTool]: + if server_name == slow.name: + try: + await slow_release.wait() + except asyncio.CancelledError: + await slow_teardown.wait() + raise + return [MCPTool(name="echo", description="Echo text back", inputSchema={"type": "object"})] + + async def call_regular_mcp_tool(*, tasks: list[asyncio.Task[object]], **_: object) -> CallToolResult: + await asyncio.gather(*tasks) + return CallToolResult(content=[TextContent(type="text", text="echoed")]) + + arguments = {"tool_id": mcp_proxy_tool_id(served_as)} + with ( + patch.dict(manager.registry, {fast.server_id: fast, slow.server_id: slow}), + patch.dict(manager.tool_name_to_mcp_server_name_mapping), + patch.object(mcp_operations, "_get_allowed_mcp_servers", AsyncMock(return_value=[slow, fast])), + patch.object(manager, "create_mcp_client", AsyncMock(return_value=object())), + patch.object(manager, "_fetch_tools_with_timeout", fetch_tools), + patch.object(manager, "pre_call_tool_check", AsyncMock(return_value={})), + patch.object(manager, "_call_regular_mcp_tool", call_regular_mcp_tool), + ): + try: + caplog.set_level(logging.ERROR, logger="LiteLLM") + result = await asyncio.wait_for( + handle_mcp_proxy_tool( + name=meta_tool, + arguments={**arguments, "arguments": {}} if meta_tool == "call_tool" else arguments, + user_api_key_dict=auth, + ), + timeout=15, + ) + finally: + slow_release.set() + slow_teardown.set() + await asyncio.sleep(0) + manager._drop_listed_tools(fast.server_id) + manager._drop_listed_tools(slow.server_id) + + assert result.is_error is False + assert [record.getMessage() for record in caplog.records if slow.name in record.getMessage()] == [] + payload = result.content[0].text + if meta_tool == "call_tool": + assert payload == "echoed" + else: + assert json.loads(payload)["tool_id"] == arguments["tool_id"]