From a30e1f6e3d718f2b20fd6f62073fd9c9a6b4ab7d Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 18 Aug 2026 16:41:40 -0700 Subject: [PATCH 1/2] fix(mcp): bind tool existence check to the selected server --- .../mcp_server/mcp_server_manager.py | 45 ++++++------ .../mcp_server/test_mcp_server.py | 2 +- .../mcp_server/test_mcp_server_manager.py | 72 +++++++++++++++++++ 3 files changed, 95 insertions(+), 24 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 7fff6c12fe0..9835c7f01ed 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -2321,23 +2321,30 @@ class MCPServerManager: openapi_key_prefix: Final = prefix_root + MCP_TOOL_PREFIX_SEPARATOR global_mcp_tool_registry.unregister_tools_with_prefix(openapi_key_prefix) - owned_raw: Final[set[str]] = set() - for p in iter_known_server_prefixes(server): - if p: - owned_raw.add(p) - if server.name: - owned_raw.add(server.name) + owned_normalized: Final = self._owned_mapping_values(server) - owned_normalized: Final = {normalize_server_name(x) for x in owned_raw} - - stale_mapping_keys: Final[list[str]] = [] - for tool_name, mapped_server in list(self.tool_name_to_mcp_server_name_mapping.items()): - if mapped_server in owned_raw or normalize_server_name(str(mapped_server)) in owned_normalized: - stale_mapping_keys.append(tool_name) + stale_mapping_keys: Final = tuple( + tool_name + for tool_name, mapped_server in self.tool_name_to_mcp_server_name_mapping.items() + if normalize_server_name(str(mapped_server)) in owned_normalized + ) for key in stale_mapping_keys: del self.tool_name_to_mcp_server_name_mapping[key] + def _owned_mapping_values(self, server: MCPServer) -> frozenset[str]: + return frozenset( + 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: + owned: Final = self._owned_mapping_values(server) + mapped_owners: Final = ( + self.tool_name_to_mcp_server_name_mapping.get(spelling) + for spelling in iter_known_tool_name_spellings(tool_name, server) + ) + return any(owner is not None and normalize_server_name(owner) in owned for owner in mapped_owners) + def remove_server(self, mcp_server: LiteLLM_MCPServerTable): """ Remove a server from the registry @@ -5463,13 +5470,8 @@ class MCPServerManager: if mcp_server is None: raise ValueError(f"Tool {name} not found") - if resolved_by_server_name_only: - tool_known: Final = ( - name in self.tool_name_to_mcp_server_name_mapping - or prefixed_tool_name in self.tool_name_to_mcp_server_name_mapping - ) - if not tool_known: - raise ValueError(f"Tool {name} not found") + 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 @@ -5847,10 +5849,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 ( - original_tool_name in self.tool_name_to_mcp_server_name_mapping - or tool_name in self.tool_name_to_mcp_server_name_mapping - ): + if matched_server is not None and self._server_exposes_tool(matched_server, original_tool_name): return matched_server return None 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 3392203dbab..20aa19b1d32 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 @@ -6406,7 +6406,7 @@ async def test_execute_mcp_tool_rest_server_id_injects_requested_server_credenti with ( patch.dict( mcp_module.global_mcp_server_manager.tool_name_to_mcp_server_name_mapping, - {"echo": collision_server.name}, + {"echo": collision_server.name, "echo_requested-echo": requested_server.name}, ), patch.object( mcp_module.global_mcp_server_manager, 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 cd1ef7320dd..043458f599b 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 @@ -4833,6 +4833,78 @@ class TestMCPServerManager: with pytest.raises(ValueError, match="Tool missing_tool not found"): manager._resolve_mcp_server_for_tool_call("github", "missing_tool") + @staticmethod + def _manager_with_deepwiki_and_huggingface() -> MCPServerManager: + manager = MCPServerManager() + deepwiki = MCPServer(server_id="deepwiki-id", name="deepwiki", server_name="deepwiki", transport=MCPTransport.http) + huggingface = MCPServer( + server_id="huggingface-id", name="huggingface", server_name="huggingface", transport=MCPTransport.http + ) + manager.registry = {"deepwiki-id": deepwiki, "huggingface-id": huggingface} + manager.tool_name_to_mcp_server_name_mapping.update( + { + "read_wiki_structure": "deepwiki", + "deepwiki-read_wiki_structure": "deepwiki", + "hub_repo_search": "huggingface", + "huggingface-hub_repo_search": "huggingface", + } + ) + return manager + + def test_resolve_mcp_server_for_tool_call_rejects_tool_exposed_only_by_another_server(self): + manager = self._manager_with_deepwiki_and_huggingface() + + with pytest.raises(ValueError, match="Tool read_wiki_structure not found"): + manager._resolve_mcp_server_for_tool_call("huggingface", "read_wiki_structure") + with pytest.raises(ValueError, match="Tool hub_repo_search not found"): + manager._resolve_mcp_server_for_tool_call("deepwiki", "hub_repo_search") + + assert manager._resolve_mcp_server_for_tool_call("deepwiki", "read_wiki_structure") is manager.registry["deepwiki-id"] + assert manager._resolve_mcp_server_for_tool_call("huggingface", "hub_repo_search") is manager.registry["huggingface-id"] + + def test_get_mcp_server_from_tool_name_rejects_other_servers_prefix(self): + manager = self._manager_with_deepwiki_and_huggingface() + + assert manager._get_mcp_server_from_tool_name("huggingface-read_wiki_structure") is None + assert manager._get_mcp_server_from_tool_name("deepwiki-hub_repo_search") is None + assert manager._get_mcp_server_from_tool_name("deepwiki-read_wiki_structure") is manager.registry["deepwiki-id"] + assert manager._get_mcp_server_from_tool_name("huggingface-hub_repo_search") is manager.registry["huggingface-id"] + + def test_resolve_mcp_server_for_tool_call_shared_bare_name_resolves_via_own_prefixed_spelling(self): + manager = MCPServerManager() + zapier = MCPServer(server_id="zapier-id", name="zapier", alias="zapier-alias", transport=MCPTransport.http) + other = MCPServer(server_id="other-id", name="other", server_name="other", transport=MCPTransport.http) + manager.registry = {"zapier-id": zapier, "other-id": other} + manager.tool_name_to_mcp_server_name_mapping.update( + { + "create_zap": "other", + "other-create_zap": "other", + "zapier-alias-create_zap": "zapier-alias", + } + ) + + assert manager._resolve_mcp_server_for_tool_call("zapier", "create_zap") is zapier + assert manager._resolve_mcp_server_for_tool_call("other", "create_zap") is other + + def test_remove_server_drops_only_its_own_tool_mapping_rows(self): + manager = self._manager_with_deepwiki_and_huggingface() + + manager.remove_server( + LiteLLM_MCPServerTable( + server_id="huggingface-id", + alias="huggingface", + url="https://huggingface.co/mcp", + transport=MCPTransport.http, + created_at=datetime.now(), + updated_at=datetime.now(), + ) + ) + + assert manager.tool_name_to_mcp_server_name_mapping == { + "read_wiki_structure": "deepwiki", + "deepwiki-read_wiki_structure": "deepwiki", + } + @pytest.mark.asyncio async def test_resolve_oauth2_headers_skipped_when_not_user_oauth(self): """Returns input headers unchanged when server does not need user OAuth.""" From 9018a9503724989ae430d8ca5f07a2c26f7c3dd0 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 18 Aug 2026 17:35:40 -0700 Subject: [PATCH 2/2] test(mcp): build fixture mapping state without in-place mutation --- .../mcp_server/test_mcp_server_manager.py | 26 ++++++++----------- 1 file changed, 11 insertions(+), 15 deletions(-) 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 043458f599b..14400ef6376 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 @@ -4841,14 +4841,12 @@ class TestMCPServerManager: server_id="huggingface-id", name="huggingface", server_name="huggingface", transport=MCPTransport.http ) manager.registry = {"deepwiki-id": deepwiki, "huggingface-id": huggingface} - manager.tool_name_to_mcp_server_name_mapping.update( - { - "read_wiki_structure": "deepwiki", - "deepwiki-read_wiki_structure": "deepwiki", - "hub_repo_search": "huggingface", - "huggingface-hub_repo_search": "huggingface", - } - ) + manager.tool_name_to_mcp_server_name_mapping = { + "read_wiki_structure": "deepwiki", + "deepwiki-read_wiki_structure": "deepwiki", + "hub_repo_search": "huggingface", + "huggingface-hub_repo_search": "huggingface", + } return manager def test_resolve_mcp_server_for_tool_call_rejects_tool_exposed_only_by_another_server(self): @@ -4875,13 +4873,11 @@ class TestMCPServerManager: zapier = MCPServer(server_id="zapier-id", name="zapier", alias="zapier-alias", transport=MCPTransport.http) other = MCPServer(server_id="other-id", name="other", server_name="other", transport=MCPTransport.http) manager.registry = {"zapier-id": zapier, "other-id": other} - manager.tool_name_to_mcp_server_name_mapping.update( - { - "create_zap": "other", - "other-create_zap": "other", - "zapier-alias-create_zap": "zapier-alias", - } - ) + manager.tool_name_to_mcp_server_name_mapping = { + "create_zap": "other", + "other-create_zap": "other", + "zapier-alias-create_zap": "zapier-alias", + } assert manager._resolve_mcp_server_for_tool_call("zapier", "create_zap") is zapier assert manager._resolve_mcp_server_for_tool_call("other", "create_zap") is other