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

@@ -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
 

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)

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,

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)
 

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"
 

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"]