From 2a1527c22fc2c4deafb324b925d2c5e453e92580 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 2 Sep 2026 16:03:31 -0700 Subject: [PATCH] fix(mcp): match gateway-served /mcp/ URLs the way the gateway resolves them and narrow only gateway references --- .../proxy/_experimental/mcp_server/server.py | 2 +- .../proxy/hooks/mcp_semantic_filter/hook.py | 23 +++++++------ .../mcp/litellm_proxy_mcp_handler.py | 32 +++++++++++++++---- .../mcp_server/test_mcp_server.py | 27 ++++++++++++++++ .../mcp_server/test_semantic_tool_filter.py | 23 +++++++++++++ .../mcp/test_litellm_proxy_mcp_handler.py | 31 ++++++++++++++++++ 6 files changed, 122 insertions(+), 16 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index b3b574673db..1100b986f15 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -2037,7 +2037,7 @@ if MCP_AVAILABLE: mcp_servers=mcp_servers, client_ip=client_ip, ) - if mcp_servers is not None and not allowed_mcp_servers: + 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, diff --git a/litellm/proxy/hooks/mcp_semantic_filter/hook.py b/litellm/proxy/hooks/mcp_semantic_filter/hook.py index 5c55b35c968..afc430173ad 100644 --- a/litellm/proxy/hooks/mcp_semantic_filter/hook.py +++ b/litellm/proxy/hooks/mcp_semantic_filter/hook.py @@ -5,7 +5,7 @@ Pre-call hook that filters MCP tools semantically before LLM inference. Reduces context window size and improves tool selection accuracy. """ -from collections.abc import Iterable, Mapping, Sequence +from collections.abc import Awaitable, Callable, Collection, Iterable, Mapping, Sequence from typing import TYPE_CHECKING, Final, Optional from fastapi import HTTPException @@ -173,7 +173,11 @@ class SemanticToolFilterHook(CustomLogger): return [name for name in names if name] @staticmethod - def _narrow_mcp_references(tools: Sequence[Mapping[str, object]], selected_tool_names: list[str]) -> list[object]: + async def _narrow_mcp_references( + tools: Sequence[Mapping[str, object]], + selected_tool_names: list[str], + served_names: Callable[[Collection[str]], Awaitable[frozenset[str]]] | None = None, + ) -> list[object]: """ Restrict each litellm_proxy MCP reference to the semantically selected tools. @@ -192,13 +196,14 @@ class SemanticToolFilterHook(CustomLogger): LiteLLM_Proxy_MCP_Handler, ) + via_gateway: Final = await ( + LiteLLM_Proxy_MCP_Handler._routes_through_gateway(tools, served_names) + if served_names is not None + else LiteLLM_Proxy_MCP_Handler._routes_through_gateway(tools) + ) return [ - ( - {**tool, "allowed_tools": selected_tool_names} - if isinstance(tool, dict) and LiteLLM_Proxy_MCP_Handler._should_use_litellm_mcp_gateway([tool]) - else tool - ) - for tool in tools + {**tool, "allowed_tools": selected_tool_names} if isinstance(tool, dict) and routed else tool + for tool, routed in zip(tools, via_gateway, strict=True) ] def _is_mcp_tool(self, tool: object) -> bool: @@ -325,7 +330,7 @@ class SemanticToolFilterHook(CustomLogger): filtered_expanded_tools = await self._filter_expanded_tools(data=data, expanded_tools=expanded_tools) selected_tool_names: Final = self._selected_tool_names(filtered_expanded_tools) - narrowed_tools: Final = self._narrow_mcp_references(tools, selected_tool_names) + narrowed_tools: Final = await self._narrow_mcp_references(tools, selected_tool_names) data["tools"] = narrowed_tools self._emit_filter_metadata_safe( data=data, diff --git a/litellm/responses/mcp/litellm_proxy_mcp_handler.py b/litellm/responses/mcp/litellm_proxy_mcp_handler.py index 11e87a72813..5db6bac10a1 100644 --- a/litellm/responses/mcp/litellm_proxy_mcp_handler.py +++ b/litellm/responses/mcp/litellm_proxy_mcp_handler.py @@ -11,6 +11,7 @@ from litellm._logging import verbose_logger from litellm.constants import MAXIMUM_TRACEBACK_LINES_TO_LOG from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.proxy._experimental.mcp_server.utils import ( + iter_known_server_prefixes, logging_safe_mcp_headers, split_server_prefix_from_name, strip_known_server_prefix, @@ -87,8 +88,10 @@ def _registered_mcp_servers() -> Collection[MCPServer]: def _registry_serves(name: str, servers: Collection[MCPServer]) -> bool: + requested: Final = name.lower() return any( - name in (server.alias, server.server_name, server.name) or name in (server.access_groups or ()) + requested in (known.lower() for known in (*iter_known_server_prefixes(server), server.name)) + or name in (server.access_groups or ()) for server in servers ) @@ -117,6 +120,13 @@ async def _gateway_served_names( return frozenset([name for name in names if _registry_serves(name, registered) or await toolset_exists(name)]) +async def _served_mcp_path_names( + tools: Collection[ToolParam], served_names: Callable[[Collection[str]], Awaitable[frozenset[str]]] +) -> frozenset[str]: + names: Final = frozenset(name for name in map(_proxy_path_mcp_name, tools) if name is not None) + return await served_names(names) if names else frozenset[str]() + + class LiteLLM_Proxy_MCP_Handler: """ Helper class with static methods for MCP integration with Responses API. @@ -156,16 +166,26 @@ class LiteLLM_Proxy_MCP_Handler: tools: Iterable[Mapping[str, object]] | None, 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]() + items: Final = tuple(tools or ()) + served: Final = await _served_mcp_path_names(items, served_names) 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 + {**tool, "server_url": f"{LITELLM_PROXY_MCP_SERVER_URL_PREFIX}{name}"} + if (name := _proxy_path_mcp_name(tool)) in served + else tool + for tool in items ] ) + @staticmethod + async def _routes_through_gateway( + tools: Iterable[Mapping[str, object]] | None, + served_names: Callable[[Collection[str]], Awaitable[frozenset[str]]] = _gateway_served_names, + ) -> tuple[bool, ...]: + items: Final = tuple(tools or ()) + served: Final = await _served_mcp_path_names(items, served_names) + return tuple(_names_gateway_explicitly(tool) or _proxy_path_mcp_name(tool) in served for tool in items) + @staticmethod async def _apply_toolset_permissions( resolved_toolset_ids: list[str], diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py index ea64fdf2a5d..52b8fa3d849 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py @@ -1397,6 +1397,33 @@ async def test_scoped_list_denied_by_agent_binding_raises_403_naming_agent(): assert resolver.await_args_list[1].kwargs["mcp_servers"] == ["github"] +@pytest.mark.asyncio +async def test_empty_scope_lists_nothing_instead_of_raising_a_nameless_denial(): + """An empty ``x-mcp-servers`` header scopes to no servers; that is an empty listing, not a 403.""" + pytest.importorskip("litellm.proxy._experimental.mcp_server.server") + from litellm.proxy._experimental.mcp_server.server import _get_tools_from_mcp_servers + + user_api_key_auth = UserAPIKeyAuth(api_key="test_key", user_id="test_user") + resolver = AsyncMock(return_value=[]) + + with ( + patch( # test-quality-ok: the permission resolver is a module-level function; the suite's only seam + "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers", + resolver, + ), + patch( # test-quality-ok: the server registry is a module-level singleton; the suite's only seam + "litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager", + _denied_scope_manager({"github": "srv-github"}), + ), + ): + listing = await _get_tools_from_mcp_servers( + user_api_key_auth=user_api_key_auth, mcp_auth_header=None, mcp_servers=[] + ) + + assert listing.tools == [] + resolver.assert_awaited_once() + + @pytest.mark.asyncio async def test_scoped_list_denied_for_non_agent_key_raises_generic_403(): """A denial for a key with no agent binding stays generic and skips the agent-stripped rerun.""" diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_semantic_tool_filter.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_semantic_tool_filter.py index 054146d474d..7ea90dda971 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_semantic_tool_filter.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_semantic_tool_filter.py @@ -2190,3 +2190,26 @@ async def test_top_k_above_router_default_is_respected(): assert len(filtered) == 6 print("✅ Configured top_k above the semantic-router default of 5 is honored") + + +@pytest.mark.asyncio +async def test_semantic_filter_hook_narrows_only_references_the_gateway_serves(): + from litellm.proxy.hooks.mcp_semantic_filter import SemanticToolFilterHook + + gateway_reference = {"type": "mcp", "server_url": "litellm_proxy", "require_approval": "never"} + external_tool = { + "type": "mcp", + "server_label": "zapier", + "server_url": "https://mcp.zapier.com/api/mcp/mcp", + "allowed_tools": ["zapier_send_email"], + } + + async def served_names(names): + assert names == {"mcp"} + return frozenset() + + narrowed = await SemanticToolFilterHook._narrow_mcp_references( + [gateway_reference, external_tool], ["srv-tool_1"], served_names=served_names + ) + + assert narrowed == [{**gateway_reference, "allowed_tools": ["srv-tool_1"]}, external_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 73d245df651..4efe4d9bfd9 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 @@ -777,6 +777,37 @@ async def test_gateway_served_names_matches_alias_server_name_name_access_group_ assert served == {"github", "github-server", "github-name", "deepwiki", "prod-group", "my-toolset"} +@pytest.mark.asyncio +async def test_gateway_served_names_matches_server_id_short_prefix_and_alias_case_like_the_gateway(): + from litellm.proxy._experimental.mcp_server.utils import compute_short_server_prefix + from litellm.responses.mcp.litellm_proxy_mcp_handler import _gateway_served_names + + server_id = "0b9ae4ca-1bd2-4faa-b183-7dd812597e3b" + short_prefix = compute_short_server_prefix(server_id) + servers = (_registered(server_id, "github-name", alias="github", access_groups=["prod-group"]),) + + served = await _gateway_served_names( + {server_id, short_prefix, "GitHub", "PROD-GROUP", "nope"}, servers=lambda: servers, toolset_exists=_no_toolset + ) + + assert served == {server_id, short_prefix, "GitHub"} + + +@pytest.mark.asyncio +async def test_routes_through_gateway_flags_explicit_and_served_tools_only(): + served_tool = {"type": "mcp", "server_label": "github", "server_url": "http://localhost:4000/mcp/github"} + + async def served_names(names): + assert names == {"github", "mcp"} + return frozenset({"github"}) + + flags = await LiteLLM_Proxy_MCP_Handler._routes_through_gateway( + [ZAPIER_TOOL, EXPLICIT_GATEWAY_TOOL, served_tool, FUNCTION_TOOL], served_names=served_names + ) + + assert flags == (False, True, True, False) + + @pytest.mark.asyncio async def test_split_mcp_tools_leaves_external_mcp_path_urls_for_the_provider():