diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/mcp_handler.py b/litellm/llms/anthropic/experimental_pass_through/messages/mcp_handler.py index c1f10c245f8..d9cc65e730f 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/mcp_handler.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/mcp_handler.py @@ -82,7 +82,7 @@ async def anthropic_messages_with_mcp( LiteLLM_Proxy_MCP_Handler, ) - mcp_references, other_tools = LiteLLM_Proxy_MCP_Handler._parse_mcp_tools(tools) + mcp_references, other_tools = await LiteLLM_Proxy_MCP_Handler._split_mcp_tools(tools) if not mcp_references: return await _AnthropicMessagesCall(fn=litellm.anthropic_messages).fn( diff --git a/litellm/proxy/hooks/mcp_semantic_filter/hook.py b/litellm/proxy/hooks/mcp_semantic_filter/hook.py index 8b82842353c..5c55b35c968 100644 --- a/litellm/proxy/hooks/mcp_semantic_filter/hook.py +++ b/litellm/proxy/hooks/mcp_semantic_filter/hook.py @@ -104,7 +104,7 @@ class SemanticToolFilterHook(CustomLogger): ) # Parse to separate MCP tools from other tools - mcp_tools, _ = LiteLLM_Proxy_MCP_Handler._parse_mcp_tools(tools) + mcp_tools, _ = await LiteLLM_Proxy_MCP_Handler._split_mcp_tools(tools) if not mcp_tools: return [] diff --git a/litellm/responses/main.py b/litellm/responses/main.py index f012ec8f07b..bae7c491926 100644 --- a/litellm/responses/main.py +++ b/litellm/responses/main.py @@ -175,7 +175,7 @@ async def aresponses_api_with_mcp( ( mcp_tools_with_litellm_proxy, other_tools, - ) = LiteLLM_Proxy_MCP_Handler._parse_mcp_tools(tools) + ) = await LiteLLM_Proxy_MCP_Handler._split_mcp_tools(tools) # Process MCP tools through the complete pipeline (fetch + filter + deduplicate + transform) # Extract user_api_key_auth from litellm_metadata (where it's added by add_user_api_key_auth_to_request_metadata) @@ -234,6 +234,7 @@ async def aresponses_api_with_mcp( "timeout": timeout, "custom_llm_provider": custom_llm_provider, **kwargs, + "_skip_mcp_handler": True, } # Handle MCP streaming if requested @@ -839,13 +840,14 @@ def _responses_try_dispatch_mcp_gateway( custom_llm_provider: str | None, kwargs: dict[str, object], _is_async: bool, + skip_mcp_handler: bool, ) -> Any | None: """Return a response when MCP gateway handles the call; otherwise None.""" from litellm.responses.mcp.litellm_proxy_mcp_handler import ( LiteLLM_Proxy_MCP_Handler, ) - if not LiteLLM_Proxy_MCP_Handler._should_use_litellm_mcp_gateway(tools=tools): + if skip_mcp_handler or not LiteLLM_Proxy_MCP_Handler._should_use_litellm_mcp_gateway(tools=tools): return None mcp_call_kwargs: Final = { "input": input, @@ -1015,6 +1017,7 @@ def responses( litellm_logging_obj: Final[LiteLLMLoggingObj] = kwargs.get("litellm_logging_obj") litellm_call_id: Final[str | None] = kwargs.get("litellm_call_id", None) _is_async: Final = kwargs.pop("aresponses", False) is True + skip_mcp_handler: Final = kwargs.pop("_skip_mcp_handler", False) use_chat_completions_api = _pop_use_chat_completions_api_kw(kwargs) client_headers: Final = kwargs.get("headers") @@ -1109,6 +1112,7 @@ def responses( custom_llm_provider=custom_llm_provider, kwargs=kwargs, _is_async=_is_async, + skip_mcp_handler=skip_mcp_handler, ) if _mcp_dispatch is not None: return _mcp_dispatch diff --git a/litellm/responses/mcp/chat_completions_handler.py b/litellm/responses/mcp/chat_completions_handler.py index a75b3768636..ae18d5f6f1b 100644 --- a/litellm/responses/mcp/chat_completions_handler.py +++ b/litellm/responses/mcp/chat_completions_handler.py @@ -106,7 +106,7 @@ async def acompletion_with_mcp( ( mcp_tools_with_litellm_proxy, other_tools, - ) = LiteLLM_Proxy_MCP_Handler._parse_mcp_tools(tools) + ) = await LiteLLM_Proxy_MCP_Handler._split_mcp_tools(tools) if not mcp_tools_with_litellm_proxy: # No MCP tools, proceed with regular completion @@ -114,6 +114,7 @@ async def acompletion_with_mcp( model=model, messages=messages, tools=tools, + _skip_mcp_handler=True, **kwargs, ) diff --git a/litellm/responses/mcp/litellm_proxy_mcp_handler.py b/litellm/responses/mcp/litellm_proxy_mcp_handler.py index 367915156d1..11e87a72813 100644 --- a/litellm/responses/mcp/litellm_proxy_mcp_handler.py +++ b/litellm/responses/mcp/litellm_proxy_mcp_handler.py @@ -1,6 +1,6 @@ import re import traceback -from collections.abc import Iterable, Mapping, Sequence +from collections.abc import Awaitable, Callable, Collection, Iterable, Mapping, Sequence from datetime import datetime from typing import TYPE_CHECKING, Any, Final, Literal, Optional, TypeAlias, TypedDict, overload @@ -23,6 +23,7 @@ from litellm.types.llms.openai import ( ResponsesAPIStreamingResponse, ) from litellm.types.llms.openai import ToolParam as ResponsesToolParam +from litellm.types.mcp_server.mcp_server_manager import MCPServer from litellm.types.utils import ( CallTypes, ChatCompletionMessageCustomToolCall, @@ -45,6 +46,7 @@ else: # NOTE: We intentionally keep ToolParam as a broad type here to avoid tight coupling ToolParam: TypeAlias = Mapping[str, object] +SplitTools: TypeAlias = tuple[list[ToolParam], list[Any]] class MCPToolResult(TypedDict): @@ -56,14 +58,65 @@ class MCPToolResult(TypedDict): LITELLM_PROXY_MCP_SERVER_URL: Final = "litellm_proxy" LITELLM_PROXY_MCP_SERVER_URL_PREFIX: Final = f"{LITELLM_PROXY_MCP_SERVER_URL}/mcp/" -# Matches any URL whose path ends with /mcp/ — covers both root-path -# (http://host:port/mcp/name) and sub-path (http://host/base/mcp/name) proxy deployments. -# A false-positive match (e.g. an external URL that happens to end with /mcp/) results -# in a "server not found" error from the internal gateway, not a silent failure or data leak, -# so this broad pattern is intentional and preferred over anchoring to localhost only. _PROXY_MCP_PATH_RE: Final = re.compile(r"^https?://.+/mcp/([^/]+)$") +def _mcp_server_url(tool: ToolParam) -> str | None: + if not isinstance(tool, dict) or tool.get("type") != "mcp": + return None + server_url: Final = tool.get("server_url") + return server_url if isinstance(server_url, str) else None + + +def _names_gateway_explicitly(tool: ToolParam) -> bool: + return (_mcp_server_url(tool) or "").startswith(LITELLM_PROXY_MCP_SERVER_URL) + + +def _proxy_path_mcp_name(tool: ToolParam) -> str | None: + server_url: Final = _mcp_server_url(tool) + match: Final = None if server_url is None else _PROXY_MCP_PATH_RE.match(server_url) + return None if match is None else match.group(1) + + +def _registered_mcp_servers() -> Collection[MCPServer]: + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + + return global_mcp_server_manager.get_registry().values() + + +def _registry_serves(name: str, servers: Collection[MCPServer]) -> bool: + return any( + name in (server.alias, server.server_name, server.name) or name in (server.access_groups or ()) + for server in servers + ) + + +async def _toolset_exists(name: str) -> bool: + try: + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + from litellm.proxy.proxy_server import prisma_client + + if prisma_client is None: + return False + return await global_mcp_server_manager.get_toolset_by_name_cached(prisma_client, name) is not None + except Exception as e: + verbose_logger.debug("Could not resolve '%s' as toolset: %s", name, e) + return False + + +async def _gateway_served_names( + names: Collection[str], + servers: Callable[[], Collection[MCPServer]] = _registered_mcp_servers, + toolset_exists: Callable[[str], Awaitable[bool]] = _toolset_exists, +) -> frozenset[str]: + registered: Final = tuple(servers()) if names else () + return frozenset([name for name in names if _registry_serves(name, registered) or await toolset_exists(name)]) + + class LiteLLM_Proxy_MCP_Handler: """ Helper class with static methods for MCP integration with Responses API. @@ -87,57 +140,31 @@ class LiteLLM_Proxy_MCP_Handler: @staticmethod def _should_use_litellm_mcp_gateway(tools: Iterable[ToolParam] | None) -> bool: - """ - Returns True if any MCP tool should be handled via the litellm proxy MCP gateway. - This includes tools with server_url="litellm_proxy" as well as URLs ending in /mcp/. - """ - if tools: - for tool in tools: - if isinstance(tool, dict) and tool.get("type") == "mcp": - server_url = tool.get("server_url", "") - if isinstance(server_url, str) and server_url.startswith(LITELLM_PROXY_MCP_SERVER_URL): - return True - if isinstance(server_url, str) and _PROXY_MCP_PATH_RE.match(server_url): - return True - return False + """True when a tool may name this gateway: server_url "litellm_proxy..." or an http(s) URL ending in + /mcp/. `_split_mcp_tools` then settles which of the latter the gateway actually serves.""" + return any(_names_gateway_explicitly(tool) or _proxy_path_mcp_name(tool) is not None for tool in tools or ()) @staticmethod - def _parse_mcp_tools( + def _parse_mcp_tools(tools: Iterable[Mapping[str, object]] | None) -> SplitTools: + items: Final = tuple(tools or ()) + gateway_tools: Final[list[ToolParam]] = [tool for tool in items if _names_gateway_explicitly(tool)] + other_tools: Final[list[Any]] = [tool for tool in items if not _names_gateway_explicitly(tool)] + return gateway_tools, other_tools + + @staticmethod + async def _split_mcp_tools( tools: Iterable[Mapping[str, object]] | None, - ) -> tuple[list[ToolParam], list[Any]]: - """ - Parse tools and separate MCP tools with litellm_proxy from other tools. - - Returns: - Tuple of (mcp_tools_with_litellm_proxy, other_tools) - """ - mcp_tools_with_litellm_proxy: Final[list[ToolParam]] = [] - other_tools: Final[list[Any]] = [] - - if tools: - for tool in tools: - if isinstance(tool, dict) and tool.get("type") == "mcp": - server_url = tool.get("server_url", "") - if isinstance(server_url, str) and server_url.startswith(LITELLM_PROXY_MCP_SERVER_URL): - mcp_tools_with_litellm_proxy.append(tool) - elif isinstance(server_url, str): - # Also intercept URLs like http://localhost:4000/mcp/atlassian_test - # by rewriting them to the internal litellm_proxy format. - m = _PROXY_MCP_PATH_RE.match(server_url) - if m: - rewritten = { - **tool, - "server_url": f"{LITELLM_PROXY_MCP_SERVER_URL_PREFIX}{m.group(1)}", - } - mcp_tools_with_litellm_proxy.append(rewritten) - else: - other_tools.append(tool) - else: - other_tools.append(tool) - else: - other_tools.append(tool) - - return mcp_tools_with_litellm_proxy, other_tools + served_names: Callable[[Collection[str]], Awaitable[frozenset[str]]] = _gateway_served_names, + ) -> SplitTools: + resolved: Final = tuple((tool, _proxy_path_mcp_name(tool)) for tool in tools or ()) + names: Final = frozenset(name for _, name in resolved if name is not None) + served: Final = await served_names(names) if names else frozenset[str]() + return LiteLLM_Proxy_MCP_Handler._parse_mcp_tools( + [ + {**tool, "server_url": f"{LITELLM_PROXY_MCP_SERVER_URL_PREFIX}{name}"} if name in served else tool + for tool, name in resolved + ] + ) @staticmethod async def _apply_toolset_permissions( diff --git a/tests/test_litellm/responses/mcp/test_chat_completions_handler.py b/tests/test_litellm/responses/mcp/test_chat_completions_handler.py index ab4c5185057..2c1845f7b92 100644 --- a/tests/test_litellm/responses/mcp/test_chat_completions_handler.py +++ b/tests/test_litellm/responses/mcp/test_chat_completions_handler.py @@ -1,6 +1,13 @@ +import json +import sys +import types + import pytest +import respx +from httpx import Response from unittest.mock import AsyncMock, patch +import litellm from litellm.types.utils import ModelResponse from litellm.responses.mcp import chat_completions_handler @@ -1344,3 +1351,39 @@ async def test_acompletion_with_mcp_streaming_drains_inner_stream_after_exhausti assert len(all_chunks) == 3 assert initial_stream.drained_after_exhaustion is True + + +@pytest.mark.asyncio +@respx.mock +async def test_acompletion_with_mcp_forwards_unserved_external_mcp_tool_to_the_provider(monkeypatch): + from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager + + zapier_tool = {"type": "mcp", "server_label": "zapier", "server_url": "https://mcp.zapier.com/api/mcp/mcp"} + monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", types.SimpleNamespace(prisma_client=None)) + monkeypatch.setattr(global_mcp_server_manager, "get_registry", lambda: {}) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + provider = respx.post("https://api.openai.com/v1/chat/completions").mock( + return_value=Response( + 200, + json={ + "id": "chatcmpl-zapier", + "object": "chat.completion", + "created": 0, + "model": "gpt-4.1", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, + }, + ) + ) + + result = await acompletion_with_mcp( + model="openai/gpt-4.1", + messages=[{"role": "user", "content": "hello"}], + tools=[zapier_tool], + api_key="sk-test", + acompletion=True, + ) + + assert isinstance(result, ModelResponse) + assert result.id == "chatcmpl-zapier" + assert json.loads(provider.calls.last.request.content)["tools"] == [zapier_tool] diff --git a/tests/test_litellm/responses/mcp/test_litellm_proxy_mcp_handler.py b/tests/test_litellm/responses/mcp/test_litellm_proxy_mcp_handler.py index b33ed3bb581..73d245df651 100644 --- a/tests/test_litellm/responses/mcp/test_litellm_proxy_mcp_handler.py +++ b/tests/test_litellm/responses/mcp/test_litellm_proxy_mcp_handler.py @@ -6,6 +6,7 @@ from unittest.mock import AsyncMock, MagicMock import pytest from fastapi import HTTPException +from openai.types.responses.tool_param import Mcp import importlib from litellm.proxy._experimental.mcp_server.faults.list_outcomes import AggregateToolListing @@ -719,3 +720,144 @@ def test_extract_tool_call_details_still_prefers_openai_arguments(): assert name == "get_weather" assert call_id == "call_123" assert arguments == '{"city": "Paris"}' + + +def _registered( + server_id: str, + name: str, + alias: str | None = None, + server_name: str | None = None, + access_groups: list[str] | None = None, +): + from litellm.types.mcp import MCPTransport + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + return MCPServer( + server_id=server_id, + name=name, + alias=alias, + server_name=server_name, + transport=MCPTransport.http, + access_groups=access_groups, + ) + + +async def _no_toolset(_: str) -> bool: + return False + + +ZAPIER_TOOL: Mcp = { + "type": "mcp", + "server_label": "zapier", + "server_url": "https://mcp.zapier.com/api/mcp/mcp", + "require_approval": "never", +} +EXPLICIT_GATEWAY_TOOL = {"type": "mcp", "server_label": "github", "server_url": "litellm_proxy/mcp/github"} +FUNCTION_TOOL = {"type": "function", "name": "get_weather", "parameters": {}} + + +@pytest.mark.asyncio +async def test_gateway_served_names_matches_alias_server_name_name_access_group_and_toolset(): + from litellm.responses.mcp.litellm_proxy_mcp_handler import _gateway_served_names + + servers = ( + _registered("id-1", "github-name", alias="github", server_name="github-server", access_groups=["prod-group"]), + _registered("id-2", "deepwiki"), + ) + + async def toolset_exists(name: str) -> bool: + return name == "my-toolset" + + served = await _gateway_served_names( + {"github", "github-server", "github-name", "deepwiki", "prod-group", "my-toolset", "mcp", "nope"}, + servers=lambda: servers, + toolset_exists=toolset_exists, + ) + + assert served == {"github", "github-server", "github-name", "deepwiki", "prod-group", "my-toolset"} + + +@pytest.mark.asyncio +async def test_split_mcp_tools_leaves_external_mcp_path_urls_for_the_provider(): + + async def served_names(names): + assert names == {"mcp"} + return frozenset() + + gateway_tools, other_tools = await LiteLLM_Proxy_MCP_Handler._split_mcp_tools( + [ZAPIER_TOOL, EXPLICIT_GATEWAY_TOOL, FUNCTION_TOOL], served_names=served_names + ) + + assert gateway_tools == [EXPLICIT_GATEWAY_TOOL] + assert other_tools == [ZAPIER_TOOL, FUNCTION_TOOL] + + +@pytest.mark.asyncio +async def test_split_mcp_tools_repoints_served_proxy_urls_at_the_gateway(): + served_tool = { + "type": "mcp", + "server_label": "toolset", + "server_url": "http://localhost:4000/mcp/my-toolset", + "require_approval": "never", + "allowed_tools": ["get_me"], + } + unserved_tool = {"type": "mcp", "server_label": "typo", "server_url": "http://localhost:4000/mcp/githb"} + + async def served_names(names): + return frozenset({"my-toolset"}) + + gateway_tools, other_tools = await LiteLLM_Proxy_MCP_Handler._split_mcp_tools( + [served_tool, unserved_tool], served_names=served_names + ) + + assert gateway_tools == [{**served_tool, "server_url": "litellm_proxy/mcp/my-toolset"}] + assert other_tools == [unserved_tool] + + +@pytest.mark.asyncio +async def test_split_mcp_tools_skips_resolution_when_nothing_points_at_the_proxy(): + async def served_names(names): + raise AssertionError("no lookup expected") + + gateway_tools, other_tools = await LiteLLM_Proxy_MCP_Handler._split_mcp_tools( + [EXPLICIT_GATEWAY_TOOL, FUNCTION_TOOL], served_names=served_names + ) + + assert gateway_tools == [EXPLICIT_GATEWAY_TOOL] + assert other_tools == [FUNCTION_TOOL] + + +def test_should_use_gateway_still_triggers_on_http_mcp_path(): + assert LiteLLM_Proxy_MCP_Handler._should_use_litellm_mcp_gateway([ZAPIER_TOOL]) is True + assert LiteLLM_Proxy_MCP_Handler._should_use_litellm_mcp_gateway([EXPLICIT_GATEWAY_TOOL]) is True + assert LiteLLM_Proxy_MCP_Handler._should_use_litellm_mcp_gateway([FUNCTION_TOOL]) is False + assert LiteLLM_Proxy_MCP_Handler._should_use_litellm_mcp_gateway(None) is False + + +@pytest.mark.asyncio +async def test_aresponses_api_with_mcp_forwards_unserved_external_mcp_tool_to_the_provider(monkeypatch): + from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager + from litellm.responses import main as responses_main + from litellm.types.llms.openai import ResponsesAPIResponse + + monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", types.SimpleNamespace(prisma_client=None)) + monkeypatch.setattr(global_mcp_server_manager, "get_registry", lambda: {}) + provider_tools: list[object] = [] + + def fake_provider(**kwargs: object) -> object: + request_params = cast(dict[str, object], kwargs["response_api_optional_request_params"]) + provider_tools.append(request_params.get("tools")) + + async def respond() -> ResponsesAPIResponse: + return ResponsesAPIResponse(id="resp_zapier", created_at=0, output=[]) + + return respond() + + monkeypatch.setattr(responses_main.base_llm_http_handler, "response_api_handler", fake_provider) + + response = await responses_main.aresponses_api_with_mcp( + input="Reply with the single word ok.", model="openai/gpt-4.1", tools=[ZAPIER_TOOL] + ) + + assert isinstance(response, ResponsesAPIResponse) + assert provider_tools == [[ZAPIER_TOOL]]