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

* Fix UI MCP testing tool regression

* Fixed linting

Author: yuneng-jiang

litellm/proxy/_experimental/mcp_server/server.py

@@ -640,24 +640,31 @@ async def call_mcp_tool(
 
         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 @@ async def call_mcp_tool(
         # 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

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