From e71410470267f64a00f61d7b75094a8628054b0e Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Thu, 19 Mar 2026 23:10:44 -0700 Subject: [PATCH] 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",