From f3b17686160893b888fdf27f15290b2dbcbe2fe0 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Thu, 19 Mar 2026 17:40:50 -0700 Subject: [PATCH 01/11] 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 9e86680e355..0c2cb8c827e 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 efc42d3355c..092b22cee6f 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 3e5b729cea6..b94bd875963 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 00000000000..20cb337cd15 --- /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 From cc7e27e094d0d0c3f1e3b31a73675d1cb096823c Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Thu, 19 Mar 2026 17:55:06 -0700 Subject: [PATCH 02/11] feat: update MCP CRUD endpoints to support mcp_server_manager role Co-Authored-By: Claude Opus 4.6 --- .../mcp_management_endpoints.py | 89 +++++++++++++------ .../test_mcp_manager_role.py | 84 +++++++++++++++++ 2 files changed, 147 insertions(+), 26 deletions(-) diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index b94bd875963..3400f90f489 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -1273,16 +1273,26 @@ if MCP_AVAILABLE: # Validate and normalize payload fields validate_and_normalize_mcp_server_payload(payload) - # AuthZ - restrict only proxy admins to create mcp servers - if LitellmUserRoles.PROXY_ADMIN != user_api_key_dict.user_role: - raise HTTPException( - status_code=status.HTTP_403_FORBIDDEN, - detail={ - "error": "User does not have permission to create mcp servers. You can only create mcp servers if you are a PROXY_ADMIN." - }, + # AuthZ - proxy admins or team MCP managers can create MCP servers + is_proxy_admin = LitellmUserRoles.PROXY_ADMIN == user_api_key_dict.user_role + manager_team_id: Optional[str] = None + + if not is_proxy_admin: + # Check if the user is an MCP manager for the specified team + if payload.team_id is None: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail={ + "error": "team_id is required when creating MCP servers as a team MCP manager." + }, + ) + manager_team_id = await _assert_can_manage_team_mcp_server( + user_api_key_dict=user_api_key_dict, + team_id=payload.team_id, ) - elif payload.server_id is not None: - # fail if the mcp server with id already exists + + # Fail if the MCP server with this id already exists + if payload.server_id is not None: mcp_server = await get_mcp_server(prisma_client, payload.server_id) if mcp_server is not None: raise HTTPException( @@ -1291,7 +1301,8 @@ if MCP_AVAILABLE: "error": f"MCP Server with id {payload.server_id} already exists. Cannot create another." }, ) - elif ( + + if ( SpecialMCPServerName.all_team_servers == payload.server_id or SpecialMCPServerName.all_proxy_servers == payload.server_id ): @@ -1322,6 +1333,40 @@ if MCP_AVAILABLE: # Ensure registry is up to date by reloading from database await global_mcp_server_manager.reload_servers_from_database() + + # If created by an MCP manager, auto-assign the server to their team + if manager_team_id is not None: + from litellm.proxy.proxy_server import user_api_key_cache + + team_obj = await get_team_object( + team_id=manager_team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + check_db_only=False, + ) + existing_permission_id = getattr( + team_obj, "object_permission_id", None + ) + + # Read existing mcp_servers list and append the new server + existing_mcp_servers: list = [] + if team_obj.object_permission is not None: + existing_mcp_servers = ( + team_obj.object_permission.mcp_servers or [] + ) + updated_mcp_servers = list( + set(existing_mcp_servers + [new_mcp_server.server_id]) + ) + + await handle_update_object_permission_common( + data_json={ + "object_permission": { + "mcp_servers": updated_mcp_servers, + } + }, + existing_object_permission_id=existing_permission_id, + prisma_client=prisma_client, + ) except Exception as e: verbose_proxy_logger.exception(f"Error creating mcp server: {str(e)}") raise HTTPException( @@ -1540,15 +1585,11 @@ if MCP_AVAILABLE: "Database not connected. Connect a database to your proxy - https://docs.litellm.ai/docs/simple_proxy#managing-auth---virtual-keys" ) - # Authz - restrict only admins to delete mcp servers + # Authz - proxy admins or team MCP managers can delete MCP servers if LitellmUserRoles.PROXY_ADMIN != user_api_key_dict.user_role: - raise HTTPException( - status_code=status.HTTP_403_FORBIDDEN, - detail={ - "error": "Call not allowed to delete MCP server. User is not a proxy admin. route={}".format( - "DELETE /v1/mcp/server" - ) - }, + await _assert_can_manage_team_mcp_server( + user_api_key_dict=user_api_key_dict, + server_id=server_id, ) # try to delete the mcp server @@ -1865,15 +1906,11 @@ if MCP_AVAILABLE: # Validate and normalize payload fields validate_and_normalize_mcp_server_payload(payload) - # Authz - restrict only admins to delete mcp servers + # Authz - proxy admins or team MCP managers can update MCP servers if LitellmUserRoles.PROXY_ADMIN != user_api_key_dict.user_role: - raise HTTPException( - status_code=status.HTTP_403_FORBIDDEN, - detail={ - "error": "Call not allowed to update MCP server. User is not a proxy admin. route={}".format( - "PUT /v1/mcp/server" - ) - }, + await _assert_can_manage_team_mcp_server( + user_api_key_dict=user_api_key_dict, + server_id=payload.server_id, ) # try to update the mcp server 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 index 20cb337cd15..860b255f9ca 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_mcp_manager_role.py +++ b/tests/test_litellm/proxy/management_endpoints/test_mcp_manager_role.py @@ -185,3 +185,87 @@ class TestAssertCanManageTeamMcpServer: user_api_key_dict=user_auth, server_id="server1" ) assert exc_info.value.status_code == 403 + + +@pytest.mark.asyncio +class TestCreateMcpServerAsManager: + async def test_create_auto_assigns_to_team(self): + """MCP manager creating a server should auto-assign it to their team's ObjectPermissionTable.""" + from litellm.proxy._types import NewMCPServerRequest + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + add_mcp_server, + ) + + user_auth = UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + user_id="user1", + api_key="sk-test", + team_id="team1", + ) + + payload = NewMCPServerRequest( + server_name="test-server", + url="https://example.com/mcp", + team_id="team1", + ) + + mock_team = LiteLLM_TeamTableCachedObj( + team_id="team1", + members_with_roles=[ + Member(user_id="user1", role="mcp_server_manager"), + ], + object_permission_id="perm1", + ) + mock_team.object_permission = MagicMock(mcp_servers=["existing_server"]) + + created_server = MagicMock() + created_server.server_id = "new_server_id" + created_server.credentials = None + + mock_handle_update = AsyncMock() + + with ( + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", + return_value=MagicMock(), + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.validate_and_normalize_mcp_server_payload", + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints._assert_can_manage_team_mcp_server", + AsyncMock(return_value="team1"), + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_mcp_server", + AsyncMock(return_value=None), + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.create_mcp_server", + AsyncMock(return_value=created_server), + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager", + MagicMock( + add_server=AsyncMock(), + reload_servers_from_database=AsyncMock(), + ), + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_team_object", + AsyncMock(return_value=mock_team), + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.handle_update_object_permission_common", + mock_handle_update, + ), + ): + await add_mcp_server(payload=payload, user_api_key_dict=user_auth) + + # Verify handle_update_object_permission_common was called with merged server list + mock_handle_update.assert_called_once() + call_kwargs = mock_handle_update.call_args.kwargs + mcp_servers = call_kwargs["data_json"]["object_permission"]["mcp_servers"] + assert "existing_server" in mcp_servers + assert "new_server_id" in mcp_servers + assert call_kwargs["existing_object_permission_id"] == "perm1" From 971cc2a9af40fa14b7a50efb428aab9098cc13ec Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Thu, 19 Mar 2026 17:56:08 -0700 Subject: [PATCH 03/11] fix: add team_id field to NewMCPServerRequest Co-Authored-By: Claude Opus 4.6 --- litellm/proxy/_types.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 0c2cb8c827e..67ae31bea8e 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -1144,6 +1144,10 @@ class NewMCPServerRequest(LiteLLMPydanticObjectBase): None, description="Server-managed: set by the endpoint; caller values are overridden.", ) + team_id: Optional[str] = Field( + None, + description="Team ID to assign the MCP server to. Required for team MCP managers.", + ) @model_validator(mode="before") @classmethod From f4e5fb2b4ae4d388d62e66eaba64c33a162bf8d6 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Thu, 19 Mar 2026 22:50:35 -0700 Subject: [PATCH 04/11] fix: exclude team_id from Prisma create payload in _prepare_mcp_server_data team_id is a request-level field, not a DB column. Excluding it prevents Prisma from rejecting the create call. Co-Authored-By: Claude Opus 4.6 --- litellm/proxy/_experimental/mcp_server/db.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/db.py b/litellm/proxy/_experimental/mcp_server/db.py index fbef33c32ed..0bc102889bc 100644 --- a/litellm/proxy/_experimental/mcp_server/db.py +++ b/litellm/proxy/_experimental/mcp_server/db.py @@ -40,8 +40,8 @@ def _prepare_mcp_server_data( """ from litellm.litellm_core_utils.safe_json_dumps import safe_dumps - # Convert model to dict - data_dict = data.model_dump(exclude_none=True) + # Convert model to dict, excluding fields not in the DB schema + data_dict = data.model_dump(exclude_none=True, exclude={"team_id"}) # Ensure alias is always present in the dict (even if None) if "alias" not in data_dict: data_dict["alias"] = getattr(data, "alias", None) From 6c06e6dc80fd7a2870a3f6fef57a8d651de1a82b Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Thu, 19 Mar 2026 22:57:27 -0700 Subject: [PATCH 05/11] fix: link new object_permission_id back to team on first MCP server create When a team has no object_permission_id yet, the auto-assign logic creates an ObjectPermissionTable row but never linked it to the team. Now updates the team's object_permission_id after creation. Co-Authored-By: Claude Opus 4.6 --- .../management_endpoints/mcp_management_endpoints.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index 3400f90f489..522a533d43f 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -1358,7 +1358,7 @@ if MCP_AVAILABLE: set(existing_mcp_servers + [new_mcp_server.server_id]) ) - await handle_update_object_permission_common( + new_permission_id = await handle_update_object_permission_common( data_json={ "object_permission": { "mcp_servers": updated_mcp_servers, @@ -1367,6 +1367,13 @@ if MCP_AVAILABLE: existing_object_permission_id=existing_permission_id, prisma_client=prisma_client, ) + + # If the team had no object_permission_id, link the new one + if existing_permission_id is None and new_permission_id is not None: + await prisma_client.db.litellm_teamtable.update( + where={"team_id": manager_team_id}, + data={"object_permission_id": new_permission_id}, + ) except Exception as e: verbose_proxy_logger.exception(f"Error creating mcp server: {str(e)}") raise HTTPException( From dfc4401e1c305f11c898e989991503122668bede Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Thu, 19 Mar 2026 23:00:38 -0700 Subject: [PATCH 06/11] fix: use check_db_only=True in _assert_can_manage_team_mcp_server Cached team objects may have stale object_permission data. Bypassing cache ensures the server-in-team check uses fresh data. Co-Authored-By: Claude Opus 4.6 --- litellm/proxy/management_endpoints/mcp_management_endpoints.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index 522a533d43f..a3b0b3ab286 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -179,7 +179,7 @@ if MCP_AVAILABLE: team_id=resolved_team_id, prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, - check_db_only=False, + check_db_only=True, # bypass cache to get fresh object_permission ) if not _is_user_team_mcp_manager(user_api_key_dict, team_obj): From e71410470267f64a00f61d7b75094a8628054b0e Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Thu, 19 Mar 2026 23:10:44 -0700 Subject: [PATCH 07/11] fix: address Greptile review feedback - P0: Separate try/except for auto-assign so server creation succeeds even if team permission update fails - P1: Clean up team permission entry on MCP server delete - P1: Add MCP_AVAILABLE skip guard to tests - P2: Return team_obj from _assert_can_manage_team_mcp_server to eliminate redundant get_team_object call in create endpoint Co-Authored-By: Claude Opus 4.6 --- .../mcp_management_endpoints.py | 80 +++++++++++++------ .../test_mcp_manager_role.py | 21 +++-- 2 files changed, 72 insertions(+), 29 deletions(-) diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index a3b0b3ab286..f137902f1b7 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -21,7 +21,7 @@ import json import os from dataclasses import dataclass from datetime import datetime, timedelta, timezone -from typing import Any, Dict, Iterable, List, Literal, Optional +from typing import Any, Dict, Iterable, List, Literal, Optional, Tuple from fastapi import ( APIRouter, @@ -147,12 +147,12 @@ if MCP_AVAILABLE: user_api_key_dict: UserAPIKeyAuth, team_id: Optional[str] = None, server_id: Optional[str] = None, - ) -> str: + ) -> Tuple[str, "LiteLLM_TeamTableCachedObj"]: """ 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. + Returns a tuple of (team_id, team_obj) for downstream use. 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. """ @@ -165,6 +165,7 @@ if MCP_AVAILABLE: detail={"error": "team_id is required for MCP server manager operations."}, ) + # When the API key is team-scoped, ensure the request team_id matches if ( team_id and user_api_key_dict.team_id @@ -200,7 +201,7 @@ if MCP_AVAILABLE: }, ) - return resolved_team_id + return resolved_team_id, team_obj @dataclass class _TemporaryMCPServerEntry: @@ -1277,6 +1278,7 @@ if MCP_AVAILABLE: is_proxy_admin = LitellmUserRoles.PROXY_ADMIN == user_api_key_dict.user_role manager_team_id: Optional[str] = None + manager_team_obj = None if not is_proxy_admin: # Check if the user is an MCP manager for the specified team if payload.team_id is None: @@ -1286,7 +1288,7 @@ if MCP_AVAILABLE: "error": "team_id is required when creating MCP servers as a team MCP manager." }, ) - manager_team_id = await _assert_can_manage_team_mcp_server( + manager_team_id, manager_team_obj = await _assert_can_manage_team_mcp_server( user_api_key_dict=user_api_key_dict, team_id=payload.team_id, ) @@ -1334,25 +1336,26 @@ if MCP_AVAILABLE: # Ensure registry is up to date by reloading from database await global_mcp_server_manager.reload_servers_from_database() - # If created by an MCP manager, auto-assign the server to their team - if manager_team_id is not None: - from litellm.proxy.proxy_server import user_api_key_cache + except Exception as e: + verbose_proxy_logger.exception(f"Error creating mcp server: {str(e)}") + raise HTTPException( + status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, + detail={"error": f"Error creating mcp server: {str(e)}"}, + ) - team_obj = await get_team_object( - team_id=manager_team_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - check_db_only=False, - ) + # Auto-assign the server to the manager's team (separate from create + # so a failure here doesn't mask the successfully created server). + if manager_team_id is not None and manager_team_obj is not None: + try: existing_permission_id = getattr( - team_obj, "object_permission_id", None + manager_team_obj, "object_permission_id", None ) # Read existing mcp_servers list and append the new server existing_mcp_servers: list = [] - if team_obj.object_permission is not None: + if manager_team_obj.object_permission is not None: existing_mcp_servers = ( - team_obj.object_permission.mcp_servers or [] + manager_team_obj.object_permission.mcp_servers or [] ) updated_mcp_servers = list( set(existing_mcp_servers + [new_mcp_server.server_id]) @@ -1374,12 +1377,10 @@ if MCP_AVAILABLE: where={"team_id": manager_team_id}, data={"object_permission_id": new_permission_id}, ) - except Exception as e: - verbose_proxy_logger.exception(f"Error creating mcp server: {str(e)}") - raise HTTPException( - status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, - detail={"error": f"Error creating mcp server: {str(e)}"}, - ) + except Exception as e: + verbose_proxy_logger.exception( + f"MCP server created but failed to auto-assign to team {manager_team_id}: {str(e)}" + ) return _redact_mcp_credentials(new_mcp_server) @router.post( @@ -1593,8 +1594,9 @@ if MCP_AVAILABLE: ) # Authz - proxy admins or team MCP managers can delete MCP servers + manager_team_obj = None if LitellmUserRoles.PROXY_ADMIN != user_api_key_dict.user_role: - await _assert_can_manage_team_mcp_server( + _, manager_team_obj = await _assert_can_manage_team_mcp_server( user_api_key_dict=user_api_key_dict, server_id=server_id, ) @@ -1612,6 +1614,36 @@ if MCP_AVAILABLE: # Ensure registry is up to date by reloading from database await global_mcp_server_manager.reload_servers_from_database() + # Remove server from the manager's team permission list + if manager_team_obj is not None: + try: + existing_permission_id = getattr( + manager_team_obj, "object_permission_id", None + ) + if ( + existing_permission_id is not None + and manager_team_obj.object_permission is not None + ): + existing_mcp_servers = ( + manager_team_obj.object_permission.mcp_servers or [] + ) + updated_mcp_servers = [ + s for s in existing_mcp_servers if s != server_id + ] + await handle_update_object_permission_common( + data_json={ + "object_permission": { + "mcp_servers": updated_mcp_servers, + } + }, + existing_object_permission_id=existing_permission_id, + prisma_client=prisma_client, + ) + except Exception as e: + verbose_proxy_logger.exception( + f"MCP server deleted but failed to remove from team permissions: {str(e)}" + ) + # TODO: Enterprise: Finish audit log trail if litellm.store_audit_logs: pass 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 index 860b255f9ca..f4e91157f21 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_mcp_manager_role.py +++ b/tests/test_litellm/proxy/management_endpoints/test_mcp_manager_role.py @@ -9,6 +9,9 @@ from litellm.proxy._types import ( from litellm.proxy.management_endpoints.common_utils import ( _is_user_team_mcp_manager, ) +from litellm.proxy.management_endpoints import ( + mcp_management_endpoints as mgmt_endpoints, +) class TestIsUserTeamMcpManager: @@ -70,6 +73,9 @@ from litellm.proxy._types import LiteLLM_TeamTableCachedObj @pytest.mark.asyncio +@pytest.mark.skipif( + not mgmt_endpoints.MCP_AVAILABLE, reason="MCP module not installed" +) class TestAssertCanManageTeamMcpServer: async def test_mcp_manager_with_team_id_succeeds(self): from litellm.proxy.management_endpoints.mcp_management_endpoints import ( @@ -86,10 +92,11 @@ class TestAssertCanManageTeamMcpServer: "litellm.proxy.management_endpoints.mcp_management_endpoints.get_team_object", AsyncMock(return_value=mock_team), ): - result = await _assert_can_manage_team_mcp_server( + team_id, team_obj = await _assert_can_manage_team_mcp_server( user_api_key_dict=user_auth, team_id="team1" ) - assert result == "team1" + assert team_id == "team1" + assert team_obj == mock_team async def test_regular_user_gets_403(self): from litellm.proxy.management_endpoints.mcp_management_endpoints import ( @@ -154,10 +161,11 @@ class TestAssertCanManageTeamMcpServer: AsyncMock(return_value={"server1", "server2"}), ), ): - result = await _assert_can_manage_team_mcp_server( + team_id, team_obj = await _assert_can_manage_team_mcp_server( user_api_key_dict=user_auth, server_id="server1" ) - assert result == "team1" + assert team_id == "team1" + assert team_obj == mock_team async def test_mcp_manager_server_not_in_team_gets_403(self): from litellm.proxy.management_endpoints.mcp_management_endpoints import ( @@ -188,6 +196,9 @@ class TestAssertCanManageTeamMcpServer: @pytest.mark.asyncio +@pytest.mark.skipif( + not mgmt_endpoints.MCP_AVAILABLE, reason="MCP module not installed" +) class TestCreateMcpServerAsManager: async def test_create_auto_assigns_to_team(self): """MCP manager creating a server should auto-assign it to their team's ObjectPermissionTable.""" @@ -234,7 +245,7 @@ class TestCreateMcpServerAsManager: ), patch( "litellm.proxy.management_endpoints.mcp_management_endpoints._assert_can_manage_team_mcp_server", - AsyncMock(return_value="team1"), + AsyncMock(return_value=("team1", mock_team)), ), patch( "litellm.proxy.management_endpoints.mcp_management_endpoints.get_mcp_server", From 65bfd449e980ac08e6339af7a569a69dfa38b192 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Thu, 19 Mar 2026 23:23:18 -0700 Subject: [PATCH 08/11] address greptile review feedback (greploop iteration 2) - Remove dead get_team_object mock in test (now reuses team_obj from _assert) - Add test for existing_permission_id=None branch (team linkage) - Remaining P1s are by-design per spec (admin blocked from MCP, team_id ignored for proxy admin) Co-Authored-By: Claude Opus 4.6 --- .../test_mcp_manager_role.py | 81 ++++++++++++++++++- 1 file changed, 77 insertions(+), 4 deletions(-) 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 index f4e91157f21..67e9dda4cef 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_mcp_manager_role.py +++ b/tests/test_litellm/proxy/management_endpoints/test_mcp_manager_role.py @@ -262,10 +262,6 @@ class TestCreateMcpServerAsManager: reload_servers_from_database=AsyncMock(), ), ), - patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints.get_team_object", - AsyncMock(return_value=mock_team), - ), patch( "litellm.proxy.management_endpoints.mcp_management_endpoints.handle_update_object_permission_common", mock_handle_update, @@ -280,3 +276,80 @@ class TestCreateMcpServerAsManager: assert "existing_server" in mcp_servers assert "new_server_id" in mcp_servers assert call_kwargs["existing_object_permission_id"] == "perm1" + + async def test_create_links_new_permission_to_team_when_none_exists(self): + """When team has no object_permission_id, create should link the new one.""" + from litellm.proxy._types import NewMCPServerRequest + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + add_mcp_server, + ) + + user_auth = UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + user_id="user1", + api_key="sk-test", + team_id="team1", + ) + + payload = NewMCPServerRequest( + server_name="test_server", + url="https://example.com/mcp", + team_id="team1", + ) + + # Team with NO object_permission_id + mock_team = LiteLLM_TeamTableCachedObj( + team_id="team1", + members_with_roles=[ + Member(user_id="user1", role="mcp_server_manager"), + ], + object_permission_id=None, + ) + mock_team.object_permission = None + + created_server = MagicMock() + created_server.server_id = "new_server_id" + created_server.credentials = None + + mock_prisma = MagicMock() + mock_prisma.db.litellm_teamtable.update = AsyncMock() + + with ( + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", + return_value=mock_prisma, + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.validate_and_normalize_mcp_server_payload", + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints._assert_can_manage_team_mcp_server", + AsyncMock(return_value=("team1", mock_team)), + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_mcp_server", + AsyncMock(return_value=None), + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.create_mcp_server", + AsyncMock(return_value=created_server), + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager", + MagicMock( + add_server=AsyncMock(), + reload_servers_from_database=AsyncMock(), + ), + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.handle_update_object_permission_common", + AsyncMock(return_value="new_perm_id"), + ), + ): + await add_mcp_server(payload=payload, user_api_key_dict=user_auth) + + # Verify the team was updated with the new object_permission_id + mock_prisma.db.litellm_teamtable.update.assert_called_once_with( + where={"team_id": "team1"}, + data={"object_permission_id": "new_perm_id"}, + ) From b8c9bf7d256ed90b1f8cfdbe8b57d79f118009f4 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Fri, 20 Mar 2026 09:05:18 -0700 Subject: [PATCH 09/11] refactor: extract _auto_assign and _remove helpers, use team_endpoints helper - Replace raw prisma_client.db.litellm_teamtable.update with handle_update_object_permission from team_endpoints (follows established helper-function pattern) - Extract _auto_assign_mcp_server_to_team and _remove_mcp_server_from_team helpers for reuse and testability - Update tests to mock at the correct boundaries Co-Authored-By: Claude Opus 4.6 --- .../mcp_management_endpoints.py | 119 ++++++++++-------- .../test_mcp_manager_role.py | 89 ++++--------- 2 files changed, 96 insertions(+), 112 deletions(-) diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index f137902f1b7..e137ad15f1b 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -203,6 +203,69 @@ if MCP_AVAILABLE: return resolved_team_id, team_obj + async def _auto_assign_mcp_server_to_team( + server_id: str, + team_id: str, + team_obj: "LiteLLM_TeamTableCachedObj", + prisma_client: Any, + ) -> None: + """ + Add an MCP server to a team's ObjectPermissionTable and link the + permission back to the team if it didn't have one yet. + + Uses handle_update_object_permission (the team-endpoint helper) so + the object_permission_id linkage follows the same pattern as + team_endpoints.update_team. + """ + from litellm.proxy.management_endpoints.team_endpoints import ( + handle_update_object_permission, + ) + + existing_mcp_servers: list = [] + if team_obj.object_permission is not None: + existing_mcp_servers = team_obj.object_permission.mcp_servers or [] + updated_mcp_servers = list(set(existing_mcp_servers + [server_id])) + + # Build the data dict in the same shape team_endpoints uses + data_json: Dict[str, Any] = { + "object_permission": {"mcp_servers": updated_mcp_servers}, + } + data_json = await handle_update_object_permission( + data_json=data_json, + existing_team_row=team_obj, + ) + + # If handle_update_object_permission produced an object_permission_id, + # persist it on the team row (it sets data_json["object_permission_id"]). + if "object_permission_id" in data_json: + await prisma_client.db.litellm_teamtable.update( + where={"team_id": team_id}, + data={"object_permission_id": data_json["object_permission_id"]}, + ) + + async def _remove_mcp_server_from_team( + server_id: str, + team_obj: "LiteLLM_TeamTableCachedObj", + ) -> None: + """Remove a server ID from a team's ObjectPermissionTable.mcp_servers list.""" + from litellm.proxy.proxy_server import prisma_client + + existing_permission_id = getattr(team_obj, "object_permission_id", None) + if ( + existing_permission_id is None + or team_obj.object_permission is None + ): + return + + existing_mcp_servers = team_obj.object_permission.mcp_servers or [] + updated_mcp_servers = [s for s in existing_mcp_servers if s != server_id] + + await handle_update_object_permission_common( + data_json={"object_permission": {"mcp_servers": updated_mcp_servers}}, + existing_object_permission_id=existing_permission_id, + prisma_client=prisma_client, + ) + @dataclass class _TemporaryMCPServerEntry: server: MCPServer @@ -1347,36 +1410,12 @@ if MCP_AVAILABLE: # so a failure here doesn't mask the successfully created server). if manager_team_id is not None and manager_team_obj is not None: try: - existing_permission_id = getattr( - manager_team_obj, "object_permission_id", None - ) - - # Read existing mcp_servers list and append the new server - existing_mcp_servers: list = [] - if manager_team_obj.object_permission is not None: - existing_mcp_servers = ( - manager_team_obj.object_permission.mcp_servers or [] - ) - updated_mcp_servers = list( - set(existing_mcp_servers + [new_mcp_server.server_id]) - ) - - new_permission_id = await handle_update_object_permission_common( - data_json={ - "object_permission": { - "mcp_servers": updated_mcp_servers, - } - }, - existing_object_permission_id=existing_permission_id, + await _auto_assign_mcp_server_to_team( + server_id=new_mcp_server.server_id, + team_id=manager_team_id, + team_obj=manager_team_obj, prisma_client=prisma_client, ) - - # If the team had no object_permission_id, link the new one - if existing_permission_id is None and new_permission_id is not None: - await prisma_client.db.litellm_teamtable.update( - where={"team_id": manager_team_id}, - data={"object_permission_id": new_permission_id}, - ) except Exception as e: verbose_proxy_logger.exception( f"MCP server created but failed to auto-assign to team {manager_team_id}: {str(e)}" @@ -1617,28 +1656,10 @@ if MCP_AVAILABLE: # Remove server from the manager's team permission list if manager_team_obj is not None: try: - existing_permission_id = getattr( - manager_team_obj, "object_permission_id", None + await _remove_mcp_server_from_team( + server_id=server_id, + team_obj=manager_team_obj, ) - if ( - existing_permission_id is not None - and manager_team_obj.object_permission is not None - ): - existing_mcp_servers = ( - manager_team_obj.object_permission.mcp_servers or [] - ) - updated_mcp_servers = [ - s for s in existing_mcp_servers if s != server_id - ] - await handle_update_object_permission_common( - data_json={ - "object_permission": { - "mcp_servers": updated_mcp_servers, - } - }, - existing_object_permission_id=existing_permission_id, - prisma_client=prisma_client, - ) except Exception as e: verbose_proxy_logger.exception( f"MCP server deleted but failed to remove from team permissions: {str(e)}" 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 index 67e9dda4cef..ed601dfffa4 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_mcp_manager_role.py +++ b/tests/test_litellm/proxy/management_endpoints/test_mcp_manager_role.py @@ -233,7 +233,7 @@ class TestCreateMcpServerAsManager: created_server.server_id = "new_server_id" created_server.credentials = None - mock_handle_update = AsyncMock() + mock_auto_assign = AsyncMock() with ( patch( @@ -263,38 +263,23 @@ class TestCreateMcpServerAsManager: ), ), patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints.handle_update_object_permission_common", - mock_handle_update, + "litellm.proxy.management_endpoints.mcp_management_endpoints._auto_assign_mcp_server_to_team", + mock_auto_assign, ), ): await add_mcp_server(payload=payload, user_api_key_dict=user_auth) - # Verify handle_update_object_permission_common was called with merged server list - mock_handle_update.assert_called_once() - call_kwargs = mock_handle_update.call_args.kwargs - mcp_servers = call_kwargs["data_json"]["object_permission"]["mcp_servers"] - assert "existing_server" in mcp_servers - assert "new_server_id" in mcp_servers - assert call_kwargs["existing_object_permission_id"] == "perm1" + # Verify _auto_assign_mcp_server_to_team was called with the right args + mock_auto_assign.assert_called_once() + call_kwargs = mock_auto_assign.call_args.kwargs + assert call_kwargs["server_id"] == "new_server_id" + assert call_kwargs["team_id"] == "team1" + assert call_kwargs["team_obj"] == mock_team - async def test_create_links_new_permission_to_team_when_none_exists(self): - """When team has no object_permission_id, create should link the new one.""" - from litellm.proxy._types import NewMCPServerRequest + async def test_auto_assign_links_new_permission_to_team(self): + """_auto_assign_mcp_server_to_team should create permission and link to team.""" from litellm.proxy.management_endpoints.mcp_management_endpoints import ( - add_mcp_server, - ) - - user_auth = UserAPIKeyAuth( - user_role=LitellmUserRoles.INTERNAL_USER, - user_id="user1", - api_key="sk-test", - team_id="team1", - ) - - payload = NewMCPServerRequest( - server_name="test_server", - url="https://example.com/mcp", - team_id="team1", + _auto_assign_mcp_server_to_team, ) # Team with NO object_permission_id @@ -307,46 +292,24 @@ class TestCreateMcpServerAsManager: ) mock_team.object_permission = None - created_server = MagicMock() - created_server.server_id = "new_server_id" - created_server.credentials = None - mock_prisma = MagicMock() mock_prisma.db.litellm_teamtable.update = AsyncMock() - with ( - patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", - return_value=mock_prisma, - ), - patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints.validate_and_normalize_mcp_server_payload", - ), - patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints._assert_can_manage_team_mcp_server", - AsyncMock(return_value=("team1", mock_team)), - ), - patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints.get_mcp_server", - AsyncMock(return_value=None), - ), - patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints.create_mcp_server", - AsyncMock(return_value=created_server), - ), - patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager", - MagicMock( - add_server=AsyncMock(), - reload_servers_from_database=AsyncMock(), - ), - ), - patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints.handle_update_object_permission_common", - AsyncMock(return_value="new_perm_id"), - ), + # handle_update_object_permission sets object_permission_id in data_json + async def fake_handle(data_json, existing_team_row): + data_json["object_permission_id"] = "new_perm_id" + return data_json + + with patch( + "litellm.proxy.management_endpoints.team_endpoints.handle_update_object_permission", + side_effect=fake_handle, ): - await add_mcp_server(payload=payload, user_api_key_dict=user_auth) + await _auto_assign_mcp_server_to_team( + server_id="new_server_id", + team_id="team1", + team_obj=mock_team, + prisma_client=mock_prisma, + ) # Verify the team was updated with the new object_permission_id mock_prisma.db.litellm_teamtable.update.assert_called_once_with( From 1cd7a48c33cd72239913ccf73ecd87c3be4d5a30 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Fri, 20 Mar 2026 09:19:53 -0700 Subject: [PATCH 10/11] Add tests for edit and delete MCP server manager paths Addresses Greptile feedback about missing integration tests for PUT/DELETE when invoked by mcp_server_manager role. Adds tests for edit success/403, delete success with team cleanup/403, and the _remove_mcp_server_from_team helper directly. Co-Authored-By: Claude Opus 4.6 --- .../test_mcp_manager_role.py | 239 ++++++++++++++++++ 1 file changed, 239 insertions(+) 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 index ed601dfffa4..36cecde65b9 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_mcp_manager_role.py +++ b/tests/test_litellm/proxy/management_endpoints/test_mcp_manager_role.py @@ -69,6 +69,7 @@ class TestIsUserTeamMcpManager: from unittest.mock import AsyncMock, MagicMock, patch +from fastapi import HTTPException from litellm.proxy._types import LiteLLM_TeamTableCachedObj @@ -316,3 +317,241 @@ class TestCreateMcpServerAsManager: where={"team_id": "team1"}, data={"object_permission_id": "new_perm_id"}, ) + + +@pytest.mark.asyncio +@pytest.mark.skipif( + not mgmt_endpoints.MCP_AVAILABLE, reason="MCP module not installed" +) +class TestEditMcpServerAsManager: + async def test_edit_succeeds_for_mcp_manager(self): + """MCP manager should be able to edit a server assigned to their team.""" + from litellm.proxy._types import UpdateMCPServerRequest + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + edit_mcp_server, + ) + + user_auth = UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + user_id="user1", + api_key="sk-test", + team_id="team1", + ) + + payload = UpdateMCPServerRequest( + server_id="server1", + description="Updated description", + ) + + mock_team = LiteLLM_TeamTableCachedObj( + team_id="team1", + members_with_roles=[ + Member(user_id="user1", role="mcp_server_manager"), + ], + ) + + updated_server = MagicMock() + updated_server.server_id = "server1" + updated_server.credentials = None + + with ( + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", + return_value=MagicMock(), + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.validate_and_normalize_mcp_server_payload", + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints._assert_can_manage_team_mcp_server", + AsyncMock(return_value=("team1", mock_team)), + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.update_mcp_server", + AsyncMock(return_value=updated_server), + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager", + MagicMock( + update_server=AsyncMock(), + reload_servers_from_database=AsyncMock(), + ), + ), + ): + result = await edit_mcp_server(payload=payload, user_api_key_dict=user_auth) + # Should not raise — edit succeeded + assert result is not None + + async def test_edit_fails_for_server_not_in_team(self): + """MCP manager should get 403 when editing a server not in their team.""" + from litellm.proxy._types import UpdateMCPServerRequest + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + edit_mcp_server, + ) + + user_auth = UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + user_id="user1", + api_key="sk-test", + team_id="team1", + ) + + payload = UpdateMCPServerRequest( + server_id="server_not_in_team", + description="Updated description", + ) + + with ( + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", + return_value=MagicMock(), + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.validate_and_normalize_mcp_server_payload", + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints._assert_can_manage_team_mcp_server", + AsyncMock(side_effect=HTTPException(status_code=403, detail="Not in team")), + ), + ): + with pytest.raises(HTTPException) as exc_info: + await edit_mcp_server(payload=payload, user_api_key_dict=user_auth) + assert exc_info.value.status_code == 403 + + +@pytest.mark.asyncio +@pytest.mark.skipif( + not mgmt_endpoints.MCP_AVAILABLE, reason="MCP module not installed" +) +class TestDeleteMcpServerAsManager: + async def test_delete_succeeds_and_cleans_up_team(self): + """MCP manager deleting a server should also remove it from team permissions.""" + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + remove_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"), + ], + ) + + deleted_server = MagicMock() + deleted_server.server_id = "server1" + + mock_remove_from_team = AsyncMock() + + with ( + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", + return_value=MagicMock(), + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints._assert_can_manage_team_mcp_server", + AsyncMock(return_value=("team1", mock_team)), + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.delete_mcp_server", + AsyncMock(return_value=deleted_server), + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager", + MagicMock( + remove_server=MagicMock(), + reload_servers_from_database=AsyncMock(), + ), + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints._remove_mcp_server_from_team", + mock_remove_from_team, + ), + ): + response = await remove_mcp_server( + server_id="server1", user_api_key_dict=user_auth + ) + assert response.status_code == 202 + + # Verify team cleanup was called + mock_remove_from_team.assert_called_once() + call_kwargs = mock_remove_from_team.call_args.kwargs + assert call_kwargs["server_id"] == "server1" + assert call_kwargs["team_obj"] == mock_team + + async def test_delete_fails_for_server_not_in_team(self): + """MCP manager should get 403 when deleting a server not in their team.""" + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + remove_mcp_server, + ) + + user_auth = UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + user_id="user1", + api_key="sk-test", + team_id="team1", + ) + + with ( + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", + return_value=MagicMock(), + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints._assert_can_manage_team_mcp_server", + AsyncMock(side_effect=HTTPException(status_code=403, detail="Not in team")), + ), + ): + with pytest.raises(HTTPException) as exc_info: + await remove_mcp_server( + server_id="server_not_in_team", user_api_key_dict=user_auth + ) + assert exc_info.value.status_code == 403 + + async def test_remove_mcp_server_from_team_helper(self): + """_remove_mcp_server_from_team should update the permission list without the deleted server.""" + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + _remove_mcp_server_from_team, + ) + + mock_team = LiteLLM_TeamTableCachedObj( + team_id="team1", + members_with_roles=[ + Member(user_id="user1", role="mcp_server_manager"), + ], + object_permission_id="perm1", + ) + mock_team.object_permission = MagicMock( + mcp_servers=["server1", "server2", "server3"] + ) + + mock_handle_common = AsyncMock() + + with ( + patch( + "litellm.proxy.proxy_server.prisma_client", + MagicMock(), + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.handle_update_object_permission_common", + mock_handle_common, + ), + ): + await _remove_mcp_server_from_team( + server_id="server2", + team_obj=mock_team, + ) + + mock_handle_common.assert_called_once() + call_kwargs = mock_handle_common.call_args.kwargs + updated_servers = call_kwargs["data_json"]["object_permission"]["mcp_servers"] + assert "server2" not in updated_servers + assert "server1" in updated_servers + assert "server3" in updated_servers + assert call_kwargs["existing_object_permission_id"] == "perm1" From 700fd86de971f0ca50bafafed129ab17af661f3a Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Fri, 20 Mar 2026 12:13:09 -0700 Subject: [PATCH 11/11] Fix importorskip guard and add LiteLLM_TeamTableCachedObj import - Add pytest.importorskip("mcp") at module level so tests skip cleanly in CI environments without the mcp package (instead of ImportError) - Import LiteLLM_TeamTableCachedObj into MCP_AVAILABLE block so type annotations resolve for static analysis and get_type_hints() - Remove string quotes from type annotations now that the import exists Co-Authored-By: Claude Opus 4.6 --- .../proxy/management_endpoints/mcp_management_endpoints.py | 7 ++++--- .../proxy/management_endpoints/test_mcp_manager_role.py | 2 ++ 2 files changed, 6 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index e137ad15f1b..8c8bd629884 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -112,6 +112,7 @@ if MCP_AVAILABLE: ) from litellm.proxy._types import ( LiteLLM_MCPServerTable, + LiteLLM_TeamTableCachedObj, LitellmUserRoles, MakeMCPServersPublicRequest, MCPApprovalStatus, @@ -147,7 +148,7 @@ if MCP_AVAILABLE: user_api_key_dict: UserAPIKeyAuth, team_id: Optional[str] = None, server_id: Optional[str] = None, - ) -> Tuple[str, "LiteLLM_TeamTableCachedObj"]: + ) -> Tuple[str, LiteLLM_TeamTableCachedObj]: """ Verify that the caller is an MCP server manager for a team and (for edit/delete) that the target server belongs to that team. @@ -206,7 +207,7 @@ if MCP_AVAILABLE: async def _auto_assign_mcp_server_to_team( server_id: str, team_id: str, - team_obj: "LiteLLM_TeamTableCachedObj", + team_obj: LiteLLM_TeamTableCachedObj, prisma_client: Any, ) -> None: """ @@ -245,7 +246,7 @@ if MCP_AVAILABLE: async def _remove_mcp_server_from_team( server_id: str, - team_obj: "LiteLLM_TeamTableCachedObj", + team_obj: LiteLLM_TeamTableCachedObj, ) -> None: """Remove a server ID from a team's ObjectPermissionTable.mcp_servers list.""" from litellm.proxy.proxy_server import prisma_client 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 index 36cecde65b9..86b0419fb22 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_mcp_manager_role.py +++ b/tests/test_litellm/proxy/management_endpoints/test_mcp_manager_role.py @@ -1,5 +1,7 @@ import pytest +pytest.importorskip("mcp", reason="mcp package not installed") + from litellm.proxy._types import ( LiteLLM_TeamTable, LitellmUserRoles,