From 65ae86b40763f90ee301642159c3c1f0705f3725 Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Mon, 23 Mar 2026 10:35:38 -0700 Subject: [PATCH] fix(mcp): toolset access control, asyncio fix, and real unit tests - server.py: _apply_toolset_scope now enforces that non-admin keys must have the requested toolset_id in their mcp_toolsets grant list; admin keys always bypass the check. - mcp_management_endpoints.py: three access-control fixes: * fetch_mcp_toolsets: non-admin keys with mcp_toolsets=None now return [] instead of all toolsets (only admins get 'all' when the field is absent) * fetch_mcp_toolset: non-admin keys that haven't been granted the requested toolset_id now get 403 instead of the full result * add_mcp_toolset: duplicate toolset_name now returns 409 Conflict instead of an opaque 500 - proxy_server.py: use asyncio.get_running_loop() instead of get_event_loop() inside an already-running coroutine (Python 3.10+). - test_mcp_toolset_scope.py: replace four hollow tests that only asserted local variable properties with real tests that call the production fetch_mcp_toolsets() and handle_streamable_http_mcp() functions with mocked dependencies. --- .../proxy/_experimental/mcp_server/server.py | 19 +- .../mcp_management_endpoints.py | 39 +++- litellm/proxy/proxy_server.py | 2 +- .../mcp_server/test_mcp_toolset_scope.py | 207 +++++++++++++++--- 4 files changed, 223 insertions(+), 44 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index f4f1f66ad5c..8ef83a03c97 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -2417,11 +2417,26 @@ if MCP_AVAILABLE: Restrict a key's MCP permissions to a single toolset. When a request arrives via /toolset/{name}/mcp we override the key's - object_permission so that only the toolset's tools are visible, - regardless of what the key's normal permissions are. + object_permission so that only the toolset's tools are visible. + + Raises HTTPException(403) if the key has an explicit toolset grant list + that does not include toolset_id (i.e. mcp_toolsets is set but empty, + or set to a list that omits this toolset). Admin keys always pass. """ from litellm.proxy._types import LiteLLM_ObjectPermissionTable + # Access control: non-admin keys must have this toolset in their grant list. + is_admin = getattr(user_api_key_auth, "user_role", None) == "proxy_admin" + if not is_admin: + op = user_api_key_auth.object_permission + granted = getattr(op, "mcp_toolsets", None) if op else None + # granted=None → no restriction (allow); granted=[] or list without toolset_id → deny + if granted is not None and toolset_id not in granted: + raise HTTPException( + status_code=403, + detail=f"API key does not have access to toolset '{toolset_id}'.", + ) + tool_permissions = ( await global_mcp_server_manager.resolve_toolset_tool_permissions( toolset_ids=[toolset_id] diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index f2d3eb709d0..4c9102f04c1 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -2072,7 +2072,17 @@ if MCP_AVAILABLE: touched_by = ( litellm_changed_by or user_api_key_dict.user_id or LITELLM_PROXY_ADMIN_NAME ) - result = await create_mcp_toolset(prisma_client, payload, touched_by) + try: + result = await create_mcp_toolset(prisma_client, payload, touched_by) + except Exception as e: + if "UniqueViolationError" in type(e).__name__ or "unique" in str(e).lower(): + raise HTTPException( + status_code=status.HTTP_409_CONFLICT, + detail={ + "error": f"A toolset named '{payload.toolset_name}' already exists." + }, + ) + raise from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( global_mcp_server_manager, ) @@ -2092,20 +2102,16 @@ if MCP_AVAILABLE: prisma_client = get_prisma_client_or_throw( "Database not connected. Connect a database to your proxy" ) - # Admins with no explicit restriction see all toolsets - if _user_has_admin_view(user_api_key_dict): - op = user_api_key_dict.object_permission - if op is None or getattr(op, "mcp_toolsets", None) is None: - return await list_mcp_toolsets(prisma_client) - + is_admin = _user_has_admin_view(user_api_key_dict) op = user_api_key_dict.object_permission - # Distinguish None (field absent = no restriction) from [] (explicitly empty = zero allowed). + # mcp_toolsets=None means the field was never set. + # For admins: None → no restriction → return all. + # For non-admins: None → no toolsets explicitly granted → return nothing. raw_toolsets = getattr(op, "mcp_toolsets", None) if op else None - # raw_toolsets is None → field not set → no restriction, return all - # raw_toolsets is [] → explicitly empty → return nothing - # raw_toolsets is [ids] → return only those if raw_toolsets is None: - return await list_mcp_toolsets(prisma_client) + if is_admin: + return await list_mcp_toolsets(prisma_client) + return [] if not raw_toolsets: return [] return await list_mcp_toolsets(prisma_client, toolset_ids=raw_toolsets) @@ -2122,6 +2128,15 @@ if MCP_AVAILABLE: prisma_client = get_prisma_client_or_throw( "Database not connected. Connect a database to your proxy" ) + # Non-admin keys may only fetch toolsets they've been explicitly granted. + if not _user_has_admin_view(user_api_key_dict): + op = user_api_key_dict.object_permission + granted = getattr(op, "mcp_toolsets", None) if op else None + if granted is None or toolset_id not in granted: + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail={"error": "API key does not have access to this toolset."}, + ) toolset = await get_mcp_toolset(prisma_client, toolset_id) if toolset is None: raise HTTPException( diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 4cdbfefc9be..28c8d8d4cc5 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -13514,7 +13514,7 @@ async def _stream_mcp_asgi_response( """ from starlette.responses import StreamingResponse - headers_ready: asyncio.Future = asyncio.get_event_loop().create_future() + headers_ready: asyncio.Future = asyncio.get_running_loop().create_future() body_queue: asyncio.Queue = asyncio.Queue() async def bridging_send(message): diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_toolset_scope.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_toolset_scope.py index 3f28fb15eb4..4dfd33cbbd4 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_toolset_scope.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_toolset_scope.py @@ -6,7 +6,11 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest -from litellm.proxy._types import LiteLLM_ObjectPermissionTable, UserAPIKeyAuth +from litellm.proxy._types import ( + LiteLLM_ObjectPermissionTable, + LitellmUserRoles, + UserAPIKeyAuth, +) def _make_auth( @@ -42,7 +46,11 @@ class TestApplyToolsetScope: "global_mcp_server_manager.resolve_toolset_tool_permissions", new=AsyncMock(return_value=toolset_perms), ): - auth = _make_auth(mcp_servers=["server-a", "server-b", "server-c"]) + # Key has been explicitly granted toolset-123 — access check passes. + auth = _make_auth( + mcp_servers=["server-a", "server-b", "server-c"], + mcp_toolsets=["toolset-123"], + ) result = await _apply_toolset_scope(auth, "toolset-123") op = result.object_permission @@ -72,25 +80,111 @@ class TestApplyToolsetScope: class TestFetchMCPToolsetsAccess: """Tests for GET /v1/mcp/toolset access control.""" - def test_empty_toolsets_returns_empty(self): + @pytest.mark.asyncio + async def test_non_admin_empty_grants_returns_empty(self): """Non-admin key with mcp_toolsets=[] must not see any toolsets.""" - # Simulate what fetch_mcp_toolsets does with raw_toolsets=[] - raw_toolsets: Optional[List[str]] = [] - # raw_toolsets is [] → return nothing - assert raw_toolsets is not None - assert not raw_toolsets # empty list → return [] + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + fetch_mcp_toolsets, + ) - def test_none_toolsets_returns_all(self): - """Key where mcp_toolsets is absent (None) should return all toolsets.""" - raw_toolsets: Optional[List[str]] = None - # raw_toolsets is None → no restriction → return all - assert raw_toolsets is None + auth = _make_auth(mcp_toolsets=[]) + mock_client = MagicMock() - def test_populated_toolsets_filters(self): - """Key with explicit toolset IDs should only see those.""" - raw_toolsets: Optional[List[str]] = ["ts-1", "ts-2"] - assert raw_toolsets is not None - assert len(raw_toolsets) == 2 + with ( + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", + return_value=mock_client, + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.list_mcp_toolsets", + new=AsyncMock(return_value=[]), + ) as mock_list, + ): + result = await fetch_mcp_toolsets(user_api_key_dict=auth) + + assert result == [] + mock_list.assert_not_called() + + @pytest.mark.asyncio + async def test_admin_unrestricted_returns_all(self): + """Admin key with mcp_toolsets absent (None) gets all toolsets.""" + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + fetch_mcp_toolsets, + ) + + auth = UserAPIKeyAuth( + api_key="sk-test", + user_role=LitellmUserRoles.PROXY_ADMIN, + object_permission=None, + ) + fake_toolsets = [MagicMock(), MagicMock()] + mock_client = MagicMock() + + with ( + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", + return_value=mock_client, + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.list_mcp_toolsets", + new=AsyncMock(return_value=fake_toolsets), + ) as mock_list, + ): + result = await fetch_mcp_toolsets(user_api_key_dict=auth) + + assert result == fake_toolsets + mock_list.assert_called_once_with(mock_client) + + @pytest.mark.asyncio + async def test_non_admin_none_grants_returns_empty(self): + """Non-admin key with no object_permission (field absent) gets no toolsets.""" + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + fetch_mcp_toolsets, + ) + + auth = UserAPIKeyAuth(api_key="sk-test", object_permission=None) + mock_client = MagicMock() + + with ( + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", + return_value=mock_client, + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.list_mcp_toolsets", + new=AsyncMock(return_value=[]), + ) as mock_list, + ): + result = await fetch_mcp_toolsets(user_api_key_dict=auth) + + assert result == [] + mock_list.assert_not_called() + + @pytest.mark.asyncio + async def test_populated_grants_filters_toolsets(self): + """Key with explicit toolset IDs fetches only those IDs from the DB.""" + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + fetch_mcp_toolsets, + ) + + auth = _make_auth(mcp_toolsets=["ts-1", "ts-2"]) + fake_toolsets = [MagicMock(toolset_id="ts-1"), MagicMock(toolset_id="ts-2")] + mock_client = MagicMock() + + with ( + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", + return_value=mock_client, + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.list_mcp_toolsets", + new=AsyncMock(return_value=fake_toolsets), + ) as mock_list, + ): + result = await fetch_mcp_toolsets(user_api_key_dict=auth) + + assert len(result) == 2 + mock_list.assert_called_once_with(mock_client, toolset_ids=["ts-1", "ts-2"]) class TestMCPActiveToolsetContextVar: @@ -109,14 +203,69 @@ class TestMCPActiveToolsetContextVar: _mcp_active_toolset_id.reset(token) assert _mcp_active_toolset_id.get() is None - def test_client_header_is_stripped(self): - """x-mcp-toolset-id header is removed from scope headers before auth runs.""" - # Simulate the stripping logic from handle_streamable_http_mcp - headers = [ - (b"authorization", b"Bearer sk-test"), - (b"x-mcp-toolset-id", b"evil-toolset"), - (b"content-type", b"application/json"), - ] - stripped = [(k, v) for k, v in headers if k.lower() != b"x-mcp-toolset-id"] - assert (b"x-mcp-toolset-id", b"evil-toolset") not in stripped - assert len(stripped) == 2 + @pytest.mark.asyncio + async def test_client_header_is_stripped_in_scope(self): + """handle_streamable_http_mcp strips x-mcp-toolset-id from scope before passing to session manager.""" + from litellm.proxy._experimental.mcp_server.server import ( + handle_streamable_http_mcp, + ) + + scope = { + "type": "http", + "path": "/mcp", + "method": "GET", + "query_string": b"", + "headers": [ + (b"authorization", b"Bearer sk-test"), + (b"x-mcp-toolset-id", b"evil-toolset"), + (b"content-type", b"application/json"), + ], + } + mock_auth = UserAPIKeyAuth(api_key="sk-test") + + async def fake_receive(): + return {"type": "http.disconnect"} + + async def fake_send(msg): + pass + + with ( + patch( + "litellm.proxy._experimental.mcp_server.server.extract_mcp_auth_context", + new=AsyncMock( + return_value=(mock_auth, None, [], {}, {}, scope["headers"]) + ), + ), + patch( + "litellm.proxy._experimental.mcp_server.server.IPAddressUtils", + MagicMock(get_mcp_client_ip=MagicMock(return_value="127.0.0.1")), + ), + patch( + "litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager", + MagicMock(get_mcp_server_by_name=MagicMock(return_value=None)), + ), + patch( + "litellm.proxy._experimental.mcp_server.server.MCPDebug", + MagicMock( + maybe_build_debug_headers=MagicMock(return_value=None), + ), + ), + patch( + "litellm.proxy._experimental.mcp_server.server.set_auth_context", + MagicMock(), + ), + patch( + "litellm.proxy._experimental.mcp_server.server._SESSION_MANAGERS_INITIALIZED", + True, + ), + patch( + "litellm.proxy._experimental.mcp_server.server._handle_stale_mcp_session", + new=AsyncMock(return_value=True), + ), + ): + await handle_streamable_http_mcp(scope, fake_receive, fake_send) + + header_keys = [k for k, _ in scope["headers"]] + assert b"x-mcp-toolset-id" not in header_keys + assert b"authorization" in header_keys + assert b"content-type" in header_keys