From 00ecb99f3e4210a4d76c367f1d5d26d6b3cefc2b Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Sat, 15 Nov 2025 15:31:04 -0800 Subject: [PATCH] [Fix] UI MCP Tool Test Regression (#16695) * Fix UI MCP testing tool regression * Fixed linting --- .../proxy/_experimental/mcp_server/server.py | 45 +++++++----- tests/mcp_tests/test_mcp_server.py | 72 ++++++++++++++++++- 2 files changed, 98 insertions(+), 19 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index ce29d2d32e..019e55b910 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -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 {} diff --git a/tests/mcp_tests/test_mcp_server.py b/tests/mcp_tests/test_mcp_server.py index 953dc373ab..6e09e890ca 100644 --- a/tests/mcp_tests/test_mcp_server.py +++ b/tests/mcp_tests/test_mcp_server.py @@ -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()