Merge e8d21eb72e into 2c6b280755

litellm/constants.py

@@ -256,6 +256,7 @@ MCP_PER_USER_TOKEN_EXPIRY_BUFFER_SECONDS: Final = get_env_int("MCP_PER_USER_TOKE
 
 # MCP timeout defaults (seconds). Override via env vars for slow/custom MCP servers.
 MCP_CLIENT_TIMEOUT: Final = get_env_float("LITELLM_MCP_CLIENT_TIMEOUT", 60.0)
+MCP_TOOL_CALL_TEARDOWN_GRACE: Final = get_env_float("LITELLM_MCP_TOOL_CALL_TEARDOWN_GRACE", 15.0)
 MCP_TOOL_LISTING_TIMEOUT: Final = get_env_float("LITELLM_MCP_TOOL_LISTING_TIMEOUT", 30.0)
 MCP_UPSTREAM_RATE_LIMIT_MAX_RETRIES: Final = get_env_int("LITELLM_MCP_UPSTREAM_RATE_LIMIT_MAX_RETRIES", 3)
 MCP_UPSTREAM_RATE_LIMIT_BASE_DELAY_SECONDS: Final = get_env_float(

litellm/proxy/_experimental/mcp_server/mcp_server_manager.py

@@ -69,6 +69,7 @@ from litellm.constants import (
     MCP_CLIENT_TIMEOUT,
     MCP_HEALTH_CHECK_TIMEOUT,
     MCP_METADATA_TIMEOUT,
+    MCP_TOOL_CALL_TEARDOWN_GRACE,
     MCP_TOOL_LISTING_TIMEOUT,
     MCP_UPSTREAM_SESSION_SHUTDOWN_TIMEOUT_SECONDS,
 )
@@ -1847,6 +1848,55 @@ def _mcp_discovery_cache_ttl() -> float:
     return ttl
 
 
+async def _await_mcp_tool_call_tasks(
+    during_hook_tasks: Sequence[asyncio.Task[object]],
+    tool_call: Callable[[], Awaitable[object]],
+    timeout: float,
+) -> tuple[object, ...]:
+    """Run one MCP tool call beside the during-call hook tasks and await both.
+
+    A cancelled asyncio.gather (or wait_for) re-delivers a native Task.cancel into its children
+    on every event loop iteration, which bypasses the AnyIO shields the MCP SDK's stdio shutdown
+    relies on: the child process stop gets interrupted mid-flight, the stdout drain then waits
+    for an EOF that never arrives, and the abandoned request handler and the tool task spin on
+    each other forever (#46657). asyncio.wait never cancels what it waits on, so the timeout and
+    a cancelled caller instead cancel the tool call's own AnyIO cancel scope, a cancellation
+    level the SDK's shielded shutdown honours, and the tool task gets a bounded window to finish
+    that shutdown before the caller's exception is raised. A tool task still unwinding past that
+    window is left to finish detached rather than cancelled natively.
+    """
+    tool_cancel_scopes: Final[list[anyio.CancelScope]] = []  # mutable-ok: scope slot shared with the abandon path
+
+    async def _run_tool_call() -> object:
+        with anyio.CancelScope() as scope:
+            tool_cancel_scopes.append(scope)
+            return await tool_call()
+
+    tool_task: Final = asyncio.create_task(_run_tool_call())
+
+    async def _abandon_tool_call() -> None:
+        for scope in tool_cancel_scopes:
+            scope.cancel()
+        if not tool_cancel_scopes:
+            tool_task.cancel()
+        for hook_task in during_hook_tasks:
+            hook_task.cancel()
+        with anyio.CancelScope(shield=True):
+            await asyncio.wait({tool_task}, timeout=MCP_TOOL_CALL_TEARDOWN_GRACE)
+        if tool_task.done() and not tool_task.cancelled():
+            tool_task.exception()
+
+    try:
+        _, pending = await asyncio.wait([*during_hook_tasks, tool_task], timeout=timeout)
+    except asyncio.CancelledError:
+        await _abandon_tool_call()
+        raise
+    if pending:
+        await _abandon_tool_call()
+        raise asyncio.TimeoutError
+    return tuple(task.result() for task in (*during_hook_tasks, tool_task))
+
+
 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,
@@ -6688,11 +6738,13 @@ class MCPServerManager:
 
             tool_call_coro = _call_tool_via_client(client, call_tool_params)
 
-        tasks.append(asyncio.create_task(tool_call_coro))
-
         _timeout: Final = mcp_server.timeout if mcp_server.timeout is not None else MCP_CLIENT_TIMEOUT
         try:
-            mcp_responses: Final = await asyncio.wait_for(asyncio.gather(*tasks), timeout=_timeout)
+            mcp_responses: Final = await _await_mcp_tool_call_tasks(
+                during_hook_tasks=tuple(cast("list[asyncio.Task[object]]", tasks)),  # cast-ok: bare-list tasks param
+                tool_call=lambda: tool_call_coro,
+                timeout=_timeout,
+            )
         except asyncio.TimeoutError:
             raise HTTPException(
                 status_code=504,

tests/unit/proxy/_experimental/mcp_server/test_mcp_server_tool_calls_and_headers.py

@@ -4,11 +4,13 @@ import contextlib
 import contextvars
 import json
 import os
+from collections.abc import Callable
 from datetime import datetime, timedelta
 from types import SimpleNamespace
 from typing import Callable, Final, Literal
 from unittest.mock import AsyncMock, MagicMock, patch
 
+import anyio
 import httpx
 import pytest
 from fastapi import HTTPException
@@ -19,6 +21,7 @@ from mcp.types import (
     INVALID_REQUEST,
     METHOD_NOT_FOUND,
     BlobResourceContents,
+    CallToolRequestParams as MCPCallToolRequestParams,
     CallToolResult,
     ListToolsResult,
     PaginatedRequestParams,
@@ -33,6 +36,7 @@ from mcp_types.version import HANDSHAKE_PROTOCOL_VERSIONS, LATEST_HANDSHAKE_VERS
 from pydantic import JsonValue, TypeAdapter
 from starlette.types import Message, Receive, Scope, Send
 
+from litellm.experimental_mcp_client.client import MCPClient, PersistentMCPSession
 from litellm.proxy._experimental.mcp_server.mcp_context import active_mcp_request_ctx_var
 from litellm.proxy._experimental.mcp_server.mcp_server_manager import ListedToolsCaller
 from litellm.proxy._experimental.mcp_server.faults.list_outcomes import AggregateToolListing
@@ -866,6 +870,99 @@ async def test_call_tool_m2m_skips_authorization_headers():
     assert extra_headers.get("X-Custom") == "trace"
 
 
+class _ShieldedTeardownStdioClient(MCPClient):
+    """Stands in for the SDK's stdio session: the tool call parks on the server reply, and the
+    transport teardown is an anyio-shielded block that stops the child process and only then
+    returns. A native Task.cancel thrown into that block interrupts the stop mid-flight, so the
+    caller must deliver cancellation at anyio level instead (#46657)"""
+
+    def __init__(self) -> None:
+        super().__init__(server_url="http://upstream/mcp", transport_type=MCPTransport.http)
+        self.call_in_flight = asyncio.Event()
+        self.child_stopped = False
+        self._server_reply = asyncio.Event()
+        self._teardown_unblocked = False
+
+    def release(self) -> None:
+        self._teardown_unblocked = True
+
+    async def call_tool(
+        self,
+        call_tool_request_params: MCPCallToolRequestParams,
+        host_progress_callback: Callable | None = None,
+        raise_on_error: bool = False,
+        allow_input_required: bool = False,
+        persistent_session: PersistentMCPSession | None = None,
+    ) -> CallToolResult:
+        del call_tool_request_params, host_progress_callback, raise_on_error, allow_input_required, persistent_session
+        try:
+            self.call_in_flight.set()
+            await self._server_reply.wait()
+            return CallToolResult(content=[], isError=False)
+        finally:
+            with anyio.CancelScope(shield=True):
+                while not self.child_stopped and not self._teardown_unblocked:
+                    try:
+                        await asyncio.sleep(0.05)
+                        self.child_stopped = True
+                    except asyncio.CancelledError:
+                        continue
+
+
+@pytest.mark.asyncio
+async def test_abandoned_stdio_tool_call_stops_the_child_and_releases_the_caller():
+    """A cancelled tools/call handler must cancel the in-flight call at a cancellation level the
+    SDK's shielded stdio teardown honours, so the child process is stopped and the cancelled
+    handler is released instead of the two spinning on each other forever (#46657)"""
+    try:
+        from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
+            MCPServerManager,
+        )
+    except ImportError:
+        pytest.skip("MCP server not available")
+
+    manager: Final = MCPServerManager()
+    server: Final = MCPServer(
+        server_id="slow-stdio",
+        name="slow-stdio",
+        server_name="slow-stdio",
+        transport=MCPTransport.http,
+        auth_type=MCPAuth.none,
+    )
+    client: Final = _ShieldedTeardownStdioClient()
+
+    with patch.object(manager, "create_mcp_client", new=AsyncMock(return_value=client)):
+
+        async def abandoned_call() -> None:
+            async with anyio.create_task_group() as tg:
+
+                async def handler() -> None:
+                    await manager._call_regular_mcp_tool(
+                        mcp_server=server,
+                        original_tool_name="execute_sql",
+                        arguments={},
+                        tasks=[],
+                        mcp_auth_header=None,
+                        mcp_server_auth_headers=None,
+                        oauth2_headers=None,
+                        raw_headers=None,
+                        proxy_logging_obj=None,
+                    )
+
+                tg.start_soon(handler)
+                await client.call_in_flight.wait()
+                tg.cancel_scope.cancel()
+
+        scenario: Final = asyncio.create_task(abandoned_call())
+        _, pending = await asyncio.wait({scenario}, timeout=10)
+        try:
+            assert not pending, "the abandoned tools/call never released the cancelled request handler"
+            assert client.child_stopped, "the abandoned tools/call never stopped the stdio child process"
+        finally:
+            client.release()
+            await asyncio.wait({scenario}, timeout=5)
+
+
 @pytest.mark.asyncio
 async def test_get_prompts_from_mcp_servers_success():
     try: