From b0ea4aeca556b7f38e71d2e0a20d83c6bed7e9ad Mon Sep 17 00:00:00 2001 From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Date: Sun, 4 Oct 2026 19:45:50 -0700 Subject: [PATCH] fix(mcp): authorize local handlers by registered server ownership --- .../mcp_server/mcp_server_manager.py | 8 ++-- .../_experimental/mcp_server/tool_registry.py | 3 ++ litellm/types/mcp_server/tool_registry.py | 3 +- .../mcp_server/test_mcp_server_manager.py | 40 +++++++++++++++++++ 4 files changed, 50 insertions(+), 4 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 6886e36cb00..9623ad2a2b6 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -2899,6 +2899,7 @@ class MCPServerManager: description=description, input_schema=input_schema, handler=tool_func, + server_id=server.server_id, ) # Update tool name to server name mapping (for both prefixed and base names) @@ -6675,9 +6676,10 @@ class MCPServerManager: """ from litellm.proxy._experimental.mcp_server.tool_registry import global_mcp_tool_registry - # Local handlers belong to their registered namespace, even if a native - # discovery result or an older cached route claims the same spelling. - if global_mcp_tool_registry.get_tool(tool_name) is not None: + registered_tool: Final = global_mcp_tool_registry.get_tool(tool_name) + if registered_tool is not None: + if registered_tool.server_id is not None: + return self.get_mcp_server_by_id(registered_tool.server_id) return self.server_owning_tool_name_prefix(tool_name) registry_servers: Final = list(self.get_registry().values()) diff --git a/litellm/proxy/_experimental/mcp_server/tool_registry.py b/litellm/proxy/_experimental/mcp_server/tool_registry.py index fee9537097d..dd25f09cc57 100644 --- a/litellm/proxy/_experimental/mcp_server/tool_registry.py +++ b/litellm/proxy/_experimental/mcp_server/tool_registry.py @@ -58,6 +58,8 @@ class MCPToolRegistry: description: str, input_schema: dict[str, Any], handler: Callable, + *, + server_id: str | None = None, ) -> None: """ Register a new tool in the registry @@ -67,6 +69,7 @@ class MCPToolRegistry: description=description, input_schema=input_schema, handler=handler, + server_id=server_id, ) verbose_logger.debug("Registered tool: %s", name) diff --git a/litellm/types/mcp_server/tool_registry.py b/litellm/types/mcp_server/tool_registry.py index 79dedfd7653..d937905b414 100644 --- a/litellm/types/mcp_server/tool_registry.py +++ b/litellm/types/mcp_server/tool_registry.py @@ -1,7 +1,7 @@ from collections.abc import Callable from typing import Any, ClassVar -from pydantic import BaseModel, ConfigDict +from pydantic import BaseModel, ConfigDict, Field class MCPTool(BaseModel): @@ -10,6 +10,7 @@ class MCPTool(BaseModel): description: str input_schema: dict[str, Any] handler: Callable + server_id: str | None = Field(default=None, frozen=True) class ToolSchema(BaseModel): diff --git a/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py index c2b6bff962e..0a816ab5f59 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -18499,3 +18499,43 @@ async def test_cached_catalog_excludes_routes_for_servers_outside_its_snapshot(m assert "search" not in manager.tool_name_to_mcp_server_name_mapping assert manager._get_mcp_server_from_tool_name("known").server_id == "pinned" read_rows.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_overlapping_server_prefix_cannot_authorize_registered_openapi_handler(monkeypatch, tmp_path): + from datetime import datetime + + from fastapi import HTTPException + + from litellm.proxy._experimental.mcp_server import operations + from litellm.proxy._experimental.mcp_server.tool_registry import global_mcp_tool_registry + + manager = MCPServerManager() + private = MCPServer(server_id="private", name="billing", alias="billing", transport=MCPTransport.http) + allowed = MCPServer(server_id="allowed", name="billing_admin", alias="billing-admin", transport=MCPTransport.http) + manager.config_mcp_servers = {server.server_id: server for server in (private, allowed)} + monkeypatch.setattr(global_mcp_tool_registry, "published_tools", {}) + monkeypatch.setattr(operations, "global_mcp_server_manager", manager) + spec = tmp_path / "spec.json" + spec.write_text(json.dumps({"openapi": "3.0.0", "paths": {"/export": {"get": {"operationId": "admin-export"}}}})) + await manager._register_openapi_tools(str(spec), private, "https://example.com") + tool = global_mcp_tool_registry.get_tool("billing-admin-export") + assert tool is not None + handler = AsyncMock(return_value="private result") + tool.handler = handler + check = AsyncMock(return_value={}) + monkeypatch.setattr(manager, "pre_call_tool_check", check) + + with pytest.raises(HTTPException) as denied: + await operations._execute_mcp_tool( + name="billing-admin-export", arguments={}, allowed_mcp_servers=[allowed], start_time=datetime.now() + ) + assert denied.value.status_code == 403 + handler.assert_not_awaited() + check.assert_not_awaited() + result = await operations._execute_mcp_tool( + name="billing-admin-export", arguments={}, allowed_mcp_servers=[private], start_time=datetime.now() + ) + assert result.is_error is False + handler.assert_awaited_once_with() + assert check.await_args.kwargs["server"].server_id == private.server_id