From f3b17686160893b888fdf27f15290b2dbcbe2fe0 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Thu, 19 Mar 2026 17:40:50 -0700 Subject: [PATCH] feat: add _assert_can_manage_team_mcp_server auth helper with tests Co-Authored-By: Claude Opus 4.6 --- litellm/proxy/_types.py | 3 +- .../management_endpoints/common_utils.py | 12 ++ .../mcp_management_endpoints.py | 69 ++++++- .../test_mcp_manager_role.py | 187 ++++++++++++++++++ 4 files changed, 269 insertions(+), 2 deletions(-) create mode 100644 tests/test_litellm/proxy/management_endpoints/test_mcp_manager_role.py diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 9e86680e35..0c2cb8c827 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -1608,8 +1608,9 @@ class Member(MemberBase): role: Literal[ "admin", "user", + "mcp_server_manager", ] = Field( - description="The role of the user within the team. 'admin' users can manage team settings and members, 'user' is a regular team member" + description="The role of the user within the team. 'admin' users can manage team settings and members, 'user' is a regular team member, 'mcp_server_manager' can manage MCP servers for the team" ) diff --git a/litellm/proxy/management_endpoints/common_utils.py b/litellm/proxy/management_endpoints/common_utils.py index efc42d3355..092b22cee6 100644 --- a/litellm/proxy/management_endpoints/common_utils.py +++ b/litellm/proxy/management_endpoints/common_utils.py @@ -41,6 +41,18 @@ def _is_user_team_admin( return False +def _is_user_team_mcp_manager( + user_api_key_dict: UserAPIKeyAuth, team_obj: LiteLLM_TeamTable +) -> bool: + for member in team_obj.members_with_roles: + if ( + member.user_id is not None and member.user_id == user_api_key_dict.user_id + ) and member.role == "mcp_server_manager": + return True + + return False + + async def _is_user_org_admin_for_team( user_api_key_dict: UserAPIKeyAuth, team_obj: LiteLLM_TeamTable ) -> bool: diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index 3e5b729cea..b94bd87596 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -130,11 +130,78 @@ if MCP_AVAILABLE: ) from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_utils.http_parsing_utils import _read_request_body - from litellm.proxy.management_endpoints.common_utils import _user_has_admin_view + from litellm.proxy.auth.auth_checks import get_team_object + from litellm.proxy.management_endpoints.common_utils import ( + _is_user_team_mcp_manager, + _user_has_admin_view, + ) + from litellm.proxy.management_helpers.object_permission_utils import ( + _get_team_allowed_mcp_servers, + handle_update_object_permission_common, + ) from litellm.proxy.management_helpers.utils import management_endpoint_wrapper from litellm.types.mcp import MCPCredentials from litellm.types.mcp_server.mcp_server_manager import MCPServer + async def _assert_can_manage_team_mcp_server( + user_api_key_dict: UserAPIKeyAuth, + team_id: Optional[str] = None, + server_id: Optional[str] = None, + ) -> str: + """ + Verify that the caller is an MCP server manager for a team and (for edit/delete) + that the target server belongs to that team. + + Returns the team_id the caller is managing. + Raises HTTPException(400) if no team_id can be determined. + Raises HTTPException(403) if the caller is not an MCP manager or server not in team. + """ + from litellm.proxy.proxy_server import prisma_client, user_api_key_cache + + resolved_team_id = team_id or user_api_key_dict.team_id + if not resolved_team_id: + raise HTTPException( + status_code=400, + detail={"error": "team_id is required for MCP server manager operations."}, + ) + + if ( + team_id + and user_api_key_dict.team_id + and team_id != user_api_key_dict.team_id + ): + raise HTTPException( + status_code=403, + detail={"error": "team_id does not match the API key's team."}, + ) + + team_obj = await get_team_object( + team_id=resolved_team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + check_db_only=False, + ) + + if not _is_user_team_mcp_manager(user_api_key_dict, team_obj): + raise HTTPException( + status_code=403, + detail={ + "error": f"User does not have mcp_server_manager role in team {resolved_team_id}." + }, + ) + + if server_id is not None: + team_server_ids = await _get_team_allowed_mcp_servers(team_obj) + if server_id not in team_server_ids: + raise HTTPException( + status_code=403, + detail={ + "error": f"MCP server {server_id} is not assigned to team {resolved_team_id}." + }, + ) + + return resolved_team_id + @dataclass class _TemporaryMCPServerEntry: server: MCPServer diff --git a/tests/test_litellm/proxy/management_endpoints/test_mcp_manager_role.py b/tests/test_litellm/proxy/management_endpoints/test_mcp_manager_role.py new file mode 100644 index 0000000000..20cb337cd1 --- /dev/null +++ b/tests/test_litellm/proxy/management_endpoints/test_mcp_manager_role.py @@ -0,0 +1,187 @@ +import pytest + +from litellm.proxy._types import ( + LiteLLM_TeamTable, + LitellmUserRoles, + Member, + UserAPIKeyAuth, +) +from litellm.proxy.management_endpoints.common_utils import ( + _is_user_team_mcp_manager, +) + + +class TestIsUserTeamMcpManager: + def test_mcp_server_manager_role_returns_true(self): + user_auth = UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + user_id="user1", + api_key="sk-test", + ) + team = LiteLLM_TeamTable( + team_id="team1", + members_with_roles=[ + Member(user_id="user1", role="mcp_server_manager") + ], + ) + assert _is_user_team_mcp_manager(user_auth, team) is True + + def test_regular_user_role_returns_false(self): + user_auth = UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + user_id="user1", + api_key="sk-test", + ) + team = LiteLLM_TeamTable( + team_id="team1", + members_with_roles=[Member(user_id="user1", role="user")], + ) + assert _is_user_team_mcp_manager(user_auth, team) is False + + def test_admin_role_returns_false(self): + user_auth = UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + user_id="user1", + api_key="sk-test", + ) + team = LiteLLM_TeamTable( + team_id="team1", + members_with_roles=[Member(user_id="user1", role="admin")], + ) + assert _is_user_team_mcp_manager(user_auth, team) is False + + def test_user_not_in_team_returns_false(self): + user_auth = UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + user_id="user2", + api_key="sk-test", + ) + team = LiteLLM_TeamTable( + team_id="team1", + members_with_roles=[ + Member(user_id="user1", role="mcp_server_manager") + ], + ) + assert _is_user_team_mcp_manager(user_auth, team) is False + + +from unittest.mock import AsyncMock, MagicMock, patch +from litellm.proxy._types import LiteLLM_TeamTableCachedObj + + +@pytest.mark.asyncio +class TestAssertCanManageTeamMcpServer: + async def test_mcp_manager_with_team_id_succeeds(self): + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + _assert_can_manage_team_mcp_server, + ) + user_auth = UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, user_id="user1", api_key="sk-test", + ) + mock_team = LiteLLM_TeamTableCachedObj( + team_id="team1", + members_with_roles=[Member(user_id="user1", role="mcp_server_manager")], + ) + with patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_team_object", + AsyncMock(return_value=mock_team), + ): + result = await _assert_can_manage_team_mcp_server( + user_api_key_dict=user_auth, team_id="team1" + ) + assert result == "team1" + + async def test_regular_user_gets_403(self): + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + _assert_can_manage_team_mcp_server, + ) + user_auth = UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, user_id="user1", api_key="sk-test", + ) + mock_team = LiteLLM_TeamTableCachedObj( + team_id="team1", + members_with_roles=[Member(user_id="user1", role="user")], + ) + with patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_team_object", + AsyncMock(return_value=mock_team), + ): + with pytest.raises(Exception) as exc_info: + await _assert_can_manage_team_mcp_server( + user_api_key_dict=user_auth, team_id="team1" + ) + assert exc_info.value.status_code == 403 + + async def test_admin_gets_403(self): + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + _assert_can_manage_team_mcp_server, + ) + user_auth = UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, user_id="user1", api_key="sk-test", + ) + mock_team = LiteLLM_TeamTableCachedObj( + team_id="team1", + members_with_roles=[Member(user_id="user1", role="admin")], + ) + with patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_team_object", + AsyncMock(return_value=mock_team), + ): + with pytest.raises(Exception) as exc_info: + await _assert_can_manage_team_mcp_server( + user_api_key_dict=user_auth, team_id="team1" + ) + assert exc_info.value.status_code == 403 + + async def test_mcp_manager_server_in_team_succeeds(self): + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + _assert_can_manage_team_mcp_server, + ) + user_auth = UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, user_id="user1", api_key="sk-test", team_id="team1", + ) + mock_team = LiteLLM_TeamTableCachedObj( + team_id="team1", + members_with_roles=[Member(user_id="user1", role="mcp_server_manager")], + ) + with ( + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_team_object", + AsyncMock(return_value=mock_team), + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints._get_team_allowed_mcp_servers", + AsyncMock(return_value={"server1", "server2"}), + ), + ): + result = await _assert_can_manage_team_mcp_server( + user_api_key_dict=user_auth, server_id="server1" + ) + assert result == "team1" + + async def test_mcp_manager_server_not_in_team_gets_403(self): + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + _assert_can_manage_team_mcp_server, + ) + user_auth = UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, user_id="user1", api_key="sk-test", team_id="team1", + ) + mock_team = LiteLLM_TeamTableCachedObj( + team_id="team1", + members_with_roles=[Member(user_id="user1", role="mcp_server_manager")], + ) + with ( + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_team_object", + AsyncMock(return_value=mock_team), + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints._get_team_allowed_mcp_servers", + AsyncMock(return_value={"server2", "server3"}), + ), + ): + with pytest.raises(Exception) as exc_info: + await _assert_can_manage_team_mcp_server( + user_api_key_dict=user_auth, server_id="server1" + ) + assert exc_info.value.status_code == 403