From 172b0354b628c899abb837a10005d0081154d302 Mon Sep 17 00:00:00 2001 From: moe-berri <322602897+moe-berri@users.noreply.github.com> Date: Fri, 30 Oct 2026 11:06:23 -0700 Subject: [PATCH] fix(mcp): preserve scoped REST tool names --- litellm/proxy/mcp_server/rest_endpoints.py | 15 ++++++--- .../integration/mcp/test_mcp_code_dispatch.py | 32 +++++++++++++++---- .../proxy/mcp_server/test_rest_endpoints.py | 6 ++++ 3 files changed, 43 insertions(+), 10 deletions(-) diff --git a/litellm/proxy/mcp_server/rest_endpoints.py b/litellm/proxy/mcp_server/rest_endpoints.py index 41a58d25e39..8531605f22a 100644 --- a/litellm/proxy/mcp_server/rest_endpoints.py +++ b/litellm/proxy/mcp_server/rest_endpoints.py @@ -777,6 +777,7 @@ if MCP_AVAILABLE: extra_headers: dict[str, str] | None = None, apply_tool_filters: bool = True, client_ip: str | None = None, + preserve_tool_names: bool = False, ): """Helper function to get tools for a single server. @@ -828,11 +829,16 @@ if MCP_AVAILABLE: from litellm.proxy.mcp_server.code_execution import settings as code_mode_settings + code_mode_enabled: Final = code_mode_settings().enabled code_tools: Final = ( - [_code_mode_rest_tool(tool, server) for tool in served_tools] - if code_mode_settings().enabled - else served_tools + [_code_mode_rest_tool(tool, server) for tool in served_tools] if code_mode_enabled else served_tools ) + if code_mode_enabled and preserve_tool_names: + scoped_tools: Final = [ + annotated.model_copy(update={"name": original.name}) + for annotated, original in zip(with_code_tool(code_tools, first_page=False), served_tools, strict=True) + ] + return _create_tool_response_objects(scoped_tools, server) return _create_tool_response_objects(code_tools, server) def _code_mode_rest_tool(tool: MCPTool, server: MCPServer) -> MCPTool: @@ -966,6 +972,7 @@ if MCP_AVAILABLE: extra_headers=user_oauth_extra_headers, apply_tool_filters=apply_tool_filters, client_ip=rest_client_ip, + preserve_tool_names=True, ) except MCPUpstreamAuthError: # Surface the upstream 401/403 to the caller so it can emit the @@ -992,7 +999,7 @@ if MCP_AVAILABLE: "message": f"Failed to get tools from server {server.name}: {e}", } return { - "tools": with_code_tool(list_tools_result, first_page=False), + "tools": list_tools_result, "error": None, "message": "Successfully retrieved tools", } diff --git a/tests/integration/mcp/test_mcp_code_dispatch.py b/tests/integration/mcp/test_mcp_code_dispatch.py index 6773d8980f7..6dbe5490ea1 100644 --- a/tests/integration/mcp/test_mcp_code_dispatch.py +++ b/tests/integration/mcp/test_mcp_code_dispatch.py @@ -98,18 +98,38 @@ def test_code_timeout_returns_error_and_releases_caller_slot(code_gateway: Gatew def test_rest_server_id_preserves_upstream_reserved_name(code_gateway: Gateway) -> None: with scripted_peer(echo_tool("litellm_execute_code")) as peer, code_gateway.scenario() as scenario: - identity: Final = register_mcp(scenario, peer, "reserved" + uuid.uuid4().hex[:8]) + alias: Final = "reserved" + uuid.uuid4().hex[:8] + identity: Final = register_mcp(scenario, peer, alias) key: Final = scenario.key(object_permission={"mcp_servers": [identity]}) + caller: Final = McpCaller(code_gateway, key, "rest") + scoped: Final = caller.list_tools(identity) + assert scoped.ok and scoped.tools == ("litellm_execute_code",), scoped.raw + aggregate: Final = caller.list_tools() + assert aggregate.ok and set(aggregate.tools) == {"litellm_execute_code", f"{alias}-litellm_execute_code"}, ( + aggregate.raw + ) + scoped_tools: Final = SEARCH_RESULTS.validate_python(JSON_OBJECT.validate_json(scoped.raw)["tools"]) + aggregate_tools: Final = SEARCH_RESULTS.validate_python(JSON_OBJECT.validate_json(aggregate.raw)["tools"]) + upstream_tool: Final = next(tool for tool in aggregate_tools if tool["name"] != "litellm_execute_code") + support: Final = object_value(object_value(scoped_tools[0]["_meta"])["litellm.ai/code_mode"]) + assert support["status"] == "unsupported" + assert scoped_tools[0]["_meta"] == upstream_tool["_meta"] arguments: Final = {"code": "qa_upstream_marker"} peer.drain() - upstream: Final = McpCaller(code_gateway, key, "rest").call("litellm_execute_code", arguments, identity) + upstream: Final = caller.call(scoped.tools[0], arguments, identity) assert upstream.ok and upstream.text is not None, upstream assert JSON_OBJECT.validate_json(upstream.text) == arguments + aggregate_name: Final = upstream_tool["name"] + assert isinstance(aggregate_name, str) + aggregate_upstream: Final = caller.call(aggregate_name, arguments, identity) + assert aggregate_upstream.ok and aggregate_upstream.text is not None, aggregate_upstream + assert JSON_OBJECT.validate_json(aggregate_upstream.text) == arguments observed: Final = tool_calls(peer.drain()) - assert len(observed) == 1 - request: Final = JSON_OBJECT.validate_python(observed[0]["body"]) - assert object_value(request["params"])["arguments"] == arguments - gateway_code: Final = McpCaller(code_gateway, key, "rest").call("litellm_execute_code", {"code": "text(42)"}) + assert len(observed) == 2 + for request in (JSON_OBJECT.validate_python(call["body"]) for call in observed): + assert object_value(request["params"])["arguments"] == arguments + assert object_value(request["params"])["name"] == "litellm_execute_code" + gateway_code: Final = caller.call("litellm_execute_code", {"code": "text(42)"}) assert gateway_code.ok and gateway_code.text is not None, gateway_code assert JSON_OBJECT.validate_json(gateway_code.text)["output"] == [42] assert tool_calls(peer.drain()) == () diff --git a/tests/unit/proxy/mcp_server/test_rest_endpoints.py b/tests/unit/proxy/mcp_server/test_rest_endpoints.py index fdb1ff3cd09..65b32351668 100644 --- a/tests/unit/proxy/mcp_server/test_rest_endpoints.py +++ b/tests/unit/proxy/mcp_server/test_rest_endpoints.py @@ -1345,6 +1345,7 @@ class TestListToolsRestAPI: extra_headers=None, apply_tool_filters=True, client_ip=None, + preserve_tool_names: bool = False, ): captured["called"] = True captured["server"] = server @@ -1432,6 +1433,7 @@ class TestListToolsRestAPI: extra_headers=None, apply_tool_filters=True, client_ip=None, + preserve_tool_names: bool = False, ): captured["user_api_key_auth"] = user_api_key_auth return ["tool-1"] @@ -2001,6 +2003,7 @@ class TestListToolsRestAPI: extra_headers=None, apply_tool_filters=True, client_ip=None, + preserve_tool_names: bool = False, ): captured["called"] = True captured["server_arg"] = server @@ -2138,6 +2141,7 @@ class TestListToolsRestAPI: extra_headers=None, apply_tool_filters=True, client_ip=None, + preserve_tool_names: bool = False, ): captured["called"] = True captured["server_arg"] = server @@ -2224,6 +2228,7 @@ class TestListToolsRestAPI: extra_headers=None, apply_tool_filters=True, client_ip=None, + preserve_tool_names: bool = False, ): return ["scoped-tool"] @@ -2432,6 +2437,7 @@ class TestListToolsRestAPI: extra_headers=None, apply_tool_filters=True, client_ip=None, + preserve_tool_names: bool = False, ): captured["server"] = server captured["auth_header"] = server_auth_header