mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-11 22:51:28 +00:00
fix(mcp): match gateway-served /mcp/<name> URLs the way the gateway resolves them and narrow only gateway references
This commit is contained in:
parent
f7838d7e9b
commit
2a1527c22f
6 changed files with 122 additions and 16 deletions
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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],
|
||||
|
|
|
|||
|
|
@ -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."""
|
||||
|
|
|
|||
|
|
@ -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]
|
||||
|
|
|
|||
|
|
@ -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():
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue