diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 36ecb05208b..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,6 +2875,20 @@ class MCPServerManager: ) return any(owner is not None and normalize_server_name(owner) in owned for owner in mapped_owners) + 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.""" + return { + normalize_server_name(known_prefix): server + for server in reversed(tuple(self.get_registry().values())) + for known_prefix in iter_known_server_prefixes(server) + } + + def server_owning_tool_name_prefix(self, tool_name: str) -> MCPServer | None: + prefix_to_server: Final = self._known_prefix_to_server() + matched: Final = match_known_server_prefix(tool_name, prefix_to_server.keys()) + return None if matched is None else prefix_to_server.get(matched[0]) + def remove_server(self, mcp_server: LiteLLM_MCPServerTable): """ Remove a server from the registry @@ -6114,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 @@ -6475,15 +6489,7 @@ class MCPServerManager: MCPServer if found, None otherwise """ registry_servers: Final = list(self.get_registry().values()) - - # Build prefix → server lookup covering every known form a tool name - # may take (alias / server_name / server_id / short ID). This is what - # makes the short-prefix mode work without breaking historical names. - prefix_to_server: Final[dict[str, MCPServer]] = {} - for server in registry_servers: - for known_prefix in iter_known_server_prefixes(server): - normalised = normalize_server_name(known_prefix) - prefix_to_server.setdefault(normalised, server) + prefix_to_server: Final = self._known_prefix_to_server() # First try with the original tool name if tool_name in self.tool_name_to_mcp_server_name_mapping: @@ -6501,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 4ea22ca1f01..3a9bca926b0 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -2888,6 +2888,40 @@ if MCP_AVAILABLE: headers={"WWW-Authenticate": get_byok_www_authenticate()}, ) + 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, + mcp_server_auth_headers: dict[str, dict[str, str]] | None, + oauth2_headers: dict[str, str] | None, + raw_headers: dict[str, str] | None, + ) -> None: + """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 + 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.server_exposes_tool(server, tool_name): + return + if all(allowed.server_id != server.server_id for allowed in allowed_mcp_servers): + return + try: + await _get_tools_from_mcp_servers( + user_api_key_auth=user_api_key_auth, + mcp_auth_header=mcp_auth_header, + mcp_servers=[server.server_id], + mcp_server_auth_headers=mcp_server_auth_headers, + oauth2_headers=oauth2_headers, + raw_headers=raw_headers, + ) + except Exception as e: # noqa: BLE001 # best effort: resolution below answers as it did before + verbose_logger.debug("MCP tools/call: listing %s before its first call failed: %s", server.name, e) + async def execute_mcp_tool( name: str, arguments: dict[str, object], @@ -2948,6 +2982,27 @@ if MCP_AVAILABLE: all_registry_prefixes.add(normalize_server_name(known_prefix)) name_is_prefixed = is_tool_name_prefixed(name, known_server_prefixes=all_registry_prefixes) + first_call_target: Final = ( + requested_server + 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, + mcp_server_auth_headers=mcp_server_auth_headers, + oauth2_headers=oauth2_headers, + raw_headers=raw_headers, + ) + if requested_server is not None and not name_is_prefixed: # REST callers may pass server_id with the upstream tool name (no # LiteLLM prefix). The first segment is not a registered server 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 47ec25a90f7..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 @@ -7437,6 +7437,186 @@ async def test_execute_mcp_tool_rest_server_id_authoritative_for_unprefixed_tool assert captured["name"] == "echo" +def _never_listed_passthrough_server() -> MCPServer: + return MCPServer( + server_id="lazy-map-1", + name="lazy_map", + server_name="lazy_map", + url="https://up.example.com/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.true_passthrough, + ) + + +@contextlib.contextmanager +def _worker_that_never_listed(server: MCPServer, upstream_tools: tuple[str, ...]): + """A worker whose tool rows for ``server`` are empty, in front of an upstream that answers + tools/list with ``upstream_tools`` and a managed dispatch that records what reaches it.""" + from mcp.types import Tool as MCPTool + + from litellm.proxy._experimental.mcp_server import server as mcp_module + + mcp_module.global_mcp_server_manager.registry[server.server_id] = server + dispatched: dict[str, object] = {} + + async def fake_handle_managed_mcp_tool(**kwargs): + dispatched.update(kwargs) + return CallToolResult(content=[TextContent(type="text", text="ok")], is_error=False) + + async def fake_fetch_tools(client, server_name): + return [MCPTool(name=tool_name, inputSchema={}) for tool_name in upstream_tools] + + with ( + patch.object( # test-quality-ok: the upstream MCP session is the boundary; a real one needs an initialize handshake over a live server + mcp_module.global_mcp_server_manager, + "_create_mcp_client", + new=AsyncMock(return_value=MagicMock()), + ) as create_client, + patch.object( # test-quality-ok: same boundary, this is the tools/list answer the upstream would give + mcp_module.global_mcp_server_manager, + "_fetch_tools_with_timeout", + side_effect=fake_fetch_tools, + ) as fetch_tools, + patch.object( # test-quality-ok: records the resolved server and bare name the managed call would forward upstream + mcp_module, "_handle_managed_mcp_tool", new=fake_handle_managed_mcp_tool + ), + ): + yield SimpleNamespace(create_client=create_client, fetch_tools=fetch_tools, dispatched=dispatched) + + +@pytest.mark.asyncio +async def test_execute_mcp_tool_lists_never_listed_passthrough_server_with_caller_token_first(): + """A prefixed tools/call on a worker that has not served tools/list must list that server once + with the caller's own credentials and then dispatch, instead of answering 404.""" + 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",)) as worker: + await mcp_module.execute_mcp_tool( + name="lazy_map-add", + arguments={"a": 1, "b": 2}, + allowed_mcp_servers=[server], + start_time=datetime.now(), + mcp_auth_header="Bearer caller-token", + raw_headers={"authorization": "Bearer caller-token"}, + ) + + assert worker.fetch_tools.await_count == 1 + assert "caller-token" in str(worker.create_client.await_args.kwargs.get("mcp_auth_header")) + assert worker.dispatched["server_name"] == "lazy_map" + assert worker.dispatched["name"] == "add" + + +@pytest.mark.asyncio +async def test_execute_mcp_tool_rest_server_id_lists_never_listed_server_first(): + 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",)) as worker: + await mcp_module.execute_mcp_tool( + name="add", + arguments={"a": 1, "b": 2}, + allowed_mcp_servers=[server], + start_time=datetime.now(), + mcp_auth_header="Bearer caller-token", + requested_server_id=server.server_id, + ) + + assert worker.fetch_tools.await_count == 1 + assert worker.dispatched["server_name"] == "lazy_map" + assert worker.dispatched["name"] == "add" + + +@pytest.mark.asyncio +async def test_execute_mcp_tool_unknown_tool_on_never_listed_server_lists_once_then_404s(): + 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",)) as worker, + pytest.raises(HTTPException) as exc_info, + ): + await mcp_module.execute_mcp_tool( + name="lazy_map-nope", + arguments={}, + allowed_mcp_servers=[server], + start_time=datetime.now(), + mcp_auth_header="Bearer caller-token", + ) + + assert exc_info.value.status_code == 404 + assert worker.fetch_tools.await_count == 1 + assert worker.dispatched == {} + + +@pytest.mark.asyncio +async def test_execute_mcp_tool_does_not_relist_a_server_this_worker_already_listed(): + 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",)) 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-add", + 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 == 0 + 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 + + server = _never_listed_passthrough_server() + other_server = MCPServer(server_id="other-1", name="other", transport=MCPTransport.http) + with ( + _worker_that_never_listed(server, upstream_tools=("add",)) as worker, + pytest.raises(HTTPException) as exc_info, + ): + await mcp_module.execute_mcp_tool( + name="lazy_map-add", + arguments={"a": 1, "b": 2}, + allowed_mcp_servers=[other_server], + start_time=datetime.now(), + mcp_auth_header="Bearer caller-token", + ) + + assert exc_info.value.status_code == 403 + assert worker.fetch_tools.await_count == 0 + assert worker.dispatched == {} + + @pytest.mark.asyncio async def test_execute_mcp_tool_strips_a_prefix_that_contains_the_separator(): """A server with no alias publishes its UUID server_id as the tool prefix. 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 dc1eed9ed7f..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 @@ -5556,6 +5556,59 @@ class TestMCPServerManager: ) mock_inject.assert_awaited_once() + def test_server_owning_tool_name_prefix_is_known_before_the_server_is_ever_listed(self): + manager = MCPServerManager() + server = MCPServer( + server_id="lazy-map-1", + name="lazy_map", + server_name="lazy_map", + transport=MCPTransport.http, + auth_type=MCPAuth.true_passthrough, + ) + manager.registry = {server.server_id: server} + + 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 + + manager._create_prefixed_tools([MCPTool(name="add", inputSchema={})], server) + + 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) + second = MCPServer( + server_id="second-id", name="second", server_name="second", alias="first", transport=MCPTransport.http + ) + manager.registry = {"first-id": first, "second-id": second} + + prefix_to_server = manager._known_prefix_to_server() + + assert prefix_to_server["first"] is first + assert prefix_to_server["second"] is second + assert manager.server_owning_tool_name_prefix("first-add") is first + def test_resolve_mcp_server_for_tool_call_via_prefixed_name(self): """Resolution succeeds when the prefixed tool name is in the mapping.""" manager = MCPServerManager()