tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-08-06 22:22:23 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: remove server_name prefix from list_tools when only one server is present (#14504)

Browse Source
This commit is contained in:
Yuta Saito
2025-09-19 18:29:43 +09:00
parent aa54994994
commit a8059f80f4
5 changed files with 255 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files
litellm/proxy/_experimental/mcp_server/mcp_server_manager.py
+12 -7
View File
@@ -419,6 +419,7 @@ class MCPServerManager:
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.
@@ -443,9 +444,11 @@ class MCPServerManager:
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(
@@ -514,7 +517,7 @@ class MCPServerManager:
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.
@@ -532,14 +535,16 @@ class MCPServerManager:
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
+1
View File
@@ -73,6 +73,7 @@ if MCP_AVAILABLE:
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
+16 -8
View File
@@ -384,6 +384,9 @@ if MCP_AVAILABLE:
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 @@ if MCP_AVAILABLE:
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(
@@ -639,29 +643,33 @@ if MCP_AVAILABLE:
if mcp_path_match:
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:
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]:
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('/')):
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('/'))]
parts[-1] = last_part[: -len(optional_path.lstrip("/"))]
mcp_servers_from_path = parts
else:
# 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)
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
+120
View File
@@ -448,3 +448,123 @@ 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, mcp_protocol_version=None, add_prefix=True
):
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,
mcp_protocol_version=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, mcp_protocol_version=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,
mcp_protocol_version=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
+106
View File
@@ -420,6 +420,112 @@ class TestMCPServerManager:
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__])