diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 1e7184db93..81c69dc6e6 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -371,6 +371,8 @@ class LiteLLMRoutes(enum.Enum): "/team/block", "/team/unblock", "/team/available", + "/team/permissions_list", + "/team/permissions_update", # model "/model/new", "/model/update", @@ -456,6 +458,8 @@ class LiteLLMRoutes(enum.Enum): self_managed_routes = [ "/team/member_add", "/team/member_delete", + "/team/permissions_list", + "/team/permissions_update", "/model/new", "/model/update", "/model/delete", diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index c6893ef4a0..c8200f7fed 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -1952,30 +1952,62 @@ async def team_member_permissions( """ Get the team member permissions for a team """ - from litellm.proxy.proxy_server import prisma_client + from litellm.proxy.proxy_server import ( + prisma_client, + proxy_logging_obj, + user_api_key_cache, + ) if prisma_client is None: raise HTTPException(status_code=500, detail={"error": "No db connected"}) - team_row = await prisma_client.db.litellm_teamtable.find_unique( - where={"team_id": team_id} + ## CHECK IF USER IS PROXY ADMIN OR TEAM ADMIN + existing_team_row = await get_team_object( + team_id=team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=None, + proxy_logging_obj=proxy_logging_obj, + check_cache_only=False, + check_db_only=True, ) - - if team_row is None: + if existing_team_row is None: raise HTTPException( status_code=404, - detail={"error": f"Team not found, passed team_id={team_id}"}, + detail={"error": f"Team not found for team_id={team_id}"}, ) - team_obj = LiteLLM_TeamTable(**team_row.model_dump()) - if team_obj.team_member_permissions is None: - team_obj.team_member_permissions = ( + complete_team_data = LiteLLM_TeamTable(**existing_team_row.model_dump()) + + if ( + hasattr(user_api_key_dict, "user_role") + and user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value + and not _is_user_team_admin( + user_api_key_dict=user_api_key_dict, team_obj=complete_team_data + ) + and not _is_available_team( + team_id=complete_team_data.team_id, + user_api_key_dict=user_api_key_dict, + ) + ): + raise HTTPException( + status_code=403, + detail={ + "error": "Call not allowed. User not proxy admin OR team admin. route={}, team_id={}".format( + "/team/member_add", + complete_team_data.team_id, + ) + }, + ) + + if existing_team_row.team_member_permissions is None: + existing_team_row.team_member_permissions = ( TeamMemberPermissionChecks.default_team_member_permissions() ) return GetTeamMemberPermissionsResponse( team_id=team_id, - team_member_permissions=team_obj.team_member_permissions, + team_member_permissions=existing_team_row.team_member_permissions, all_available_permissions=TeamMemberPermissionChecks.get_all_available_team_member_permissions(), ) @@ -1993,27 +2025,53 @@ async def update_team_member_permissions( """ Update the team member permissions for a team """ - from litellm.proxy.proxy_server import prisma_client + from litellm.proxy.proxy_server import ( + prisma_client, + proxy_logging_obj, + user_api_key_cache, + ) if prisma_client is None: raise HTTPException(status_code=500, detail={"error": "No db connected"}) - team_row = await prisma_client.db.litellm_teamtable.find_unique( - where={"team_id": data.team_id} + ## CHECK IF USER IS PROXY ADMIN OR TEAM ADMIN + existing_team_row = await get_team_object( + team_id=data.team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=None, + proxy_logging_obj=proxy_logging_obj, + check_cache_only=False, + check_db_only=True, ) - - if team_row is None: + if existing_team_row is None: raise HTTPException( status_code=404, - detail={"error": f"Team not found, passed team_id={data.team_id}"}, + detail={"error": f"Team not found for team_id={data.team_id}"}, ) - if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value: + complete_team_data = LiteLLM_TeamTable(**existing_team_row.model_dump()) + + if ( + hasattr(user_api_key_dict, "user_role") + and user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value + and not _is_user_team_admin( + user_api_key_dict=user_api_key_dict, team_obj=complete_team_data + ) + and not _is_available_team( + team_id=complete_team_data.team_id, + user_api_key_dict=user_api_key_dict, + ) + ): raise HTTPException( status_code=403, - detail={"error": "Only proxy admin can update team member permissions"}, + detail={ + "error": "Call not allowed. User not proxy admin OR team admin. route={}, team_id={}".format( + "/team/member_add", + complete_team_data.team_id, + ) + }, ) - # Update the team member permissions updated_team = await prisma_client.db.litellm_teamtable.update( where={"team_id": data.team_id}, diff --git a/tests/litellm/proxy/management_endpoints/test_team_endpoints.py b/tests/litellm/proxy/management_endpoints/test_team_endpoints.py index d3f3c11593..ec47196ae0 100644 --- a/tests/litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/litellm/proxy/management_endpoints/test_team_endpoints.py @@ -59,40 +59,47 @@ async def test_get_team_permissions_list_success(mock_db_client, mock_admin_auth Test successful retrieval of team member permissions. """ test_team_id = "test-team-123" + permissions = ["/key/generate", "/key/update"] mock_team_data = { "team_id": test_team_id, "team_alias": "Test Team", - "team_member_permissions": ["/key/generate", "/key/update"], + "team_member_permissions": permissions, "spend": 0.0, } mock_team_row = MagicMock() mock_team_row.model_dump.return_value = mock_team_data - mock_db_client.db.litellm_teamtable.find_unique = AsyncMock( - return_value=mock_team_row - ) - # Override the dependency for this test - app.dependency_overrides[user_api_key_auth] = lambda: mock_admin_auth + # Set attributes directly on the mock object + mock_team_row.team_id = test_team_id + mock_team_row.team_alias = "Test Team" + mock_team_row.team_member_permissions = permissions + mock_team_row.spend = 0.0 - response = client.get(f"/team/permissions_list?team_id={test_team_id}") + # Mock the get_team_object function used in the endpoint + with patch( + "litellm.proxy.management_endpoints.team_endpoints.get_team_object", + new_callable=AsyncMock, + return_value=mock_team_row, + ): + # Override the dependency for this test + app.dependency_overrides[user_api_key_auth] = lambda: mock_admin_auth - assert response.status_code == 200 - response_data = response.json() - assert response_data["team_id"] == test_team_id - assert ( - response_data["team_member_permissions"] - == mock_team_data["team_member_permissions"] - ) - assert ( - response_data["all_available_permissions"] - == TeamMemberPermissionChecks.get_all_available_team_member_permissions() - ) - mock_db_client.db.litellm_teamtable.find_unique.assert_awaited_once_with( - where={"team_id": test_team_id} - ) + response = client.get(f"/team/permissions_list?team_id={test_team_id}") - # Clean up dependency override - app.dependency_overrides = {} + assert response.status_code == 200 + response_data = response.json() + assert response_data["team_id"] == test_team_id + assert ( + response_data["team_member_permissions"] + == mock_team_data["team_member_permissions"] + ) + assert ( + response_data["all_available_permissions"] + == TeamMemberPermissionChecks.get_all_available_team_member_permissions() + ) + + # Clean up dependency override + app.dependency_overrides = {} # Test for /team/permissions_update endpoint (POST) @@ -102,15 +109,17 @@ async def test_update_team_permissions_success(mock_db_client, mock_admin_auth): Test successful update of team member permissions by an admin. """ test_team_id = "test-team-456" + update_permissions = ["/key/generate", "/key/update"] update_payload = { "team_id": test_team_id, - "team_member_permissions": ["/key/generate", "/key/update"], + "team_member_permissions": update_permissions, } + existing_permissions = ["/key/list"] mock_existing_team_data = { "team_id": test_team_id, "team_alias": "Existing Team", - "team_member_permissions": ["/key/list"], + "team_member_permissions": existing_permissions, "spend": 0.0, "models": [], } @@ -121,41 +130,50 @@ async def test_update_team_permissions_success(mock_db_client, mock_admin_auth): mock_existing_team_row = MagicMock(spec=LiteLLM_TeamTable) mock_existing_team_row.model_dump.return_value = mock_existing_team_data - # Set attributes directly if model_dump isn't enough for LiteLLM_TeamTable usage - for key, value in mock_existing_team_data.items(): - setattr(mock_existing_team_row, key, value) + + # Set attributes directly on the existing team mock + mock_existing_team_row.team_id = test_team_id + mock_existing_team_row.team_alias = "Existing Team" + mock_existing_team_row.team_member_permissions = existing_permissions + mock_existing_team_row.spend = 0.0 + mock_existing_team_row.models = [] mock_updated_team_row = MagicMock(spec=LiteLLM_TeamTable) mock_updated_team_row.model_dump.return_value = mock_updated_team_data - # Set attributes directly if model_dump isn't enough for LiteLLM_TeamTable usage - for key, value in mock_updated_team_data.items(): - setattr(mock_updated_team_row, key, value) - mock_db_client.db.litellm_teamtable.find_unique = AsyncMock( - return_value=mock_existing_team_row - ) - mock_db_client.db.litellm_teamtable.update = AsyncMock( - return_value=mock_updated_team_row - ) + # Set attributes directly on the updated team mock + mock_updated_team_row.team_id = test_team_id + mock_updated_team_row.team_alias = "Existing Team" + mock_updated_team_row.team_member_permissions = update_permissions + mock_updated_team_row.spend = 0.0 + mock_updated_team_row.models = [] - # Override the dependency for this test - app.dependency_overrides[user_api_key_auth] = lambda: mock_admin_auth + # Mock the get_team_object function used in the endpoint + with patch( + "litellm.proxy.management_endpoints.team_endpoints.get_team_object", + new_callable=AsyncMock, + return_value=mock_existing_team_row, + ): + # Mock the database update function + mock_db_client.db.litellm_teamtable.update = AsyncMock( + return_value=mock_updated_team_row + ) - response = client.post("/team/permissions_update", json=update_payload) + # Override the dependency for this test + app.dependency_overrides[user_api_key_auth] = lambda: mock_admin_auth - assert response.status_code == 200 - response_data = response.json() + response = client.post("/team/permissions_update", json=update_payload) - # Use model_dump for comparison if the endpoint returns the Prisma model directly - assert response_data == mock_updated_team_row.model_dump() + assert response.status_code == 200 + response_data = response.json() - mock_db_client.db.litellm_teamtable.find_unique.assert_awaited_once_with( - where={"team_id": test_team_id} - ) - mock_db_client.db.litellm_teamtable.update.assert_awaited_once_with( - where={"team_id": test_team_id}, - data={"team_member_permissions": update_payload["team_member_permissions"]}, - ) + # Use model_dump for comparison if the endpoint returns the Prisma model directly + assert response_data == mock_updated_team_row.model_dump() - # Clean up dependency override - app.dependency_overrides = {} + mock_db_client.db.litellm_teamtable.update.assert_awaited_once_with( + where={"team_id": test_team_id}, + data={"team_member_permissions": update_payload["team_member_permissions"]}, + ) + + # Clean up dependency override + app.dependency_overrides = {} diff --git a/ui/litellm-dashboard/src/components/edit_user.tsx b/ui/litellm-dashboard/src/components/edit_user.tsx index c47d44b25c..1cd47ffec2 100644 --- a/ui/litellm-dashboard/src/components/edit_user.tsx +++ b/ui/litellm-dashboard/src/components/edit_user.tsx @@ -18,6 +18,7 @@ import { Input, Select as Select2, message, + InputNumber, } from "antd"; import NumericalInput from "./shared/numerical_input"; @@ -113,7 +114,7 @@ const EditUserModal: React.FC = ({ visible, possibleUIRoles, tooltip="(float) - Spend of all LLM calls completed by this user" help="Across all keys (including keys with team_id)." > - + = ({ visible, possibleUIRoles, tooltip="(float) - Maximum budget of this user" help="Maximum budget of this user." > - +