scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-08 03:08:45 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 114

refactor(mcp_server_manager.py): refactor to maintain 50 LOC requirement

Browse source
This commit is contained in:
Krrish Dholakia 2025-10-06 20:15:06 -07:00
parent 543e00a886
commit 12726b0c33
1 changed files with 27 additions and 8 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

35
litellm/proxy/_experimental/mcp_server/mcp_server_manager.py
View file

@ -775,14 +775,7 @@ class MCPServerManager:
raise ValueError(f"Tool {name} not found")
# Validate that the server from prefix matches the actual server (if prefix was used)
if server_name_from_prefix:
expected_prefix = get_server_prefix(mcp_server)
if normalize_server_name(server_name_from_prefix) != normalize_server_name(
expected_prefix
):
raise ValueError(
f"Tool {name} server prefix mismatch: expected {expected_prefix}, got {server_name_from_prefix}"
)
self._validate_server_prefix_match(name, server_name_from_prefix, mcp_server)
#########################################################
# Pre MCP Tool Call Hook
@ -956,6 +949,32 @@ class MCPServerManager:
return None
def _validate_server_prefix_match(
self,
tool_name: str,
server_name_from_prefix: Optional[str],
mcp_server: MCPServer,
) -> None:
"""
Validate that the server prefix from the tool name matches the actual server.
Args:
tool_name: Original tool name provided
server_name_from_prefix: Server name extracted from tool name prefix (if any)
mcp_server: The MCP server that was found for this tool
Raises:
ValueError: If the server prefix doesn't match the expected server
"""
if server_name_from_prefix:
expected_prefix = get_server_prefix(mcp_server)
if normalize_server_name(server_name_from_prefix) != normalize_server_name(
expected_prefix
):
raise ValueError(
f"Tool {tool_name} server prefix mismatch: expected {expected_prefix}, got {server_name_from_prefix}"
)
async def _add_mcp_servers_from_db_to_in_memory_registry(self):
from litellm.proxy._experimental.mcp_server.db import get_all_mcp_servers
from litellm.proxy.management_endpoints.mcp_management_endpoints import (