Skip to content

fix(mcp): route tool calls by server_id so unprefixed names cannot pick an arbitrary server #33665

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
tin-berri wants to merge 9 commits into from
Closed

fix(mcp): route tool calls by server_id so unprefixed names cannot pick an arbitrary server #33665

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
9 commits
Select commit Hold shift + click to select a range
7fe0d95
fix(mcp): route tool calls by server_id so unprefixed names cannot pi…
tin-berri Jul 17, 2026
85eaa03
fix(mcp): discriminate tool routes on a tag so reloads cannot skip th…
tin-berri Jul 17, 2026
2f7adb7
fix(mcp): fail closed when a tool's only owners are outside the calle…
tin-berri Jul 17, 2026
1da66a3
Merge remote-tracking branch 'origin/litellm_internal_staging' into l…
tin-berri Jul 18, 2026
cd3a922
fix(mcp): keep out-of-scope servers unreachable through the credentia…
tin-berri Jul 19, 2026
b0908d5
test(mcp): seed tool ownership from the specific server under test
tin-berri Jul 21, 2026
ad888fc
fix(mcp): make the ownership map the single authority when resolving …
tin-berri Jul 21, 2026
0353b24
fix(mcp): make tool-route registration authoritative per server
tin-berri Jul 25, 2026
e9fd8fc
fix(mcp): resolve responses-API tool dispatch by server_id and gate i…
tin-berri Jul 26, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
304 changes: 218 additions & 86 deletions litellm/proxy/_experimental/mcp_server/mcp_server_manager.py
View file
Open in desktop

Large diffs are not rendered by default.

