fix(mcp): publish stable code references across clients

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
 

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']}"))

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

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

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()) == ()

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