Merge pull request #14720 from uc4w6c/feat/remove-servername-prefix-mcp_tools

feat: remove server_name prefix from list_tools

Author: krrishdholakia

litellm/proxy/_experimental/mcp_server/mcp_server_manager.py

@@ -421,6 +421,7 @@ async def _get_tools_from_server(
         self,
         server: MCPServer,
         mcp_auth_header: Optional[str] = None,
+        add_prefix: bool = True,
     ) -> List[MCPTool]:
         """
         Helper method to get tools from a single MCP server with prefixed names.
@@ -445,9 +446,11 @@ async def _get_tools_from_server(
 
             tools = await self._fetch_tools_with_timeout(client, server.name)
 
-            prefixed_tools = self._create_prefixed_tools(tools, server)
+            prefixed_or_original_tools = self._create_prefixed_tools(
+                tools, server, add_prefix=add_prefix
+            )
 
-            return prefixed_tools
+            return prefixed_or_original_tools
 
         except Exception as e:
             verbose_logger.warning(
@@ -516,7 +519,7 @@ async def _list_tools_task():
             return []
 
     def _create_prefixed_tools(
-        self, tools: List[MCPTool], server: MCPServer
+        self, tools: List[MCPTool], server: MCPServer, add_prefix: bool = True
     ) -> List[MCPTool]:
         """
         Create prefixed tools and update tool mapping.