41 changes: 39 additions & 2 deletions litellm/proxy/_experimental/mcp_server/server.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2611,7 +2611,33 @@ async def execute_mcp_tool(
original_tool_name = name
else:
# Resolve from tool name (MCP JSON-RPC or prefixed REST tool names).
mcp_server = global_mcp_server_manager._get_mcp_server_from_tool_name(name)
if requested_server is None:
route = global_mcp_server_manager.resolve_tool_route(
name,
allowed_server_ids=frozenset(server.server_id for server in allowed_mcp_servers),
)
if route.kind == "ambiguous":
candidates = sorted(
server.name
for server in (
global_mcp_server_manager.get_mcp_server_by_id(server_id) for server_id in route.server_ids
)
if server is not None
)
raise HTTPException(
status_code=409,
detail={
"error": "ambiguous_tool_name",
"message": (
f"Tool '{name}' is served by more than one MCP server "
f"({', '.join(candidates)}). Call it by its server-prefixed name to "
f"select one."
),
},
)
mcp_server = route.server if route.kind == "resolved" else None
else:
mcp_server = global_mcp_server_manager._get_mcp_server_from_tool_name(name)
Comment thread
cursor[bot] marked this conversation as resolved.
if mcp_server is None and requested_server is not None:
for known_prefix in iter_known_server_prefixes(requested_server):
candidate = global_mcp_server_manager._get_mcp_server_from_tool_name(
Expand Down Expand Up @@ -2664,8 +2690,13 @@ async def execute_mcp_tool(
litellm_logging_obj.model_call_details["model"] = f"MCP: {name}"
# Resolve the MCP server early so BYOK checks and credential injection
# apply to ALL dispatch paths (local tool registry AND managed MCP server).
# Stay within the caller's scope: a scope-blind lookup here would undo the
# scope decision already made above and dispatch to an out-of-scope server.
if mcp_server is None:
mcp_server = global_mcp_server_manager._get_mcp_server_from_tool_name(name)
allowed_server_ids = {server.server_id for server in allowed_mcp_servers}
fallback_server = global_mcp_server_manager._get_mcp_server_from_tool_name(name)
if fallback_server is not None and fallback_server.server_id in allowed_server_ids:
mcp_server = fallback_server

if mcp_server:
standard_logging_mcp_tool_call["mcp_server_cost_info"] = (mcp_server.mcp_info or {}).get(
Expand Down Expand Up @@ -2809,6 +2840,8 @@ async def execute_mcp_tool(
raw_headers=raw_headers,
litellm_logging_obj=litellm_logging_obj,
host_progress_callback=host_progress_callback,
resolved_server=mcp_server,
allowed_server_ids=frozenset(server.server_id for server in allowed_mcp_servers),
)

# Fall back to local tool registry with original name (legacy support)
Expand Down Expand Up @@ -3121,6 +3154,8 @@ async def _handle_managed_mcp_tool(
raw_headers: Optional[Dict[str, str]] = None,
litellm_logging_obj: Optional[Any] = None,
host_progress_callback: Optional[Callable] = None,
resolved_server: MCPServer | None = None,
allowed_server_ids: frozenset[str] | None = None,
) -> CallToolResult:
"""Handle tool execution for managed server tools"""
# Import here to avoid circular import
Expand All @@ -3137,6 +3172,8 @@ async def _handle_managed_mcp_tool(
raw_headers=raw_headers,
proxy_logging_obj=proxy_logging_obj,
host_progress_callback=host_progress_callback,
resolved_server=resolved_server,
allowed_server_ids=allowed_server_ids,
)
verbose_logger.debug("CALL TOOL RESULT: %s", call_tool_result)
return call_tool_result
Expand Down
25 changes: 14 additions & 11 deletions litellm/proxy/guardrails/guardrail_hooks/panw_prisma_airs/panw_prisma_airs.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -467,17 +467,20 @@ def _get_mcp_server_name(request_data: dict, mcp_tool_name: str) -> str:
)

server_id = request_data.get("server_id")
if server_id:
server = global_mcp_server_manager.get_mcp_server_by_id(server_id)
if server:
return (
getattr(server, "alias", None)
or getattr(server, "server_name", None)
or getattr(server, "name", None)
or getattr(server, "server_id", None)
or "unknown"
)
return global_mcp_server_manager.tool_name_to_mcp_server_name_mapping.get(mcp_tool_name, "unknown")
server = (
global_mcp_server_manager.get_mcp_server_by_id(server_id)
if server_id
else global_mcp_server_manager._get_mcp_server_from_tool_name(mcp_tool_name)
)
if server:
return (
getattr(server, "alias", None)
or getattr(server, "server_name", None)
or getattr(server, "name", None)
or getattr(server, "server_id", None)
or "unknown"
)
return "unknown"
except ImportError:
return "unknown"
except Exception:
Expand Down
68 changes: 46 additions & 22 deletions litellm/responses/mcp/litellm_proxy_mcp_handler.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -192,7 +192,7 @@ async def _get_mcp_tools_from_manager(

Returns:
List of MCP tools
List names of allowed MCP servers
List of server_ids for the MCP servers this caller may reach
"""
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
global_mcp_server_manager,
Expand Down Expand Up @@ -282,21 +282,18 @@ async def _get_mcp_tools_from_manager(
allowed_mcp_servers=allowed_mcp_servers,
)

server_names: List[str] = []
for server in allowed_mcp_servers:
if server is None:
continue
server_name = (
getattr(server, "server_name", None) or getattr(server, "alias", None) or getattr(server, "name", None)
)
if isinstance(server_name, str):
server_names.append(server_name)
# Carry server_id, not a display name. Names are not unique, so a name handed
# to the dispatch site gets re-resolved against the whole registry and can
# land on a server this caller cannot reach.
server_ids: List[str] = [
server.server_id for server in allowed_mcp_servers if server is not None and server.server_id
]

return tools, server_names
return tools, server_ids

@staticmethod
def _deduplicate_mcp_tools(
mcp_tools: List[MCPTool], allowed_mcp_servers: List[str]
mcp_tools: List[MCPTool], allowed_server_ids_list: List[str]
) -> tuple[List[MCPTool], dict[str, str]]:
"""
Deduplicate MCP tools by name, keeping the first occurrence of each tool.
Expand All @@ -306,11 +303,16 @@ def _deduplicate_mcp_tools(

Returns:
List of deduplicated MCP tools
The returned dictionary maps each tool_name to the server_name
The returned dictionary maps each tool_name to the owning server_id
"""
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
global_mcp_server_manager,
)

seen_names = set()
deduplicated_tools = []
tool_server_map: dict[str, str] = {}
allowed_server_ids = frozenset(allowed_server_ids_list)

for tool in mcp_tools:
if isinstance(tool, dict):
Expand All @@ -321,10 +323,16 @@ def _deduplicate_mcp_tools(
if tool_name and tool_name not in seen_names:
seen_names.add(tool_name)
deduplicated_tools.append(tool)
if len(allowed_mcp_servers) == 1:
tool_server_map[tool_name] = allowed_mcp_servers[0]
else:
_, tool_server_map[tool_name] = split_server_prefix_from_name(tool_name)
if len(allowed_server_ids_list) == 1:
tool_server_map[tool_name] = allowed_server_ids_list[0]
continue
# Same scoped resolver the MCP JSON-RPC surface uses, so a name two
# reachable servers share stays ambiguous here too instead of
# silently picking one. Unresolvable names are left out of the map
# and fail closed at dispatch.
route = global_mcp_server_manager.resolve_tool_route(tool_name, allowed_server_ids=allowed_server_ids)
if route.kind == "resolved":
tool_server_map[tool_name] = route.server.server_id

return deduplicated_tools, tool_server_map

Expand Down Expand Up @@ -658,11 +666,25 @@ async def _execute_tool_calls(
# Import here to avoid circular import
from litellm.proxy.proxy_server import proxy_logging_obj

server_name = tool_server_map[tool_name]

mcp_server = global_mcp_server_manager.get_mcp_server_by_name(
server_name
) or global_mcp_server_manager._get_mcp_server_from_tool_name(tool_name)
# tool_server_map carries the server_id resolved within this caller's
# reachable set at listing time, so look it up by identity. A name
# lookup here would walk the whole registry and could land on a
# same-named server the caller cannot reach.
server_id = tool_server_map.get(tool_name)
mcp_server = global_mcp_server_manager.get_mcp_server_by_id(server_id) if server_id else None
if mcp_server is None:
# Fail closed, and report it the way every other failure in this
# loop does so the model still sees a result for its tool call.
verbose_logger.warning(f"No reachable MCP server owns tool {tool_name}")
tool_results.append(
{
"tool_call_id": tool_call_id,
"result": f"Tool call failed: no MCP server reachable by this caller serves '{tool_name}'.",
"name": tool_name,
}
)
continue
server_name = mcp_server.name
resolved_tool_name = (
_resolve_display_name_to_original(tool_name, [mcp_server]) if mcp_server else tool_name
)
Expand Down Expand Up @@ -780,6 +802,8 @@ async def _execute_tool_calls(
oauth2_headers=oauth2_headers,
raw_headers=raw_headers,
proxy_logging_obj=proxy_logging_obj,
resolved_server=mcp_server,
Comment thread
veria-ai[bot] marked this conversation as resolved.
allowed_server_ids=frozenset(tool_server_map.values()),
)

if litellm_logging_obj:
Expand Down
59 changes: 34 additions & 25 deletions tests/mcp_tests/test_mcp_logging.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -104,7 +104,7 @@ def mock_client_constructor(*args, **kwargs):
litellm.callbacks = [test_logger]

# Initialize the tool mapping
await local_mcp_server_manager._initialize_tool_name_to_mcp_server_name_mapping()
await local_mcp_server_manager._initialize_tool_name_to_mcp_server_ids_mapping()

# Patch the global manager in both modules where it's used
with (
Expand All @@ -121,17 +121,20 @@ def mock_client_constructor(*args, **kwargs):
_set_authorized_user(local_mcp_server_manager.get_all_mcp_server_ids())

print(
"tool_name_to_mcp_server_name_mapping",
local_mcp_server_manager.tool_name_to_mcp_server_name_mapping,
"tool_name_to_mcp_server_ids_mapping",
local_mcp_server_manager.tool_name_to_mcp_server_ids_mapping,
)

# Manually add the tool mapping to ensure it's available (since mocking might not capture it properly)
local_mcp_server_manager.tool_name_to_mcp_server_name_mapping[
zapier_server = local_mcp_server_manager.get_mcp_server_by_name("zapier_gmail_server")
assert zapier_server is not None
zapier_server_ids = frozenset({zapier_server.server_id})
local_mcp_server_manager.tool_name_to_mcp_server_ids_mapping[
"add_tools"
] = "zapier_gmail_server"
local_mcp_server_manager.tool_name_to_mcp_server_name_mapping[
] = zapier_server_ids
local_mcp_server_manager.tool_name_to_mcp_server_ids_mapping[
"zapier_gmail_server-add_tools"
] = "zapier_gmail_server"
] = zapier_server_ids

# Call mcp tool
response = await mcp_server_tool_call(
Expand Down Expand Up @@ -237,21 +240,24 @@ def mock_client_constructor(*args, **kwargs):
litellm.callbacks = [test_logger]

# Initialize the tool mapping
await local_mcp_server_manager._initialize_tool_name_to_mcp_server_name_mapping()
await local_mcp_server_manager._initialize_tool_name_to_mcp_server_ids_mapping()

# Manually add the tool mapping to ensure it's available (since mocking might not capture it properly)
local_mcp_server_manager.tool_name_to_mcp_server_name_mapping[
test_server = local_mcp_server_manager.get_mcp_server_by_name("test_server")
assert test_server is not None
test_server_ids = frozenset({test_server.server_id})
local_mcp_server_manager.tool_name_to_mcp_server_ids_mapping[
"expensive_tool"
] = "test_server"
local_mcp_server_manager.tool_name_to_mcp_server_name_mapping[
] = test_server_ids
local_mcp_server_manager.tool_name_to_mcp_server_ids_mapping[
"test_server-expensive_tool"
] = "test_server"
local_mcp_server_manager.tool_name_to_mcp_server_name_mapping["cheap_tool"] = (
"test_server"
] = test_server_ids
local_mcp_server_manager.tool_name_to_mcp_server_ids_mapping["cheap_tool"] = (
test_server_ids
)
local_mcp_server_manager.tool_name_to_mcp_server_name_mapping[
local_mcp_server_manager.tool_name_to_mcp_server_ids_mapping[
"test_server-cheap_tool"
] = "test_server"
] = test_server_ids

# Patch the global manager in both modules where it's used
with (
Expand All @@ -268,8 +274,8 @@ def mock_client_constructor(*args, **kwargs):
_set_authorized_user(local_mcp_server_manager.get_all_mcp_server_ids())

print(
"tool_name_to_mcp_server_name_mapping",
local_mcp_server_manager.tool_name_to_mcp_server_name_mapping,
"tool_name_to_mcp_server_ids_mapping",
local_mcp_server_manager.tool_name_to_mcp_server_ids_mapping,
)

# Test 1: Call expensive_tool - should cost 5.0
Expand Down Expand Up @@ -401,15 +407,18 @@ def mock_client_constructor(*args, **kwargs):
litellm.callbacks = [test_logger]

# Initialize the tool mapping
await local_mcp_server_manager._initialize_tool_name_to_mcp_server_name_mapping()
await local_mcp_server_manager._initialize_tool_name_to_mcp_server_ids_mapping()

# Manually add the tool mapping to ensure it's available (since mocking might not capture it properly)
local_mcp_server_manager.tool_name_to_mcp_server_name_mapping["add_tools"] = (
"zapier_gmail_server"
zapier_server = local_mcp_server_manager.get_mcp_server_by_name("zapier_gmail_server")
assert zapier_server is not None
zapier_server_ids = frozenset({zapier_server.server_id})
local_mcp_server_manager.tool_name_to_mcp_server_ids_mapping["add_tools"] = (
zapier_server_ids
)
local_mcp_server_manager.tool_name_to_mcp_server_name_mapping[
local_mcp_server_manager.tool_name_to_mcp_server_ids_mapping[
"zapier_gmail_server-add_tools"
] = "zapier_gmail_server"
] = zapier_server_ids

# Patch the global manager in both modules where it's used
with (
Expand All @@ -426,8 +435,8 @@ def mock_client_constructor(*args, **kwargs):
_set_authorized_user(local_mcp_server_manager.get_all_mcp_server_ids())

print(
"tool_name_to_mcp_server_name_mapping",
local_mcp_server_manager.tool_name_to_mcp_server_name_mapping,
"tool_name_to_mcp_server_ids_mapping",
local_mcp_server_manager.tool_name_to_mcp_server_ids_mapping,
)

# Call mcp tool using the correct separator format (- not /)
Expand Down
Loading
Loading