From 082064569082761bbb8ce1dc1b1be2fb2a145b94 Mon Sep 17 00:00:00 2001 From: tin-berri Date: Sat, 24 Oct 2026 10:17:14 -0700 Subject: [PATCH] fix(mcp): retain toolset scope when refreshing catalog authority (#48095) --- .../mcp_server/auth/user_api_key_auth_mcp.py | 9 +- .../proxy/_experimental/mcp_server/server.py | 2 +- .../auth/test_user_api_key_auth_mcp.py | 19 +++- .../_experimental/mcp_server/test_catalog.py | 101 +++++++++++++++++- 4 files changed, 125 insertions(+), 6 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py index 53b56c7b31d..7f0d392ab1b 100644 --- a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py +++ b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py @@ -1368,7 +1368,7 @@ class MCPRequestHandler: if auth is None: return None - refreshed = auth.model_copy() + refreshed: Final = auth.model_copy(update={"requires_fresh_policy": True}) if (auth.mcp_admitted_user_subject or auth.is_session_token) and auth.user_id: # A session credential (``lite login`` or a redeemed ID-JAG) has no key row, its authority is # sealed in the token: refresh the live user policy and keep the token's team, routes and models. @@ -1386,7 +1386,12 @@ class MCPRequestHandler: refreshed.project_id = current.project_id refreshed.user_id = current.user_id refreshed.access_group_ids = current.access_group_ids - refreshed.requires_fresh_policy = True + else: + return refreshed + if refreshed.mcp_toolset_id is not None: + from litellm.proxy._experimental.mcp_server.server import apply_toolset_scope + + return await apply_toolset_scope(refreshed, refreshed.mcp_toolset_id) return refreshed @staticmethod diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index b3426089375..bb7c81f9d60 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -1797,7 +1797,7 @@ if MCP_AVAILABLE: return acting.model_copy(update={"mcp_toolset_id": toolset_id}) tool_permissions = await operations.global_mcp_server_manager.resolve_toolset_tool_permissions( - toolset_ids=[toolset_id] + toolset_ids=[toolset_id], requires_fresh_policy=acting.requires_fresh_policy ) server_ids: Final = list(tool_permissions.keys()) existing_op: Final = acting.object_permission diff --git a/tests/unit/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py b/tests/unit/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py index f73a8547158..e8689b2ba89 100644 --- a/tests/unit/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py +++ b/tests/unit/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py @@ -10906,7 +10906,19 @@ async def test_catalog_refresh_reads_current_user_org_without_losing_resource_sc @pytest.mark.asyncio @pytest.mark.parametrize("current_groups", [[], ["replacement-group"]]) async def test_catalog_refresh_uses_current_virtual_key_policy_and_keeps_session_scope(monkeypatch, current_groups): - permission = LiteLLM_ObjectPermissionTable(object_permission_id="current-policy", mcp_servers=["current-server"]) + from litellm.proxy import proxy_server + from litellm.types.mcp_server.mcp_toolset import MCPToolset + + permission = LiteLLM_ObjectPermissionTable( + object_permission_id="current-policy", mcp_toolsets=["session-toolset"] + ) + toolset: Final = MCPToolset( + toolset_id="session-toolset", toolset_name="session-toolset", + tools=[{"server_id": "current-server", "tool_name": "read"}], + ) + tables: Final = SimpleNamespace(litellm_mcptoolsettable=SimpleNamespace(find_many=AsyncMock(return_value=[toolset]))) + monkeypatch.setattr(proxy_server, "prisma_client", SimpleNamespace(db=tables, writer_db=tables)) + monkeypatch.setattr(proxy_server, "user_api_key_cache", DualCache()) current = UserAPIKeyAuth( object_permission=permission, object_permission_id="current-policy", @@ -10931,7 +10943,10 @@ async def test_catalog_refresh_uses_current_virtual_key_policy_and_keeps_session caller.mcp_toolset_id = "session-toolset" refreshed = await MCPRequestHandler.refresh_catalog_authority(caller) reload_key.assert_awaited_once_with("owned-key-hash", check_db_only=True) - assert refreshed.object_permission == permission + assert refreshed.object_permission == permission.model_copy(update={ + "mcp_servers": ["current-server"], "mcp_toolsets": [], + "mcp_tool_permissions": {"current-server": ["read"]}, + }) assert refreshed.object_permission_id == "current-policy" assert (refreshed.team_id, refreshed.org_id, refreshed.project_id, refreshed.user_id) == ( "new-team", diff --git a/tests/unit/proxy/_experimental/mcp_server/test_catalog.py b/tests/unit/proxy/_experimental/mcp_server/test_catalog.py index b383078c0ed..4f6530e5de1 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_catalog.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_catalog.py @@ -1,5 +1,5 @@ import asyncio -from collections.abc import Sequence +from collections.abc import Mapping, Sequence from types import SimpleNamespace from typing import Final, Literal from unittest.mock import AsyncMock @@ -13,6 +13,7 @@ from mcp.types import ( ListResourcesResult, ListResourceTemplatesRequest, ListResourceTemplatesResult, + ListToolsRequest, ListToolsResult, PaginatedRequestParams, Tool, @@ -34,6 +35,104 @@ CatalogKind = Literal["tools", "prompts", "resources", "templates"] OptionalCatalogResult = ListPromptsResult | ListResourcesResult | ListResourceTemplatesResult +@pytest.mark.asyncio +@pytest.mark.parametrize( + "selected,revoked,expected", + [ + ("narrow", False, ("alpha-read_item",)), + ("full", False, ("alpha-read_item", "alpha-write_item", "beta-read_item", "beta-write_item")), + (None, False, ("alpha-read_item", "alpha-write_item", "beta-read_item", "beta-write_item")), + ("narrow", True, ()), + ], +) +async def test_catalog_refresh_keeps_selected_toolset_with_current_grants_and_contents( + monkeypatch: pytest.MonkeyPatch, + selected: str | None, + revoked: bool, + expected: tuple[str, ...], +) -> None: + from fastapi import HTTPException + + from litellm.caching.dual_cache import DualCache + from litellm.proxy import proxy_server + from litellm.proxy._experimental.mcp_server import operations + from litellm.proxy._experimental.mcp_server.server import apply_toolset_scope + from litellm.proxy._experimental.mcp_server.tool_registry import global_mcp_tool_registry + from litellm.proxy._types import LiteLLM_ObjectPermissionTable + from litellm.types.mcp_server.mcp_toolset import MCPToolset, MCPToolsetTool + + grants: Final = LiteLLM_ObjectPermissionTable( + object_permission_id="catalog-grants", mcp_toolsets=["narrow", "full"] + ) + live_grants: Final = grants.model_copy(update={"mcp_toolsets": ["full"]}) if revoked else grants + tools: Final[tuple[MCPToolsetTool, ...]] = ( + {"server_id": "alpha", "tool_name": "read_item"}, + {"server_id": "alpha", "tool_name": "write_item"}, + {"server_id": "beta", "tool_name": "read_item"}, + {"server_id": "beta", "tool_name": "write_item"}, + ) + toolsets: Final = ( + MCPToolset(toolset_id="narrow", toolset_name="narrow", tools=[tools[0]]), + MCPToolset(toolset_id="full", toolset_name="full", tools=list(tools)), + ) + + async def read_toolsets(where: Mapping[str, Mapping[str, Sequence[str]]]) -> tuple[MCPToolset, ...]: + return tuple(toolset for toolset in toolsets if toolset.toolset_id in where["toolset_id"]["in"]) + + tables: Final = SimpleNamespace( + litellm_objectpermissiontable=SimpleNamespace(find_unique=AsyncMock(return_value=live_grants)), + litellm_mcptoolsettable=SimpleNamespace(find_many=read_toolsets), + ) + database: Final = SimpleNamespace( + db=tables, + writer_db=tables, + get_data=AsyncMock(return_value=UserAPIKeyAuth(object_permission_id=grants.object_permission_id)), + ) + cache: Final = DualCache() + monkeypatch.setattr(proxy_server, "prisma_client", database) + monkeypatch.setattr(proxy_server, "user_api_key_cache", cache) + monkeypatch.setattr(proxy_server, "general_settings", {"supported_db_objects": []}) + await cache.async_set_cache(key="toolset_perms:narrow", value={"alpha": ["read_item", "write_item"]}) + manager: Final = operations.global_mcp_server_manager + for server_id in ("alpha", "beta"): + monkeypatch.setitem( + manager.registry, + server_id, + MCPServer(server_id=server_id, name=server_id, transport=MCPTransport.http, spec_path="/catalog.yaml"), + ) + + async def tool_handler() -> str: + return "ok" + + for tool in tools: + global_mcp_tool_registry.register_tool( + name=f"{tool['server_id']}-{tool['tool_name']}", + description=tool["tool_name"], + input_schema={"type": "object"}, + handler=tool_handler, + ) + original: Final = UserAPIKeyAuth( + api_key="catalog-key-hash", object_permission=grants, + object_permission_id=grants.object_permission_id, + ) + original.via_virtual_key = True + caller: Final = await apply_toolset_scope(original, selected) if selected is not None else original + before: Final = caller.model_dump() + context: Final = operations.prepare_context(user_api_key_auth=caller) + if revoked: + with pytest.raises(HTTPException, match="does not have access to toolset") as denied: + await catalog.list_gateway_catalog(context, ListToolsRequest(), log_list_tools_to_spendlogs=False) + assert denied.value.status_code == 403 + else: + result: Final = await catalog.list_gateway_catalog( + context, ListToolsRequest(), log_list_tools_to_spendlogs=False + ) + assert isinstance(result, ListToolsResult) + assert tuple(sorted(tool.name for tool in result.tools)) == expected + assert caller.model_dump() == before + assert original.object_permission == grants + + def page(name: str, cursor: str | None = None, revision: str = "stable") -> ListToolsResult: return ListToolsResult( tools=[Tool(name=name, input_schema={"type": "object"})],