From a3e9ed34fe5b2e1e1d53bf97f43c6498dbf79b66 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 19 Sep 2026 19:29:33 -0700 Subject: [PATCH] fix(mcp): gate the pre-call listing per tool, not per server A tools/call on a cold worker listed the target server once and then never again, so a later caller whose credentials expose a wider upstream catalog got 404 for tools the first caller never had. Gate the pre-call listing on whether this worker already exposes the requested tool, so callers with different catalogs no longer mask each other. Removing the per-server guard also drops the empty-listing case that re-listed on every call. --- .../mcp_server/mcp_server_manager.py | 13 +++------- .../proxy/_experimental/mcp_server/server.py | 15 +++++++++--- .../mcp_server/test_mcp_server.py | 24 +++++++++++++++++++ .../mcp_server/test_mcp_server_manager.py | 22 +++++++++++++++-- 4 files changed, 59 insertions(+), 15 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 5bc793ff544..b293ab5a206 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -2867,7 +2867,7 @@ class MCPServerManager: normalize_server_name(value) for value in (*iter_known_server_prefixes(server), server.name) if value ) - def _server_exposes_tool(self, server: MCPServer, tool_name: str) -> bool: + def server_exposes_tool(self, server: MCPServer, tool_name: str) -> bool: owned: Final = self._owned_mapping_values(server) mapped_owners: Final = ( self.tool_name_to_mcp_server_name_mapping.get(spelling) @@ -2875,13 +2875,6 @@ class MCPServerManager: ) return any(owner is not None and normalize_server_name(owner) in owned for owner in mapped_owners) - def has_listed_tools(self, server: MCPServer) -> bool: - """True once this worker holds at least one tool row for ``server``.""" - owned: Final = self._owned_mapping_values(server) - return any( - normalize_server_name(owner) in owned for owner in self.tool_name_to_mcp_server_name_mapping.values() - ) - def _known_prefix_to_server(self) -> Mapping[str, MCPServer]: """Every prefix form a tool name may carry, keyed to its server; a form two servers share stays with the one registered first.""" @@ -6135,7 +6128,7 @@ class MCPServerManager: if mcp_server is None: raise ValueError(f"Tool {name} not found") - if resolved_by_server_name_only and not self._server_exposes_tool(mcp_server, name): + if resolved_by_server_name_only and not self.server_exposes_tool(mcp_server, name): raise ValueError(f"Tool {name} not found") return mcp_server @@ -6514,7 +6507,7 @@ class MCPServerManager: if matched is not None: matched_prefix, original_tool_name = matched matched_server: Final = prefix_to_server.get(matched_prefix) - if matched_server is not None and self._server_exposes_tool(matched_server, original_tool_name): + if matched_server is not None and self.server_exposes_tool(matched_server, original_tool_name): return matched_server return None diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index e9fb1c464ad..3a9bca926b0 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -2890,6 +2890,7 @@ if MCP_AVAILABLE: async def _list_tools_before_first_call( server: MCPServer | None, + tool_name: str, allowed_mcp_servers: list[MCPServer], user_api_key_auth: UserAPIKeyAuth | None, mcp_auth_header: str | None, @@ -2897,13 +2898,15 @@ if MCP_AVAILABLE: oauth2_headers: dict[str, str] | None, raw_headers: dict[str, str] | None, ) -> None: - """Fill this worker's tool rows for ``server`` with the caller's own credentials. + """List ``server`` with the caller's own credentials when it does not yet expose ``tool_name`` here. The startup fill skips a server whose upstream wants the caller's token, and mcp 2 no longer lists before an uncached tools/call, so a worker that has not served tools/list - would otherwise answer 404 for every tool on that server. + for this caller would otherwise answer 404 for a tool the caller can see. Gating on the + requested tool, not on any prior listing, keeps callers with different upstream catalogs + from masking each other. """ - if server is None or global_mcp_server_manager.has_listed_tools(server): + if server is None or global_mcp_server_manager.server_exposes_tool(server, tool_name): return if all(allowed.server_id != server.server_id for allowed in allowed_mcp_servers): return @@ -2984,8 +2987,14 @@ if MCP_AVAILABLE: if requested_server is not None and not name_is_prefixed else global_mcp_server_manager.server_owning_tool_name_prefix(name) ) + first_call_tool_name: Final = ( + name + if first_call_target is None or (requested_server is not None and not name_is_prefixed) + else strip_known_server_prefix(name, first_call_target) + ) await _list_tools_before_first_call( server=first_call_target, + tool_name=first_call_tool_name, allowed_mcp_servers=allowed_mcp_servers, user_api_key_auth=user_api_key_auth, mcp_auth_header=mcp_auth_header, 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 dd65968e41c..3668a06203c 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 @@ -7570,6 +7570,30 @@ async def test_execute_mcp_tool_does_not_relist_a_server_this_worker_already_lis assert worker.dispatched["name"] == "add" +@pytest.mark.asyncio +async def test_execute_mcp_tool_lists_a_tool_this_worker_has_not_yet_seen_on_a_listed_server(): + """A worker that already holds one of the server's tools must still list when a caller asks + for a different tool it has not cached, so callers with wider upstream catalogs are not 404ed.""" + from mcp.types import Tool as MCPTool + + from litellm.proxy._experimental.mcp_server import server as mcp_module + + server = _never_listed_passthrough_server() + with _worker_that_never_listed(server, upstream_tools=("add", "multiply")) as worker: + mcp_module.global_mcp_server_manager._create_prefixed_tools([MCPTool(name="add", inputSchema={})], server) + await mcp_module.execute_mcp_tool( + name="lazy_map-multiply", + arguments={"a": 1, "b": 2}, + allowed_mcp_servers=[server], + start_time=datetime.now(), + mcp_auth_header="Bearer caller-token", + ) + + assert worker.fetch_tools.await_count == 1 + assert worker.dispatched["server_name"] == "lazy_map" + assert worker.dispatched["name"] == "multiply" + + @pytest.mark.asyncio async def test_execute_mcp_tool_never_lists_a_server_the_caller_cannot_access(): from litellm.proxy._experimental.mcp_server import server as mcp_module diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py index 68600c5e662..9140ac61f1a 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -5570,13 +5570,31 @@ class TestMCPServerManager: assert manager._get_mcp_server_from_tool_name("lazy_map-add") is None assert manager.server_owning_tool_name_prefix("lazy_map-add") is server assert manager.server_owning_tool_name_prefix("someone_else-add") is None - assert manager.has_listed_tools(server) is False manager._create_prefixed_tools([MCPTool(name="add", inputSchema={})], server) - assert manager.has_listed_tools(server) is True assert manager._get_mcp_server_from_tool_name("lazy_map-add") is server + def test_server_exposes_tool_is_per_tool_not_per_server(self): + """A tool listed for the server does not make its unlisted siblings look exposed.""" + manager = MCPServerManager() + server = MCPServer( + server_id="lazy-map-2", + name="lazy_map", + server_name="lazy_map", + transport=MCPTransport.http, + auth_type=MCPAuth.true_passthrough, + ) + manager.registry = {server.server_id: server} + + assert manager.server_exposes_tool(server, "add") is False + + manager._create_prefixed_tools([MCPTool(name="add", inputSchema={})], server) + + assert manager.server_exposes_tool(server, "add") is True + assert manager.server_exposes_tool(server, "lazy_map-add") is True + assert manager.server_exposes_tool(server, "multiply") is False + def test_known_prefix_to_server_keeps_the_first_registered_owner_of_a_shared_prefix(self): manager = MCPServerManager() first = MCPServer(server_id="first-id", name="first", server_name="first", transport=MCPTransport.http)