fix(mcp): match gateway-served /mcp/<name> URLs the way the gateway resolves them and narrow only gateway references

This commit is contained in:
mateo-berri 2026-09-02 16:03:31 -07:00
parent f7838d7e9b
commit 2a1527c22f
6 changed files with 122 additions and 16 deletions

View file

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

View file

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

View file

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

View file

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

View file

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

View file

@ -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():