From cf08c07fbbc2e3323e9a3e4376a9feec5c9929fc Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Fri, 17 Jul 2026 13:31:21 -0700 Subject: [PATCH] fix(mcp): key every caller-visible listing surface by the display prefix, never canonical names Outcome keys in the tools/list _meta, the spend-log outcome and count maps, and the REST error messages now all use get_server_prefix (alias, or the short prefix when that mode is enabled), the same naming the caller already sees on tool names. Keying them by canonical server_name let an authenticated caller enumerate internal server names and their health or auth state that the alias and short-prefix schemes deliberately hide (Veria finding). One helper decides the key for every surface; exception messages reaching the multi-server REST error list are mapped to their fault tag with the display prefix instead of relaying exception text carrying canonical names. Server-side logs keep the real names --- .../mcp_server/rest_endpoints.py | 9 ++++- .../proxy/_experimental/mcp_server/server.py | 11 +++-- .../test_mcp_oauth_passthrough_tools.py | 1 + .../mcp_server/test_mcp_server.py | 40 +++++++++++++++---- .../mcp_server/test_rest_endpoints.py | 3 +- 5 files changed, 47 insertions(+), 17 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py index 8ab2130cc70..94271c54f4b 100644 --- a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py @@ -32,6 +32,7 @@ from litellm.proxy._experimental.mcp_server.ui_session_utils import ( ) from litellm.proxy._experimental.mcp_server.utils import ( MCPMissingUserEnvVarsError, + get_server_prefix, merge_mcp_headers, ) from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth @@ -640,7 +641,7 @@ if MCP_AVAILABLE: status_code=list_fault_http_status(fault), detail={ "error": fault.tag, - "message": f"Failed to list tools from server {server.name}", + "message": f"Failed to list tools from server {get_server_prefix(server)}", }, ) from e except Exception as e: @@ -854,7 +855,11 @@ if MCP_AVAILABLE: list_tools_result.extend(tools_result) except Exception as e: verbose_logger.exception(f"Error getting tools from {server.name}: {e}") - errors.append(f"{server.name}: {str(e)}") + errors.append( + f"{get_server_prefix(server)}: {classify_list_exception(e).tag}" + if isinstance(e, (MCPServerListError, MCPUpstreamAuthError)) + else f"{get_server_prefix(server)}: {str(e)}" + ) continue if errors and not list_tools_result: diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 608d0d21177..a8ab0937124 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -1766,12 +1766,11 @@ if MCP_AVAILABLE: _mcp_gateway_server_name.reset(server_name_token) def _aggregate_server_key(server: MCPServer) -> str: - return str( - getattr(server, "server_name", None) - or getattr(server, "alias", None) - or getattr(server, "name", None) - or "unknown" - ) + """The client-visible key for a server in listing outcomes and spend metadata: the same + display prefix (alias, or the short prefix when that mode is enabled) the caller already + sees on the tool names. Canonical internal server names never key a caller-readable + surface; when the display naming deliberately hides them, the outcome keys must too.""" + return get_server_prefix(server) or "unknown" async def _get_tools_from_mcp_servers( user_api_key_auth: Optional[UserAPIKeyAuth], diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_oauth_passthrough_tools.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_oauth_passthrough_tools.py index 2d56680b64d..fdc77d19d73 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_oauth_passthrough_tools.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_oauth_passthrough_tools.py @@ -273,6 +273,7 @@ def _http_server(server_id: str, name: str, **kwargs) -> MCPServer: return MCPServer( server_id=server_id, name=name, + alias=name, url=f"https://{name}/mcp", transport=MCPTransport.http, **kwargs, diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py index 21e7cafe046..ae4f12fc1e1 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py @@ -1096,8 +1096,8 @@ async def test_get_tools_from_mcp_servers_continues_when_one_server_fails(): # Verify that tools from the working server are returned assert len(result.tools) == 1 assert result.tools[0].name == "working_tool_1" - assert result.outcomes["working_server"].tag == "ok" - assert result.outcomes["failing_server"].tag == "internal" + assert result.outcomes["working"].tag == "ok" + assert result.outcomes["failing"].tag == "internal" # Verify failure logging mock_logger.exception.assert_any_call( @@ -1191,8 +1191,8 @@ async def test_get_tools_from_mcp_servers_handles_all_servers_failing(): # Verify that empty list is returned assert len(result.tools) == 0 - assert result.outcomes["failing_server1"].tag == "internal" - assert result.outcomes["failing_server2"].tag == "internal" + assert result.outcomes["failing1"].tag == "internal" + assert result.outcomes["failing2"].tag == "internal" # Verify failure logging for both servers mock_logger.exception.assert_any_call( @@ -7512,10 +7512,34 @@ async def test_aggregate_listing_reports_per_server_outcomes(): ) assert [tool.name for tool in listing.tools] == ["working_tool_1"] - assert listing.outcomes["working_server"].tag == "ok" - assert listing.outcomes["working_server"].tool_count == 1 - assert listing.outcomes["broken_server"].tag == "upstream_error" - assert listing.outcomes["broken_server"].status_code == 500 + assert listing.outcomes["working"].tag == "ok" + assert listing.outcomes["working"].tool_count == 1 + assert listing.outcomes["broken"].tag == "upstream_error" + assert listing.outcomes["broken"].status_code == 500 + assert "working_server" not in listing.outcomes + assert "broken_server" not in listing.outcomes + + +@pytest.mark.asyncio +async def test_outcome_keys_use_display_prefix_never_canonical_names(): + """Outcome keys are client-visible and must use the same display naming (alias or short prefix) + the caller already sees on tool names: keying them by canonical server_name would let any + authenticated caller enumerate internal server names the alias scheme deliberately hides.""" + try: + from litellm.proxy._experimental.mcp_server.server import _aggregate_server_key + except ImportError: + pytest.skip("MCP server not available") + + server = MagicMock() + server.alias = "public-alias" + server.server_name = "internal-canonical-name" + server.name = "internal-canonical-name" + server.short_prefix = None + server.server_id = "srv-1" + + key = _aggregate_server_key(server) + assert key == "public-alias" + assert "internal-canonical-name" not in key @pytest.mark.asyncio diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py index 106584435b5..d4ba66c4381 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py @@ -966,7 +966,8 @@ class TestListToolsRestAPI: assert exc_info.value.status_code == 502 assert exc_info.value.detail["error"] == "upstream_error" - assert "flaky" in exc_info.value.detail["message"] + assert "server-1" in exc_info.value.detail["message"] + assert "flaky" not in exc_info.value.detail["message"] async def test_aggregate_list_absorbs_one_server_auth_failure(self, monkeypatch): """The multi-server aggregate listing degrades a server whose upstream