From 0353b2450948cfaa3b1f815b483d504afe20d5bd Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Sat, 25 Jul 2026 14:51:54 -0700 Subject: [PATCH] fix(mcp): make tool-route registration authoritative per server Registering a server's tool routes was union-only, so a routing row could only ever gain owners. Nothing withdrew a tool while its server stayed in the registry: `_cleanup_server_tool_routing_artifacts` withdraws a server's id from every row, but it only runs when the server leaves the registry. So an upstream that stopped exposing a tool left its owner pinned. A name then served by exactly one reachable server kept resolving as ambiguous and kept returning the 409, with no way back short of restarting the proxy. A server's `tools/list` result is its complete listing, so it is the truth rather than an increment. `_replace_server_tool_routes` replaces one server's rows instead of accumulating into them, withdrawing its id from a name it no longer serves and dropping the row once no owner is left. Rows still accumulate across servers, so a genuinely shared name stays ambiguous. Treating a listing as the truth is safe because `_fetch_tools_with_timeout` raises on every failure instead of returning an empty list, and caller-scoped narrowing (`check_allowed_or_banned_tools`, semantic filtering) runs downstream, so neither a failed listing nor one caller's filtered view can evict routes another caller needs. Eviction is the same operation with an empty set, so it delegates rather than keeping its own copy of the withdrawal arithmetic. An earlier cut of this carried a per-server reverse index to avoid scanning the map, which bought under a millisecond on a path that has just made a network round trip, in exchange for a second source of truth that can disagree with the first. It already had: pointing eviction at the index broke `test_update_server_eviction_clears_openapi_routing_artifacts`, which seeds the mapping directly and asserts eviction clears rows however they were written. The scan is exhaustive by construction, so that class is gone. The OpenAPI path replaces once the whole spec has parsed, so a mid-loop failure leaves the previous routes intact rather than committing a partial set that would withdraw operations still served. The startup warm-up no longer re-registers what listing already recorded, which would have re-added names outside the replace and reintroduced rows it cannot withdraw. The OpenAPI registration call site turned out to have no test coverage at all; deleting it left the suite green. It now has three, including the mid-parse failure case. --- .../mcp_server/mcp_server_manager.py | 102 +++++-- .../mcp_server/test_mcp_server_manager.py | 278 ++++++++++++++++-- 2 files changed, 327 insertions(+), 53 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 44336cc4ca4..7bcccd1841d 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -15,7 +15,7 @@ import re import time from contextlib import asynccontextmanager from dataclasses import dataclass -from typing import Any, AsyncIterator, Callable, Literal, Optional, Union, cast +from typing import Any, AsyncIterator, Callable, Collection, Literal, Optional, Union, cast from urllib.parse import urlparse import anyio @@ -1482,6 +1482,7 @@ class MCPServerManager: paths = spec.get("paths", {}) components = spec.get("components", {}) registered_count = 0 + route_names: set[str] = set() verbose_logger.debug(f"Processing {len(paths)} paths from OpenAPI spec") @@ -1524,13 +1525,18 @@ class MCPServerManager: handler=tool_func, ) - # Update tool name to server id mapping (for both prefixed and base names) - self._register_tool_route(base_tool_name, server.server_id) - self._register_tool_route(prefixed_tool_name, server.server_id) + # Collect tool name to server id routes (both prefixed and base names) + route_names.add(base_tool_name) + route_names.add(prefixed_tool_name) registered_count += 1 verbose_logger.debug(f"Registered OpenAPI tool: {prefixed_tool_name} for server {server.name}") + # Only once the whole spec parsed: a mid-loop raise must leave the + # previous routes intact rather than commit a partial set, which would + # withdraw routes for operations later in the spec that are still served. + self._replace_server_tool_routes(server.server_id, route_names) + verbose_logger.info(f"Successfully registered {registered_count} OpenAPI tools for server {server.name}") except Exception as e: @@ -1558,12 +1564,7 @@ class MCPServerManager: openapi_key_prefix = prefix_root + MCP_TOOL_PREFIX_SEPARATOR global_mcp_tool_registry.unregister_tools_with_prefix(openapi_key_prefix) - for tool_name in list(self.tool_name_to_mcp_server_ids_mapping): - remaining = self.tool_name_to_mcp_server_ids_mapping[tool_name] - {server.server_id} - if remaining: - self.tool_name_to_mcp_server_ids_mapping[tool_name] = remaining - else: - del self.tool_name_to_mcp_server_ids_mapping[tool_name] + self._replace_server_tool_routes(server.server_id, frozenset()) def remove_server(self, mcp_server: LiteLLM_MCPServerTable): """ @@ -3821,6 +3822,9 @@ class MCPServerManager: """ prefixed_tools = [] prefix = get_server_prefix(server) + # Server-level, so resolve once instead of per tool. + known_prefixes = tuple(iter_known_server_prefixes(server)) + route_names: set[str] = set() for tool in tools: tool_copy = tool.model_copy(deep=True) @@ -3834,13 +3838,17 @@ class MCPServerManager: tool_copy.name = name_to_use prefixed_tools.append(tool_copy) - # Register every known prefix form (alias, server_name, server_id, + # Collect every known prefix form (alias, server_name, server_id, # short ID) so call_tool can resolve regardless of which form a # caller / cached client is using. - self._register_tool_route(original_name, server.server_id) - for known_prefix in iter_known_server_prefixes(server): - qualified = add_server_prefix_to_name(original_name, known_prefix) - self._register_tool_route(qualified, server.server_id) + route_names.add(original_name) + for known_prefix in known_prefixes: + route_names.add(add_server_prefix_to_name(original_name, known_prefix)) + + # `tools` is this server's complete upstream listing, so it is authoritative: + # replace rather than accumulate, or a tool the upstream dropped keeps its + # owner and makes the name look ambiguous until restart. + self._replace_server_tool_routes(server.server_id, route_names) verbose_logger.info(f"Successfully fetched {len(prefixed_tools)} tools from server {server.name}") return prefixed_tools @@ -4818,15 +4826,21 @@ class MCPServerManager: async def _initialize_tool_name_to_mcp_server_ids_mapping(self): """ - Call list_tools for each server and update the tool name to MCP server name mapping - Note: This now handles prefixed tool names + Warm the tool-name routing map by listing every server once at startup. + + Listing is what registers the routes: ``_get_tools_from_server`` hands the + server's complete tool list to ``_create_prefixed_tools``, which replaces + that server's routes via ``_replace_server_tool_routes``. OpenAPI servers + registered theirs from the spec in ``_register_openapi_tools``. So nothing is + registered here — doing so would re-add names outside the per-server replace + and reintroduce rows it cannot withdraw. """ for server in self.get_registry().values(): if server.needs_user_oauth_token: # Skip OAuth2 servers that rely on user-provided tokens continue try: - tools = await self._get_tools_from_server(server) + await self._get_tools_from_server(server) except MCPUpstreamAuthError as e: # Pass-through servers expect a user-supplied bearer token; # at startup we have none, so an upstream 401 is normal. @@ -4840,22 +4854,48 @@ class MCPServerManager: f"Failed to get tools from server {server.name} during tool name mapping initialization: {str(e)}" ) continue - for tool in tools: - # The tool.name here is already prefixed from _get_tools_from_server - # Extract original name for mapping - original_name, _ = split_server_prefix_from_name(tool.name) - self._register_tool_route(original_name, server.server_id) - self._register_tool_route(tool.name, server.server_id) - def _register_tool_route(self, tool_name: str, server_id: str) -> None: - """Record that ``server_id`` serves ``tool_name``, preserving any other owners. + def _replace_server_tool_routes(self, server_id: str, route_names: Collection[str]) -> None: + """Make ``route_names`` the complete set of routing rows owned by ``server_id``. - Routing rows accumulate rather than overwrite so that a tool name served by - more than one server stays resolvable as ambiguous instead of silently - collapsing to whichever server registered last. + Rows accumulate *across* servers, so a name several serve stays ambiguous. But + within one server this is authoritative rather than additive: a name it no + longer serves has ``server_id`` withdrawn and the row goes once unowned. + Additive registration pinned owners forever, so a dropped upstream tool kept + its name looking ambiguous, a false 409, until restart. Empty ``route_names`` + means owning nothing, which is how eviction withdraws a departing server. + + Callers must pass a *complete* listing, and both do. Treating one as the truth + is safe because ``_fetch_tools_with_timeout`` raises on every failure instead + of returning an empty list, and caller-scoped narrowing + (``check_allowed_or_banned_tools``, semantic filtering) runs downstream, so + neither a failed listing nor one caller's filtered view can evict another's + routes. Scanning every row beats a per-server index here: it runs once per + server per listing, right after that server's network round trip, so well + under a millisecond amortizes into tens of milliseconds of I/O, and a second + source of truth could disagree with this one. + + Pre-existing and unchanged: ``ClientSession.list_tools()`` is called without a + cursor, so a paginating upstream only yields page one. Every path populating + this map shares that listing, so replacing cannot drop a legitimate row. """ - owners = self.tool_name_to_mcp_server_ids_mapping.get(tool_name, frozenset()) - self.tool_name_to_mcp_server_ids_mapping[tool_name] = owners | {server_id} + desired = frozenset(route_names) + + # Mutated in place: the startup listing task is dispatched without being + # awaited and holds a reference to this dict, so rebinding drops its writes. + for tool_name in list(self.tool_name_to_mcp_server_ids_mapping): + owners = self.tool_name_to_mcp_server_ids_mapping[tool_name] + if server_id not in owners or tool_name in desired: + continue + remaining = owners - {server_id} + if remaining: + self.tool_name_to_mcp_server_ids_mapping[tool_name] = remaining + else: + del self.tool_name_to_mcp_server_ids_mapping[tool_name] + + for tool_name in desired: + owners = self.tool_name_to_mcp_server_ids_mapping.get(tool_name, frozenset()) + self.tool_name_to_mcp_server_ids_mapping[tool_name] = owners | {server_id} def resolve_tool_route( self, 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 0fdb6be4b44..f9de88610d7 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 @@ -3562,6 +3562,159 @@ class TestMCPServerManager: assert captured["headers"] is not None assert captured["headers"]["Authorization"] == "STATIC token" + @staticmethod + def _openapi_spec_with(operation_ids: list[str]) -> str: + return json.dumps( + { + "openapi": "3.0.0", + "info": {"title": "Demo", "version": "1.0.0"}, + "paths": {f"/{op}": {"get": {"operationId": op, "summary": op}} for op in operation_ids}, + } + ) + + @staticmethod + def _openapi_patches(): + async def tool_func(**kwargs): + return "ok" + + return ( + patch( + "litellm.proxy._experimental.mcp_server.openapi_to_mcp_generator.create_tool_function", + return_value=tool_func, + ), + patch( + "litellm.proxy._experimental.mcp_server.openapi_to_mcp_generator.build_input_schema", + return_value={"type": "object", "properties": {}, "required": []}, + ), + patch( + "litellm.proxy._experimental.mcp_server.tool_registry.global_mcp_tool_registry.register_tool", + return_value=None, + ), + ) + + @pytest.mark.asyncio + async def test_register_openapi_tools_records_base_and_prefixed_routes(self, tmp_path): + """OpenAPI operations must be routable by both their base and prefixed names. + + Nothing else covered this call site: deleting the registration entirely left + the whole mcp_server suite green. + """ + manager = MCPServerManager() + spec_path = tmp_path / "openapi.json" + spec_path.write_text(self._openapi_spec_with(["health_check"])) + server = MCPServer( + server_id="openapi-server", + name="openapi-server", + server_name="openapi-server", + url="https://example.com", + transport=MCPTransport.http, + auth_type=MCPAuth.none, + ) + + create_fn, build_schema, register = self._openapi_patches() + with create_fn, build_schema, register: + await manager._register_openapi_tools( + spec_path=str(spec_path), server=server, base_url="https://example.com" + ) + + assert manager.tool_name_to_mcp_server_ids_mapping["health_check"] == frozenset({"openapi-server"}) + assert manager.tool_name_to_mcp_server_ids_mapping["openapi-server-health_check"] == frozenset( + {"openapi-server"} + ) + + @pytest.mark.asyncio + async def test_reregistering_openapi_tools_withdraws_a_dropped_operation(self, tmp_path): + """Re-parsing a spec that lost an operation must withdraw that route. + + The spec is the server's complete listing, so registration replaces rather + than accumulates here for the same reason it does on the tools/list path. + """ + manager = MCPServerManager() + spec_path = tmp_path / "openapi.json" + server = MCPServer( + server_id="openapi-server", + name="openapi-server", + server_name="openapi-server", + url="https://example.com", + transport=MCPTransport.http, + auth_type=MCPAuth.none, + ) + + spec_path.write_text(self._openapi_spec_with(["health_check", "legacy_probe"])) + create_fn, build_schema, register = self._openapi_patches() + with create_fn, build_schema, register: + await manager._register_openapi_tools( + spec_path=str(spec_path), server=server, base_url="https://example.com" + ) + assert "legacy_probe" in manager.tool_name_to_mcp_server_ids_mapping + + spec_path.write_text(self._openapi_spec_with(["health_check"])) + create_fn, build_schema, register = self._openapi_patches() + with create_fn, build_schema, register: + await manager._register_openapi_tools( + spec_path=str(spec_path), server=server, base_url="https://example.com" + ) + + assert "legacy_probe" not in manager.tool_name_to_mcp_server_ids_mapping + assert "openapi-server-legacy_probe" not in manager.tool_name_to_mcp_server_ids_mapping + assert manager.tool_name_to_mcp_server_ids_mapping["health_check"] == frozenset({"openapi-server"}) + + @pytest.mark.asyncio + async def test_failed_openapi_reregistration_keeps_the_previous_routes(self, tmp_path): + """A mid-parse failure must not commit a partial set of routes. + + Replacing per operation would withdraw the operations after the failure point + even though the server still serves them. + """ + manager = MCPServerManager() + spec_path = tmp_path / "openapi.json" + server = MCPServer( + server_id="openapi-server", + name="openapi-server", + server_name="openapi-server", + url="https://example.com", + transport=MCPTransport.http, + auth_type=MCPAuth.none, + ) + + spec_path.write_text(self._openapi_spec_with(["health_check", "legacy_probe"])) + create_fn, build_schema, register = self._openapi_patches() + with create_fn, build_schema, register: + await manager._register_openapi_tools( + spec_path=str(spec_path), server=server, base_url="https://example.com" + ) + + async def tool_func(**kwargs): + return "ok" + + # Fail on the *second* operation, so the first has already been collected. + # Replacing per operation instead of after the loop would commit that partial + # set here and withdraw legacy_probe, which the server still serves. + calls = iter([tool_func, RuntimeError("spec blew up")]) + + def _second_call_explodes(*args, **kwargs): + outcome = next(calls) + if isinstance(outcome, Exception): + raise outcome + return outcome + + _, build_schema, register = self._openapi_patches() + with ( + patch( + "litellm.proxy._experimental.mcp_server.openapi_to_mcp_generator.create_tool_function", + side_effect=_second_call_explodes, + ), + build_schema, + register, + pytest.raises(RuntimeError), + ): + await manager._register_openapi_tools( + spec_path=str(spec_path), server=server, base_url="https://example.com" + ) + + assert manager.tool_name_to_mcp_server_ids_mapping["health_check"] == frozenset({"openapi-server"}) + assert manager.tool_name_to_mcp_server_ids_mapping["legacy_probe"] == frozenset({"openapi-server"}) + @pytest.mark.asyncio async def test_pre_call_tool_check_allowed_tools_list_allows_tool(self): """Test pre_call_tool_check allows tool when it's in allowed_tools list""" @@ -3897,30 +4050,65 @@ class TestMCPServerManager: resolved = manager._resolve_mcp_server_for_tool_call("zapier-alias", "create_zap") assert resolved is server - def test_register_tool_route_accumulates_owners_across_servers(self): + def test_replace_server_tool_routes_accumulates_owners_across_servers(self): """Two servers exposing one tool name are both recorded, not last-writer-wins.""" manager = MCPServerManager() - manager._register_tool_route("echo", "id-alpha") - manager._register_tool_route("echo", "id-zulu") + manager._replace_server_tool_routes("id-alpha", {"echo"}) + manager._replace_server_tool_routes("id-zulu", {"echo"}) assert manager.tool_name_to_mcp_server_ids_mapping["echo"] == frozenset({"id-alpha", "id-zulu"}) - def test_register_tool_route_is_idempotent_for_one_server(self): + def test_replace_server_tool_routes_is_idempotent_for_one_server(self): """Re-listing the same server must not make its own tool look ambiguous.""" manager = MCPServerManager() - manager._register_tool_route("echo", "id-alpha") - manager._register_tool_route("echo", "id-alpha") + manager._replace_server_tool_routes("id-alpha", {"echo"}) + manager._replace_server_tool_routes("id-alpha", {"echo"}) assert manager.tool_name_to_mcp_server_ids_mapping["echo"] == frozenset({"id-alpha"}) + def test_replace_server_tool_routes_withdraws_a_tool_the_upstream_dropped(self): + """A re-listing without a previously served tool must withdraw only that owner. + + Union-only registration would keep id-alpha pinned to "echo" forever, so a + name now served by id-zulu alone would keep resolving as ambiguous (409) + until the proxy restarted. + """ + manager = MCPServerManager() + manager._replace_server_tool_routes("id-alpha", {"echo", "still_there"}) + manager._replace_server_tool_routes("id-zulu", {"echo"}) + + # id-alpha's upstream stops exposing "echo"; "still_there" is unaffected. + manager._replace_server_tool_routes("id-alpha", {"still_there"}) + + assert manager.tool_name_to_mcp_server_ids_mapping["echo"] == frozenset({"id-zulu"}) + assert manager.tool_name_to_mcp_server_ids_mapping["still_there"] == frozenset({"id-alpha"}) + + def test_replace_server_tool_routes_drops_a_row_whose_last_owner_dropped_it(self): + """The row disappears, rather than lingering as an empty owner set.""" + manager = MCPServerManager() + manager._replace_server_tool_routes("id-alpha", {"echo"}) + manager._replace_server_tool_routes("id-alpha", set()) + + assert "echo" not in manager.tool_name_to_mcp_server_ids_mapping + + def test_replace_server_tool_routes_leaves_other_servers_rows_untouched(self): + """Replacing one server's routes must not disturb a row it never owned.""" + manager = MCPServerManager() + manager._replace_server_tool_routes("id-zulu", {"zulu_only"}) + manager._replace_server_tool_routes("id-alpha", {"alpha_only"}) + + manager._replace_server_tool_routes("id-alpha", set()) + + assert manager.tool_name_to_mcp_server_ids_mapping["zulu_only"] == frozenset({"id-zulu"}) + def test_get_mcp_server_from_tool_name_refuses_ambiguous_unprefixed_name(self): """An unprefixed name owned by several servers resolves to nothing, never to one of them.""" manager = MCPServerManager() alpha = MCPServer(server_id="id-alpha", name="echo_alpha", transport=MCPTransport.http) zulu = MCPServer(server_id="id-zulu", name="echo_zulu", transport=MCPTransport.http) manager.registry = {"id-alpha": alpha, "id-zulu": zulu} - manager._register_tool_route("echo", "id-alpha") - manager._register_tool_route("echo", "id-zulu") + manager._replace_server_tool_routes("id-alpha", {"echo"}) + manager._replace_server_tool_routes("id-zulu", {"echo"}) assert manager._get_mcp_server_from_tool_name("echo") is None @@ -3929,7 +4117,7 @@ class TestMCPServerManager: manager = MCPServerManager() alpha = MCPServer(server_id="id-alpha", name="echo_alpha", transport=MCPTransport.http) manager.registry = {"id-alpha": alpha} - manager._register_tool_route("echo", "id-alpha") + manager._replace_server_tool_routes("id-alpha", {"echo"}) assert manager._get_mcp_server_from_tool_name("echo") is alpha @@ -3943,8 +4131,8 @@ class TestMCPServerManager: first = MCPServer(server_id="id-first", name="shared", server_name="shared", transport=MCPTransport.http) second = MCPServer(server_id="id-second", name="shared", server_name="shared", transport=MCPTransport.http) manager.registry = {"id-first": first, "id-second": second} - manager._register_tool_route("shared-echo", "id-first") - manager._register_tool_route("shared-echo", "id-second") + manager._replace_server_tool_routes("id-first", {"shared-echo"}) + manager._replace_server_tool_routes("id-second", {"shared-echo"}) assert manager._get_mcp_server_from_tool_name("shared-echo") is None @@ -3954,7 +4142,7 @@ class TestMCPServerManager: alpha = MCPServer(server_id="id-alpha", name="alpha", server_name="alpha", transport=MCPTransport.http) beta = MCPServer(server_id="id-beta", name="beta", server_name="beta", transport=MCPTransport.http) manager.registry = {"id-alpha": alpha, "id-beta": beta} - manager._register_tool_route("echo", "id-alpha") + manager._replace_server_tool_routes("id-alpha", {"echo"}) assert manager._get_mcp_server_from_tool_name("beta-echo") is None assert manager._get_mcp_server_from_tool_name("alpha-echo") is alpha @@ -3970,7 +4158,7 @@ class TestMCPServerManager: alpha = MCPServer(server_id="id-alpha", name="echo_alpha", transport=MCPTransport.http) zulu = MCPServer(server_id="id-zulu", name="echo_zulu", transport=MCPTransport.http) manager.registry = {"id-alpha": alpha, "id-zulu": zulu} - manager._register_tool_route("secret_tool", "id-zulu") + manager._replace_server_tool_routes("id-zulu", {"secret_tool"}) route = manager.resolve_tool_route("secret_tool", allowed_server_ids=frozenset({"id-alpha"})) @@ -3982,8 +4170,8 @@ class TestMCPServerManager: alpha = MCPServer(server_id="id-alpha", name="echo_alpha", transport=MCPTransport.http) zulu = MCPServer(server_id="id-zulu", name="echo_zulu", transport=MCPTransport.http) manager.registry = {"id-alpha": alpha, "id-zulu": zulu} - manager._register_tool_route("echo", "id-alpha") - manager._register_tool_route("echo", "id-zulu") + manager._replace_server_tool_routes("id-alpha", {"echo"}) + manager._replace_server_tool_routes("id-zulu", {"echo"}) route = manager.resolve_tool_route("echo") @@ -3994,7 +4182,7 @@ class TestMCPServerManager: manager = MCPServerManager() alpha = MCPServer(server_id="id-alpha", name="echo_alpha", transport=MCPTransport.http) manager.registry = {"id-alpha": alpha} - manager._register_tool_route("echo", "id-alpha") + manager._replace_server_tool_routes("id-alpha", {"echo"}) route = manager.resolve_tool_route("echo") @@ -4012,8 +4200,8 @@ class TestMCPServerManager: alpha = MCPServer(server_id="id-alpha", name="echo_alpha", transport=MCPTransport.http) zulu = MCPServer(server_id="id-zulu", name="echo_zulu", transport=MCPTransport.http) manager.registry = {"id-alpha": alpha, "id-zulu": zulu} - manager._register_tool_route("echo", "id-alpha") - manager._register_tool_route("echo", "id-zulu") + manager._replace_server_tool_routes("id-alpha", {"echo"}) + manager._replace_server_tool_routes("id-zulu", {"echo"}) manager._cleanup_server_tool_routing_artifacts(zulu) del manager.registry["id-zulu"] @@ -4032,9 +4220,8 @@ class TestMCPServerManager: alpha = MCPServer(server_id="id-alpha", name="echo_alpha", transport=MCPTransport.http) zulu = MCPServer(server_id="id-zulu", name="echo_zulu", transport=MCPTransport.http) manager.registry = {"id-alpha": alpha, "id-zulu": zulu} - manager._register_tool_route("shared", "id-alpha") - manager._register_tool_route("shared", "id-zulu") - manager._register_tool_route("alpha_only", "id-alpha") + manager._replace_server_tool_routes("id-alpha", {"shared", "alpha_only"}) + manager._replace_server_tool_routes("id-zulu", {"shared"}) original_map = manager.tool_name_to_mcp_server_ids_mapping manager._cleanup_server_tool_routing_artifacts(zulu) @@ -4047,7 +4234,7 @@ class TestMCPServerManager: manager = MCPServerManager() alpha = MCPServer(server_id="id-alpha", name="echo_alpha", transport=MCPTransport.http) manager.registry = {"id-alpha": alpha} - manager._register_tool_route("echo", "id-alpha") + manager._replace_server_tool_routes("id-alpha", {"echo"}) manager._cleanup_server_tool_routing_artifacts(alpha) @@ -4422,6 +4609,53 @@ class TestMCPServerManager: assert manager.tool_name_to_mcp_server_ids_mapping["close_issue"] == frozenset({"jira"}) assert manager.tool_name_to_mcp_server_ids_mapping["jira-close_issue"] == frozenset({"jira"}) + def test_relisting_withdraws_routes_for_a_tool_the_upstream_removed(self): + """A tools/list that no longer reports a tool must withdraw that route. + + Registration is authoritative per server, so the second listing is the truth. + With union-only registration "close_issue" would keep resolving to jira after + the upstream dropped it, and a name a second server still served would keep + looking ambiguous (409) until the proxy restarted. + """ + manager = MCPServerManager() + server = MCPServer(server_id="jira", name="jira", transport=MCPTransport.http) + + def _tool(name: str) -> MCPTool: + return MCPTool(name=name, description="", inputSchema={}) + + manager._create_prefixed_tools([_tool("create_issue"), _tool("close_issue")], server) + assert manager.tool_name_to_mcp_server_ids_mapping["close_issue"] == frozenset({"jira"}) + + # Upstream drops close_issue. + manager._create_prefixed_tools([_tool("create_issue")], server) + + assert "close_issue" not in manager.tool_name_to_mcp_server_ids_mapping + assert "jira-close_issue" not in manager.tool_name_to_mcp_server_ids_mapping + assert manager.tool_name_to_mcp_server_ids_mapping["create_issue"] == frozenset({"jira"}) + assert manager.tool_name_to_mcp_server_ids_mapping["jira-create_issue"] == frozenset({"jira"}) + + def test_relisting_clears_a_false_ambiguity_once_one_owner_drops_the_tool(self): + """The reviewed edge case end to end: the 409 must clear without a restart.""" + manager = MCPServerManager() + alpha = MCPServer(server_id="id-alpha", name="echo_alpha", transport=MCPTransport.http) + zulu = MCPServer(server_id="id-zulu", name="echo_zulu", transport=MCPTransport.http) + manager.registry = {"id-alpha": alpha, "id-zulu": zulu} + scope = frozenset({"id-alpha", "id-zulu"}) + + def _tool(name: str) -> MCPTool: + return MCPTool(name=name, description="", inputSchema={}) + + manager._create_prefixed_tools([_tool("echo")], alpha) + manager._create_prefixed_tools([_tool("echo")], zulu) + assert manager.resolve_tool_route("echo", allowed_server_ids=scope).kind == "ambiguous" + + # alpha's upstream stops exposing echo, so zulu is now the sole owner. + manager._create_prefixed_tools([], alpha) + + route = manager.resolve_tool_route("echo", allowed_server_ids=scope) + assert route.kind == "resolved" + assert route.server is zulu + def test_get_mcp_server_from_tool_name_with_prefixed_and_unprefixed(self): """After mapping is populated, manager resolves both prefixed and unprefixed tool names to the same server.""" manager = MCPServerManager()