From 7a1880f000bb1ce6bb234beec7135e8d05bce741 Mon Sep 17 00:00:00 2001 From: moe-berri Date: Fri, 30 Oct 2026 01:06:29 -0700 Subject: [PATCH] fix(mcp): publish stable code references across clients --- litellm/proxy/mcp_server/README.md | 6 +- litellm/proxy/mcp_server/code_execution.py | 40 +++++++++- litellm/proxy/mcp_server/rest_endpoints.py | 11 ++- litellm/proxy/mcp_server/tool_search.py | 12 +-- .../integration/mcp/test_mcp_code_dispatch.py | 67 ++++++++++++++++ .../proxy/mcp_server/test_code_execution.py | 78 ++++++++++++++++++- 6 files changed, 195 insertions(+), 19 deletions(-) create mode 100644 tests/integration/mcp/test_mcp_code_dispatch.py diff --git a/litellm/proxy/mcp_server/README.md b/litellm/proxy/mcp_server/README.md index 6954e62446d..f2485cea549 100644 --- a/litellm/proxy/mcp_server/README.md +++ b/litellm/proxy/mcp_server/README.md @@ -11,7 +11,7 @@ Reconnect the MCP client or refresh its tool list. It should discover `litellm_e Intel macOS and Windows ARM installations need a manual source installation of `quickjs-ng` before enabling code mode. A missing runtime produces a startup error. The original `quickjs` package conflicts with `quickjs-ng`; use an environment without that package. These requirements do not affect the gateway when code mode is disabled. -The code tool accepts a `code` string. Discover a tool and its argument schema normally, then use its exact catalog name or the opaque ID returned by `search_tools` in `tools.call_tool`. Calls return the full MCP result, including `content`, `structuredContent` when provided, and `isError`. Use `text(value)` to return a JSON value to the model. For example, after discovering a tool named `example-add` that accepts `a` and `b`: +The code tool accepts a `code` string. Discover a tool and its argument schema normally, then copy the code-mode `tool_id` from its description or `search_tools` into `tools.call_tool`. Copy that reference unchanged: clients can rename the displayed tool. Exact gateway catalog names also work. Calls return the full MCP result, including `content`, `structuredContent` when provided, and `isError`. Use `text(value)` to return a JSON value to the model. For example, after discovering a tool named `example-add` that accepts `a` and `b`: ```javascript const result = await tools.call_tool({ @@ -31,7 +31,7 @@ The script itself and each nested tool call count toward request-rate limits. Th ## Tool compatibility -When code mode is enabled, tool listings include `_meta["litellm.ai/code_mode"]`; search and schema results include `code_mode`. `available` means the tool can be attempted through the JSON call bridge, not that the gateway has tested every behavior of that upstream. `unsupported` includes a reason. Existing tools remain callable directly. +When code mode is enabled, tool listings include `_meta["litellm.ai/code_mode"]`; search and schema results include `code_mode`. `available` means the tool can be attempted through the JSON call bridge, not that the gateway has tested every behavior of that upstream. `unsupported` includes a reason. Descriptions also expose the call reference or the unsupported reason, for clients that omit metadata. Existing tools remain callable directly. Task-only tools and nested orchestration are unsupported. Interactive continuations cannot be resumed inside a script. An unexpected interaction requirement returns an error instead of bypassing approval or replaying the script. Approving the outer code tool in an MCP client is not per-call approval of every nested action. @@ -51,7 +51,7 @@ Keep large identifiers as strings. Values crossing the JavaScript bridge must be Each script has a 120-second deadline, a 64 MiB JavaScript heap limit, and a maximum of 100 tool calls with eight in flight. Each proxy process admits four scripts at once; additional requests receive a busy error. Code is limited to 128 KiB, individual bridge messages to 2 MiB, and returned output to 256 KiB. Use direct calls for tools whose results or duration exceed these limits. -A script that exceeds a limit fails explicitly. The result includes `output`, `error`, and `calls_started`; the count records attempted tool calls, not confirmed upstream actions. Cancellation stops further dispatch and terminates the worker, but cannot undo a request already sent upstream. A failed script can therefore have partial effects. Check the upstream state before retrying a write; the gateway does not automatically retry whole scripts. +A script that exceeds a limit fails explicitly. The result includes `output`, `error`, and `calls_started`; the count records attempted tool calls, not confirmed upstream actions. When cancellation reaches the running request, it stops further dispatch and terminates the worker, but cannot undo a request already sent upstream. A client timeout or disconnect does not guarantee server-side cancellation; work can continue until the script deadline. A failed script can therefore have partial effects. Check the upstream state before retrying a write; the gateway does not automatically retry whole scripts. # Gateway catalog pagination diff --git a/litellm/proxy/mcp_server/code_execution.py b/litellm/proxy/mcp_server/code_execution.py index 4382a089eb6..1cb14bad64b 100644 --- a/litellm/proxy/mcp_server/code_execution.py +++ b/litellm/proxy/mcp_server/code_execution.py @@ -107,6 +107,28 @@ def compatibility(tool: Tool) -> Mapping[str, JsonValue]: ) +def _code_reference(tool: Tool) -> str: + from litellm.proxy.mcp_server.tool_search import mcp_proxy_tool_id + + try: + return mcp_proxy_tool_id(tool) + except TypeError: + return tool.name + + +def code_description(tool: Tool) -> str: + """Publish a reference clients can copy even when they rename the tool.""" + if not settings().enabled: + return tool.description or "" + support: Final = compatibility(tool) + hint: Final = ( + f"Code mode unsupported: {support['reason']}" + if support["status"] == "unsupported" + else f"Code mode tool_id: {json.dumps(_code_reference(tool))}." + ) + return f"{hint}\n\n{tool.description}" if tool.description else hint + + def code_tool() -> Tool: from litellm.proxy.mcp_server.code_runtime import CodeExecutionLimits @@ -115,7 +137,8 @@ def code_tool() -> Tool: name=CODE_TOOL_NAME, description=( "Run JavaScript to compose MCP tool calls. Use await tools.call_tool({tool_id, arguments}); " - "tool_id is an ID from search_tools or an exact tool name from the catalog. " + "Copy the exact Code mode tool_id from the tool description, or tool_id from search_tools. " + "Do not rename, normalize or prefix this reference. " "Each call returns the full MCP envelope: structuredContent, content and isError. " "Inspect isError before using results. Emit only needed results with text(value). " "No filesystem, network, imports, persistent state or interactive continuations. " @@ -135,7 +158,13 @@ def with_code_tool( if any(tool.name == CODE_TOOL_NAME for tool in tools): raise MCPError(code=-32602, message="MCP tool name conflicts with the reserved code execution tool") annotated: Final = [ - tool.model_copy(update={"meta": {**(tool.meta or {}), CODE_MODE_META: compatibility(tool)}}) for tool in tools + tool.model_copy( + update={ + "meta": {**(tool.meta or {}), CODE_MODE_META: compatibility(tool)}, + "description": code_description(tool), + } + ) + for tool in tools ] return [*annotated, *([code_tool()] if first_page else [])] @@ -281,7 +310,12 @@ async def _execute_script(arguments: Mapping[str, object], context: OperationCon return serialize_tool_result(error_result("Nested orchestration tools are unsupported")) selected: Final = await find_tool(tool_id) if selected is None: - return serialize_tool_result(error_result("Unknown, ambiguous or unauthorized tool")) + return serialize_tool_result( + error_result( + "Unknown, ambiguous or unauthorized tool. Copy the exact Code mode tool_id from " + "the tool description, or tool_id returned by search_tools." + ) + ) support: Final = compatibility(selected) if support["status"] == "unsupported": return serialize_tool_result(error_result(f"Tool is unsupported in code mode: {support['reason']}")) diff --git a/litellm/proxy/mcp_server/rest_endpoints.py b/litellm/proxy/mcp_server/rest_endpoints.py index 3f58d38979c..bcf2970c669 100644 --- a/litellm/proxy/mcp_server/rest_endpoints.py +++ b/litellm/proxy/mcp_server/rest_endpoints.py @@ -839,12 +839,11 @@ if MCP_AVAILABLE: from litellm.proxy.mcp_server.code_execution import CODE_TOOL_NAME from litellm.proxy.mcp_server.tool_search import with_mcp_proxy_identity - identified: Final = with_mcp_proxy_identity(tool, server.server_id) - if tool.name == CODE_TOOL_NAME: - return identified.model_copy( - update={"name": add_server_prefix_to_name(tool.name, get_server_prefix(server))} - ) - return identified + canonical: Final = tool.model_copy( + update={"name": add_server_prefix_to_name(tool.name, get_server_prefix(server))} + ) + identified: Final = with_mcp_proxy_identity(canonical, server.server_id) + return identified if tool.name == CODE_TOOL_NAME else identified.model_copy(update={"name": tool.name}) async def fetch_pinnable_tool_catalog( server: MCPServer, request: Request, user_api_key_dict: UserAPIKeyAuth diff --git a/litellm/proxy/mcp_server/tool_search.py b/litellm/proxy/mcp_server/tool_search.py index e7404dc6e55..ea05a34b07c 100644 --- a/litellm/proxy/mcp_server/tool_search.py +++ b/litellm/proxy/mcp_server/tool_search.py @@ -188,11 +188,11 @@ def mcp_tool_search_settings() -> MCPToolSearchSettings | ValidationError: def _tool_result(tool: Tool) -> ToolSearchResult: - from litellm.proxy.mcp_server.code_execution import compatibility, settings + from litellm.proxy.mcp_server.code_execution import code_description, compatibility, settings base: Final[ToolSearchResult] = { "name": tool.name, - "description": tool.description or "", + "description": code_description(tool), "inputSchema": tool.input_schema, } return {**base, "code_mode": compatibility(tool)} if settings().enabled else base @@ -228,12 +228,12 @@ def mcp_proxy_tool_id(tool: Tool) -> str: def _proxy_search_result(hit: MCPToolSearchHit) -> MCPProxySearchResult: - from litellm.proxy.mcp_server.code_execution import compatibility, settings + from litellm.proxy.mcp_server.code_execution import code_description, compatibility, settings base: Final[MCPProxySearchResult] = { "tool_id": mcp_proxy_tool_id(hit.tool), "name": hit.tool.name, - "description": hit.tool.description or "", + "description": code_description(hit.tool), } annotated: Final[MCPProxySearchResult] = ( {**base, "code_mode": compatibility(hit.tool)} if settings().enabled else base @@ -242,12 +242,12 @@ def _proxy_search_result(hit: MCPToolSearchHit) -> MCPProxySearchResult: def _proxy_schema_result(tool: Tool) -> MCPProxySchemaResult: - from litellm.proxy.mcp_server.code_execution import compatibility, settings + from litellm.proxy.mcp_server.code_execution import code_description, compatibility, settings base: Final[MCPProxySchemaResult] = { "tool_id": mcp_proxy_tool_id(tool), "name": tool.name, - "description": tool.description or "", + "description": code_description(tool), "inputSchema": tool.input_schema, } annotated: Final[MCPProxySchemaResult] = {**base, "code_mode": compatibility(tool)} if settings().enabled else base diff --git a/tests/integration/mcp/test_mcp_code_dispatch.py b/tests/integration/mcp/test_mcp_code_dispatch.py new file mode 100644 index 00000000000..369e02f8edc --- /dev/null +++ b/tests/integration/mcp/test_mcp_code_dispatch.py @@ -0,0 +1,67 @@ +import uuid +from collections.abc import Iterator +from pathlib import Path +from typing import Final + +import pytest +import yaml +from integration._support.client import JSON_OBJECT, Gateway, gateway_from_environment, object_value +from integration._support.mcp import McpCaller, mcp_peer, register_mcp, tool_calls +from integration._support.process import owned_proxy + + +@pytest.fixture(scope="module") +def code_gateway(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Gateway]: + directory: Final = tmp_path_factory.mktemp("mcp-code-dispatch") + configuration: Final = JSON_OBJECT.validate_python( + yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + ) + path: Final = directory / "config.yaml" + path.write_text( + yaml.safe_dump( + { + **configuration, + "litellm_settings": {**object_value(configuration["litellm_settings"]), "mcp_code_mode": True}, + } + ) + ) + with ( + gateway_from_environment() as existing, + owned_proxy(existing, directory, {}, config=path, workers=2) as gateway, + ): + yield gateway + + +def test_code_calls_chain_with_caller_permissions(code_gateway: Gateway) -> None: + with mcp_peer() as allowed, mcp_peer() as denied, code_gateway.scenario() as scenario: + alias: Final = "code" + uuid.uuid4().hex[:8] + denied_alias: Final = "denied" + uuid.uuid4().hex[:8] + identity: Final = register_mcp(scenario, allowed, alias) + register_mcp(scenario, denied, denied_alias) + key: Final = scenario.key(object_permission={"mcp_servers": [identity]}, max_parallel_requests=1) + caller: Final = McpCaller(code_gateway, key, "mcp") + assert "litellm_execute_code" in caller.list_tools().tools + allowed.drain() + denied.drain() + result: Final = caller.call( + "litellm_execute_code", + { + "code": f''' + const sum = await tools.call_tool({{tool_id: "{alias}-add", arguments: {{a: 17, b: 25}}}}); + text(await tools.call_tool({{tool_id: "{alias}-multiply", arguments: {{a: Number(sum.content[0].text), b: 3}}}})); + text(await tools.call_tool({{tool_id: "{denied_alias}-add", arguments: {{a: 1, b: 2}}}})); + ''' + }, + ) + assert result.ok and result.text is not None, result.raw + body: Final = JSON_OBJECT.validate_json(result.text) + output: Final = body["output"] + assert isinstance(output, list) and len(output) == 2, body + successful: Final = object_value(output[0]) + refused: Final = object_value(output[1]) + content: Final = successful["content"] + assert isinstance(content, list) and object_value(content[0])["text"] == "126", successful + assert successful.get("isError", False) is False + assert refused["isError"] is True + assert len(tool_calls(allowed.drain())) == 2 + assert tool_calls(denied.drain()) == () diff --git a/tests/unit/proxy/mcp_server/test_code_execution.py b/tests/unit/proxy/mcp_server/test_code_execution.py index f6d9629bd02..8d4d0109f3b 100644 --- a/tests/unit/proxy/mcp_server/test_code_execution.py +++ b/tests/unit/proxy/mcp_server/test_code_execution.py @@ -220,7 +220,9 @@ async def test_disabled_code_name_falls_through_to_existing_tool_routing(monkeyp def test_rest_keeps_nested_code_tool_visible_without_name_collision(monkeypatch: pytest.MonkeyPatch) -> None: - from litellm.proxy.mcp_server.rest_endpoints import _code_mode_rest_tool + from litellm.proxy.mcp_server.rest_endpoints import ( + _code_mode_rest_tool, # pyright: ignore[reportPrivateUsage] # exercise the REST wire-format boundary + ) from litellm.types.mcp import MCPTransport from litellm.types.mcp_server.mcp_server_manager import MCPServer @@ -232,8 +234,82 @@ def test_rest_keeps_nested_code_tool_visible_without_name_collision(monkeypatch: transport=MCPTransport.http, url="https://example.com/mcp", ) + from litellm.proxy.mcp_server.mcp_server_manager import global_mcp_server_manager + + monkeypatch.setitem(global_mcp_server_manager.registry, server.server_id, server) nested = _code_mode_rest_tool(Tool(name=CODE_TOOL_NAME, input_schema={}), server) listed = with_code_tool([nested]) assert [tool.name for tool in listed] == ["gateway-litellm_execute_code", CODE_TOOL_NAME] assert listed[0].meta is not None assert listed[0].meta[CODE_MODE_META]["status"] == "unsupported" + + +def test_visible_code_reference_survives_client_renaming_and_matches_rest(monkeypatch: pytest.MonkeyPatch) -> None: + from litellm.proxy.mcp_server.mcp_server_manager import global_mcp_server_manager + from litellm.proxy.mcp_server.rest_endpoints import ( + _code_mode_rest_tool, # pyright: ignore[reportPrivateUsage] # exercise the REST wire-format boundary + ) + from litellm.proxy.mcp_server.tool_search import mcp_proxy_tool_id, with_mcp_proxy_identity + from litellm.types.mcp import MCPTransport + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + monkeypatch.setattr(litellm, "mcp_code_mode", True) + server = MCPServer( + server_id="visible-code-reference", + name="math", + alias="math", + transport=MCPTransport.http, + url="https://example.com/mcp", + tool_name_to_display_name={"add": "Friendly Add"}, + ) + monkeypatch.setitem(global_mcp_server_manager.registry, server.server_id, server) + original_description = "Add two numbers." + raw = Tool(name="add", description=original_description, input_schema={}) + canonical = with_mcp_proxy_identity(raw.model_copy(update={"name": "math-add"}), server.server_id) + expected_id = mcp_proxy_tool_id(canonical) + displayed = canonical.model_copy(update={"name": "Friendly Add"}) + rest = _code_mode_rest_tool(raw, server) + assert rest.name == "add" + for tool in (canonical, displayed, rest): + assert mcp_proxy_tool_id(tool) == expected_id + advertised = with_code_tool([tool])[0] + renamed_by_client = advertised.model_copy(update={"name": "mcp__qa__math_add"}) + assert renamed_by_client.description == f'Code mode tool_id: "{expected_id}".\n\n{original_description}' + assert renamed_by_client.description == advertised.description + + +def test_search_shapes_publish_reference_and_unsupported_reason(monkeypatch: pytest.MonkeyPatch) -> None: + from litellm.proxy.mcp_server.tool_search import ( + MCPToolSearchHit, + _proxy_schema_result, # pyright: ignore[reportPrivateUsage] # exercise the search wire-format boundary + _proxy_search_result, # pyright: ignore[reportPrivateUsage] # exercise the search wire-format boundary + _tool_result, # pyright: ignore[reportPrivateUsage] # exercise the search wire-format boundary + mcp_proxy_tool_id, + with_mcp_proxy_identity, + ) + + monkeypatch.setattr(litellm, "mcp_code_mode", True) + tool = with_mcp_proxy_identity( + Tool(name="math-add", description="Original tool description", input_schema={}), "visible-code-reference" + ) + expected_prefix = f'Code mode tool_id: "{mcp_proxy_tool_id(tool)}".' + descriptions = ( + _tool_result(tool)["description"], + _proxy_search_result(MCPToolSearchHit(tool=tool))["description"], + _proxy_schema_result(tool)["description"], + ) + assert all(description.startswith(expected_prefix) for description in descriptions) + assert all(description.endswith("Original tool description") for description in descriptions) + monkeypatch.setattr( + litellm, + "mcp_code_mode", + { + "enabled": True, + "excluded_tools": {"math-add": "Human confirmation required"}, + }, + ) + unsupported = with_code_tool([tool])[0] + assert unsupported.description == "Code mode unsupported: Human confirmation required\n\nOriginal tool description" + assert _tool_result(tool)["description"] == unsupported.description + assert _proxy_search_result(MCPToolSearchHit(tool=tool))["description"] == unsupported.description + assert _proxy_schema_result(tool)["description"] == unsupported.description