@@ -534,14 +537,16 @@ def _create_prefixed_tools(
         for tool in tools:
             prefixed_name = add_server_prefix_to_tool_name(tool.name, prefix)
 
-            prefixed_tool = MCPTool(
-                name=prefixed_name,
+            name_to_use = prefixed_name if add_prefix else tool.name
+
+            tool_obj = MCPTool(
+                name=name_to_use,
                 description=tool.description,
                 inputSchema=tool.inputSchema,
             )
-            prefixed_tools.append(prefixed_tool)
+            prefixed_tools.append(tool_obj)
 
-            # Update tool to server mapping with both original and prefixed names
+            # Update tool to server mapping for resolution (support both forms)
             self.tool_name_to_mcp_server_name_mapping[tool.name] = prefix
             self.tool_name_to_mcp_server_name_mapping[prefixed_name] = prefix
 

litellm/proxy/_experimental/mcp_server/rest_endpoints.py

@@ -73,6 +73,7 @@ async def _get_tools_for_single_server(server, server_auth_header):
         tools = await global_mcp_server_manager._get_tools_from_server(
             server=server,
             mcp_auth_header=server_auth_header,
+            add_prefix=False,
         )
         return _create_tool_response_objects(tools, server.mcp_info)
 

litellm/proxy/_experimental/mcp_server/server.py

@@ -384,6 +384,9 @@ async def _get_tools_from_mcp_servers(
                 allowed_mcp_servers=allowed_mcp_servers,
             )
 
+        # Decide whether to add prefix based on number of allowed servers
+        add_prefix = not (len(allowed_mcp_servers) == 1)
+
         # Get tools from each allowed server
         all_tools = []
         for server_id in allowed_mcp_servers:
@@ -406,6 +409,7 @@ async def _get_tools_from_mcp_servers(
                 tools = await global_mcp_server_manager._get_tools_from_server(
                     server=server,
                     mcp_auth_header=server_auth_header,
+                    add_prefix=add_prefix,
                 )
                 all_tools.extend(tools)
                 verbose_logger.debug(
@@ -637,27 +641,35 @@ def _get_mcp_servers_in_path(path: str) -> Optional[List[str]]:
         # Server names can contain slashes (e.g., "custom_solutions/user_123")
         mcp_path_match = re.match(r"^/mcp/([^?#]+)(?:\?.*)?(?:#.*)?$", path)
         if mcp_path_match:
-            servers_and_path = mcp_path_match.group(1)
-            
-            if servers_and_path:
-                # Check if it contains commas (comma-separated servers)
-                if ',' in servers_and_path:
-                    # For comma-separated, look for a path at the end
-                    # Common patterns: /tools, /chat/completions, etc.
-                    path_match = re.search(r'/([^/,]+(?:/[^/,]+)*)$', servers_and_path)
-                    if path_match:
-                        # Path found at the end, remove it from servers
-                        path_part = '/' + path_match.group(1)
-                        servers_part = servers_and_path[:-len(path_part)]
-                        mcp_servers_from_path = [s.strip() for s in servers_part.split(',') if s.strip()]
-                    else:
-                        # No path, just comma-separated servers
-                        mcp_servers_from_path = [s.strip() for s in servers_and_path.split(',') if s.strip()]
+            mcp_servers_str = mcp_path_match.group(1)
+            optional_path = mcp_path_match.group(2)
+
+            if mcp_servers_str:
+                # First, try to split by comma for comma-separated lists
+                if "," in mcp_servers_str:
+                    # For comma-separated lists, we need to handle the case where the last item
+                    # might include the path (e.g., "zapier,group1/tools" -> ["zapier", "group1/tools"])
+                    parts = [s.strip() for s in mcp_servers_str.split(",") if s.strip()]
+
+                    # If there's an optional path AND the last part contains a slash that matches the optional path,
+                    # remove the path portion from the last server name
+                    if optional_path and len(parts) > 0 and "/" in parts[-1]:
+                        last_part = parts[-1]
+                        # Check if the last part ends with the optional path
+                        if optional_path and last_part.endswith(
+                            optional_path.lstrip("/")
+                        ):
+                            # Remove the path portion from the last server name
+                            parts[-1] = last_part[: -len(optional_path.lstrip("/"))]
+
+                    mcp_servers_from_path = parts
                 else:
-                    # Single server case - use regex approach for server/path separation
-                    # This handles cases like "custom_solutions/user_123/chat/completions"
-                    # where we want to extract "custom_solutions/user_123" as the server name
-                    single_server_match = re.match(r"^([^/]+(?:/[^/]+)?)(?:/.*)?$", servers_and_path)
+                    # For single server, it might be just a name or contain slashes
+                    # We need to determine where the server name ends and the path begins
+                    # This is tricky - let's use the original logic but handle comma cases differently
+                    single_server_match = re.match(
+                        r"^([^/]+(?:/[^/]+)?)(?:/.*)?$", mcp_servers_str
+                    )
                     if single_server_match:
                         server_name = single_server_match.group(1)
                         mcp_servers_from_path = [server_name]

tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py

@@ -102,7 +102,7 @@ async def test_get_tools_from_mcp_servers_continues_when_one_server_fails():
         working_server if server_id == "working_server" else failing_server
     )
 
-    async def mock_get_tools_from_server(server, mcp_auth_header=None):
+    async def mock_get_tools_from_server(server, mcp_auth_header=None, add_prefix=True):
         if server.name == "working_server":
             # Working server returns tools
             tool1 = MagicMock()
@@ -184,7 +184,7 @@ async def test_get_tools_from_mcp_servers_handles_all_servers_failing():
         failing_server1 if server_id == "failing_server1" else failing_server2
     )
 
-    async def mock_get_tools_from_server(server, mcp_auth_header=None):
+    async def mock_get_tools_from_server(server, mcp_auth_header=None, add_prefix=True):
         # All servers fail
         raise Exception(f"Server {server.name} connection failed")
 
@@ -448,3 +448,121 @@ async def test_mcp_routing_with_conflicting_alias_and_group_name():
     assert (
         called_servers[0].server_id == specific_server.server_id
     ), "Should have contacted the specific server alias, not the group."
+
+
+@pytest.mark.asyncio
+async def test_list_tools_single_server_unprefixed_names():
+    """When only one MCP server is allowed, list tools should return unprefixed names."""
+    try:
+        from litellm.proxy._experimental.mcp_server.server import (
+            _get_tools_from_mcp_servers,
+            set_auth_context,
+        )
+    except ImportError:
+        pytest.skip("MCP server not available")
+
+    # Mock user auth
+    user_api_key_auth = UserAPIKeyAuth(api_key="test_key", user_id="test_user")
+    set_auth_context(user_api_key_auth)
+
+    # One allowed server
+    server = MagicMock()
+    server.server_id = "server1"
+    server.name = "Zapier MCP"
+    server.alias = "zapier"
+
+    # Mock manager: allow just one server and return a tool based on add_prefix flag
+    mock_manager = MagicMock()
+    mock_manager.get_allowed_mcp_servers = AsyncMock(return_value=["server1"])
+    mock_manager.get_mcp_server_by_id = (
+        lambda server_id: server if server_id == "server1" else None
+    )
+
+    async def mock_get_tools_from_server(
+        server, mcp_auth_header=None, add_prefix=False
+    ):
+        tool = MagicMock()
+        tool.name = f"{server.alias}-toolA" if add_prefix else "toolA"
+        tool.description = "desc"
+        tool.inputSchema = {}
+        return [tool]
+
+    mock_manager._get_tools_from_server = mock_get_tools_from_server
+
+    with patch(
+        "litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager",
+        mock_manager,
+    ):
+        tools = await _get_tools_from_mcp_servers(
+            user_api_key_auth=user_api_key_auth,
+            mcp_auth_header=None,
+            mcp_servers=None,
+            mcp_server_auth_headers=None,
+        )
+
+    # Should be unprefixed since only one server is allowed
+    assert len(tools) == 1
+    assert tools[0].name == "toolA"
+
+
+@pytest.mark.asyncio
+async def test_list_tools_multiple_servers_prefixed_names():
+    """When multiple MCP servers are allowed, list tools should return prefixed names."""
+    try:
+        from litellm.proxy._experimental.mcp_server.server import (
+            _get_tools_from_mcp_servers,
+            set_auth_context,
+        )
+    except ImportError:
+        pytest.skip("MCP server not available")
+
+    # Mock user auth
+    user_api_key_auth = UserAPIKeyAuth(api_key="test_key", user_id="test_user")
+    set_auth_context(user_api_key_auth)
+
+    # Two allowed servers
+    server1 = MagicMock()
+    server1.server_id = "server1"
+    server1.name = "Zapier MCP"
+    server1.alias = "zapier"
+
+    server2 = MagicMock()
+    server2.server_id = "server2"
+    server2.name = "Jira MCP"
+    server2.alias = "jira"
+
+    # Mock manager
+    mock_manager = MagicMock()
+    mock_manager.get_allowed_mcp_servers = AsyncMock(
+        return_value=["server1", "server2"]
+    )
+    mock_manager.get_mcp_server_by_id = (
+        lambda server_id: server1 if server_id == "server1" else server2
+    )
+
+    async def mock_get_tools_from_server(
+        server, mcp_auth_header=None, add_prefix=True
+    ):
+        tool = MagicMock()
+        # When multiple servers, add_prefix should be True -> prefixed names
+        tool.name = f"{server.alias}-toolA" if add_prefix else "toolA"
+        tool.description = "desc"
+        tool.inputSchema = {}
+        return [tool]
+
+    mock_manager._get_tools_from_server = mock_get_tools_from_server
+
+    with patch(
+        "litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager",
+        mock_manager,
+    ):
+        tools = await _get_tools_from_mcp_servers(
+            user_api_key_auth=user_api_key_auth,
+            mcp_auth_header=None,
+            mcp_servers=None,
+            mcp_server_auth_headers=None,
+        )
+
+    # Should be prefixed since multiple servers are allowed
+    names = sorted([t.name for t in tools])
+    assert names == ["jira-toolA", "zapier-toolA"]

tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py

@@ -420,6 +420,112 @@ async def mock_get_tools_from_server(server, mcp_auth_header=None):
         assert result["status"] == "healthy"
         assert result["tools_count"] == 1
 
+    @pytest.mark.asyncio
+    async def test_get_tools_from_server_add_prefix(self):
+        """Verify _get_tools_from_server respects add_prefix True/False."""
+        manager = MCPServerManager()
+
+        # Create a minimal server with alias used as prefix
+        server = MCPServer(
+            server_id="zapier",
+            name="zapier",
+            transport=MCPTransport.http,
+        )
+
+        # Mock client creation and fetching tools
+        manager._create_mcp_client = MagicMock(return_value=object())
+
+        # Tools returned upstream (unprefixed from provider)
+        upstream_tool = MagicMock()
+        upstream_tool.name = "send_email"
+        upstream_tool.description = "Send an email"
+        upstream_tool.inputSchema = {}
+
+        manager._fetch_tools_with_timeout = AsyncMock(return_value=[upstream_tool])
+
+        # Case 1: add_prefix=True (default for multi-server) -> expect prefixed
+        tools_prefixed = await manager._get_tools_from_server(server, add_prefix=True)
+        assert len(tools_prefixed) == 1
+        assert tools_prefixed[0].name == "zapier-send_email"
+
+        # Case 2: add_prefix=False (single-server) -> expect unprefixed
+        tools_unprefixed = await manager._get_tools_from_server(
+            server, add_prefix=False
+        )
+        assert len(tools_unprefixed) == 1
+        assert tools_unprefixed[0].name == "send_email"
+
+    def test_create_prefixed_tools_updates_mapping_for_both_forms(self):
+        """_create_prefixed_tools should populate mapping for prefixed and original names even when not adding prefix in output."""
+        manager = MCPServerManager()
+
+        server = MCPServer(
+            server_id="jira",
+            name="jira",
+            transport=MCPTransport.http,
+        )
+
+        # Input tools as would come from upstream
+        t1 = MagicMock()
+        t1.name = "create_issue"
+        t1.description = ""
+        t1.inputSchema = {}
+        t2 = MagicMock()
+        t2.name = "close_issue"
+        t2.description = ""
+        t2.inputSchema = {}
+
+        # Do not add prefix in returned objects
+        out_tools = manager._create_prefixed_tools([t1, t2], server, add_prefix=False)
+
+        # Returned names should be unprefixed
+        names = sorted([t.name for t in out_tools])
+        assert names == ["close_issue", "create_issue"]
+
+        # Mapping should include both original and prefixed names -> resolves calls either way
+        assert manager.tool_name_to_mcp_server_name_mapping["create_issue"] == "jira"
+        assert (
+            manager.tool_name_to_mcp_server_name_mapping["jira-create_issue"] == "jira"
+        )
+        assert manager.tool_name_to_mcp_server_name_mapping["close_issue"] == "jira"
+        assert (
+            manager.tool_name_to_mcp_server_name_mapping["jira-close_issue"] == "jira"
+        )
+
+    def test_get_mcp_server_from_tool_name_with_prefixed_and_unprefixed(self):
+        """After mapping is populated, manager resolves both prefixed and unprefixed tool names to the same server."""
+        manager = MCPServerManager()
+
+        server = MCPServer(
+            server_id="zapier",
+            name="zapier",
+            server_name="zapier",
+            transport=MCPTransport.http,
+        )
+
+        # Register server so resolution can find it
+        manager.registry = {server.server_id: server}
+
+        # Populate mapping (add_prefix value doesn't matter for mapping population)
+        base_tool = MagicMock()
+        base_tool.name = "create_zap"
+        base_tool.description = ""
+        base_tool.inputSchema = {}
+        _ = manager._create_prefixed_tools([base_tool], server, add_prefix=False)
+
+        # Unprefixed resolution
+        resolved_server_unpref = manager._get_mcp_server_from_tool_name("create_zap")
+        print(resolved_server_unpref)
+        assert resolved_server_unpref is not None
+        assert resolved_server_unpref.server_id == server.server_id
+
+        # Prefixed resolution
+        resolved_server_pref = manager._get_mcp_server_from_tool_name(
+            "zapier-create_zap"
+        )
+        assert resolved_server_pref is not None
+        assert resolved_server_pref.server_id == server.server_id
+
 
 if __name__ == "__main__":
     pytest.main([__file__])