diff --git a/litellm/constants.py b/litellm/constants.py index bc94760b187..9db3fe6ee26 100644 --- a/litellm/constants.py +++ b/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( diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index f65310a77e3..8d6543f99d1 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/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, diff --git a/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_tool_calls_and_headers.py b/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_tool_calls_and_headers.py index b8e28c3f9b8..725a1ffd65e 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_tool_calls_and_headers.py +++ b/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: