tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-08-02 10:21:52 +00:00
Code Issues Packages Projects Releases Wiki Activity

[Fix] UI MCP Tool Test Regression (#16695)

Browse Source 
* Fix UI MCP testing tool regression

* Fixed linting
This commit is contained in:
yuneng-jiang
2025-11-15 15:31:04 -08:00
committed by GitHub
parent e325b1833d
commit 00ecb99f3e
2 changed files with 98 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files
litellm/proxy/_experimental/mcp_server/server.py
+27 -18
View File
@@ -640,24 +640,31 @@ if MCP_AVAILABLE:
allowed_mcp_servers = await _get_allowed_mcp_servers_from_mcp_server_names(
mcp_servers=mcp_servers,
allowed_mcp_servers=allowed_mcp_servers
allowed_mcp_servers=allowed_mcp_servers,
)
server_name: Optional[str]
if len(allowed_mcp_servers) == 1:
original_tool_name, server_name = name, allowed_mcp_servers[0].server_name
else:
# Remove prefix from tool name for logging and processing
original_tool_name, server_name = get_server_name_prefix_tool_mcp(name)
# Track resolved MCP server for both permission checks and dispatch
mcp_server: Optional[MCPServer] = None
if not server_name or not MCPRequestHandler.is_tool_allowed(
allowed_mcp_servers=[server.name for server in allowed_mcp_servers],
server_name=server_name,
):
raise HTTPException(
status_code=403,
detail=f"User not allowed to call this tool. Allowed MCP servers: {allowed_mcp_servers}",
)
# Remove prefix from tool name for logging and processing
original_tool_name, server_name = get_server_name_prefix_tool_mcp(name)
# If tool name is unprefixed, resolve its server so we can enforce permissions
if not server_name:
mcp_server = global_mcp_server_manager._get_mcp_server_from_tool_name(name)
if mcp_server:
server_name = mcp_server.name
# Only enforce server-level permissions when we can resolve a server
if server_name:
if not MCPRequestHandler.is_tool_allowed(
allowed_mcp_servers=[server.name for server in allowed_mcp_servers],
server_name=server_name,
):
raise HTTPException(
status_code=403,
detail=f"User not allowed to call this tool. Allowed MCP servers: {allowed_mcp_servers}",
)
standard_logging_mcp_tool_call: StandardLoggingMCPToolCall = (
_get_standard_logging_mcp_tool_call(
@@ -686,9 +693,11 @@ if MCP_AVAILABLE:
# Primary and recommended way to use external MCP servers
#########################################################
else:
mcp_server: Optional[
MCPServer
] = global_mcp_server_manager._get_mcp_server_from_tool_name(name)
# If we haven't already resolved the server, do it now for dispatch
if mcp_server is None:
mcp_server = global_mcp_server_manager._get_mcp_server_from_tool_name(
name
)
if mcp_server:
standard_logging_mcp_tool_call["mcp_server_cost_info"] = (
mcp_server.mcp_info or {}
tests/mcp_tests/test_mcp_server.py
+71 -1
View File
@@ -2584,8 +2584,78 @@ async def test_call_mcp_tool_uses_manager_permission_lookup():
assert result == expected_response
mock_get_allowed.assert_awaited_once()
assert mock_get_server.call_count == 2
# We call `_get_mcp_server_from_tool_name` multiple times:
# - for logging/metadata
# - for resolving the server during dispatch
# - and inside `call_tool` for guardrails/hooks
# The exact count isn't important, only that it is used.
assert mock_get_server.call_count >= 2
# First call should use the prefixed tool name
assert (
mock_get_server.call_args_list[0][0][0]
== f"{mock_server.name}/gmail_send_email"
)
@pytest.mark.asyncio
async def test_call_mcp_tool_resolves_unprefixed_tool_name_and_checks_permissions():
"""
Ensure `call_mcp_tool` correctly resolves the MCP server for an unprefixed tool
name and enforces server-level permissions using that resolved server.
"""
from litellm.proxy._experimental.mcp_server.server import (
call_mcp_tool,
global_mcp_server_manager,
)
mock_server = MCPServer(
server_id="server-123",
name="test_server",
alias="test_server",
server_name="test_server",
url="https://test-server.com/mcp",
transport=MCPTransport.http,
mcp_info={"server_name": "test_server"},
)
expected_response = [TextContent(type="text", text="ok")]
with patch.object(
global_mcp_server_manager,
"get_allowed_mcp_servers",
new_callable=AsyncMock,
) as mock_get_allowed, patch.object(
global_mcp_server_manager,
"get_mcp_servers_from_ids",
return_value=[mock_server],
), patch.object(
global_mcp_server_manager,
"_get_mcp_server_from_tool_name",
return_value=mock_server,
) as mock_get_server, patch(
"litellm.proxy._experimental.mcp_server.server.global_mcp_tool_registry"
) as mock_tool_registry, patch(
"litellm.proxy._experimental.mcp_server.server._handle_managed_mcp_tool",
new_callable=AsyncMock,
) as mock_handle_managed, patch(
"litellm.proxy._experimental.mcp_server.server.MCPRequestHandler.is_tool_allowed",
return_value=True,
) as mock_is_allowed:
mock_get_allowed.return_value = [mock_server.server_id]
mock_tool_registry.get_tool.return_value = None
mock_handle_managed.return_value = expected_response
# Call with UNPREFIXED tool name; server should be resolved via mapping
result = await call_mcp_tool(
name="gmail_send_email",
arguments={"body": "hello"},
mcp_servers=["test_server"],
)
assert result == expected_response
mock_get_allowed.assert_awaited_once()
# We should resolve the server at least once using the unprefixed name
assert mock_get_server.call_count >= 1
assert mock_get_server.call_args_list[0][0][0] == "gmail_send_email"
# Permissions check should be invoked with the resolved server name
mock_is_allowed.assert_called_once()