diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index d83ceb1b095..5fb3355dc13 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -100,6 +100,8 @@ from litellm.types.proxy.management_endpoints.common_daily_activity import ( from litellm.types.proxy.management_endpoints.team_endpoints import ( BulkTeamMemberAddRequest, BulkTeamMemberAddResponse, + BulkUpdateTeamMemberPermissionsRequest, + BulkUpdateTeamMemberPermissionsResponse, GetTeamMemberPermissionsResponse, TeamListResponse, TeamMemberAddResult, @@ -4044,6 +4046,151 @@ async def update_team_member_permissions( return updated_team +@router.post( + "/team/permissions_bulk_update", + tags=["team management"], + dependencies=[Depends(user_api_key_auth)], + response_model=BulkUpdateTeamMemberPermissionsResponse, +) +@management_endpoint_wrapper +async def bulk_update_team_member_permissions( + data: BulkUpdateTeamMemberPermissionsRequest, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): + """ + Append permissions to existing teams. + + Either pass team_ids to target specific teams, or set + apply_to_all_teams=True to update every team. For each team, + the provided permissions are merged with the team's existing + permissions (duplicates are skipped). + """ + from litellm.proxy.proxy_server import prisma_client + + if prisma_client is None: + raise HTTPException(status_code=500, detail={"error": "No db connected"}) + + if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value: + raise HTTPException( + status_code=403, + detail={"error": "Only proxy admins can bulk-update team permissions"}, + ) + + if not data.permissions: + return { + "message": "No permissions provided", + "teams_updated": 0, + } + + if not data.apply_to_all_teams and not data.team_ids: + raise HTTPException( + status_code=400, + detail={ + "error": "Must provide team_ids or set apply_to_all_teams=true" + }, + ) + + if data.apply_to_all_teams and data.team_ids: + raise HTTPException( + status_code=400, + detail={ + "error": "Cannot set both apply_to_all_teams=true and team_ids" + }, + ) + + permissions_to_add = set(data.permissions) + + if data.team_ids: + teams_updated = await _append_permissions_to_specific_teams( + prisma_client, data.team_ids, permissions_to_add + ) + else: + teams_updated = await _append_permissions_to_all_teams( + prisma_client, permissions_to_add + ) + + return { + "message": "Team permissions updated successfully", + "teams_updated": teams_updated, + "permissions_appended": data.permissions, + } + + +async def _compute_and_batch_updates(prisma_client, teams, permissions_to_add: set) -> int: + """Compute merged permissions and batch-write updates. Returns count of teams updated.""" + updates = [] + for team in teams: + existing = set(team.team_member_permissions or []) + if permissions_to_add <= existing: + continue + merged = sorted(existing | permissions_to_add) # normalise to alphabetical order + updates.append((team.team_id, merged)) + + if updates: + batcher = prisma_client.db.batch_() + for team_id, merged_perms in updates: + batcher.litellm_teamtable.update( + where={"team_id": team_id}, + data={"team_member_permissions": merged_perms}, + ) + await batcher.commit() + + return len(updates) + + +async def _append_permissions_to_specific_teams( + prisma_client, team_ids: List[str], permissions_to_add: set +) -> int: + """Fetch specific teams by ID and append permissions.""" + teams = await prisma_client.db.litellm_teamtable.find_many( + where={"team_id": {"in": team_ids}}, + ) + + found_ids = {team.team_id for team in teams} + missing_ids = set(team_ids) - found_ids + if missing_ids: + raise HTTPException( + status_code=404, + detail={"error": f"Team(s) not found: {sorted(missing_ids)}"}, + ) + + return await _compute_and_batch_updates(prisma_client, teams, permissions_to_add) + + +async def _append_permissions_to_all_teams( + prisma_client, permissions_to_add: set +) -> int: + """Paginated read + batched write across all teams.""" + teams_updated = 0 + cursor = None + BATCH_SIZE = 500 + + while True: + find_args: dict = { + "take": BATCH_SIZE, + "order": {"team_id": "asc"}, + } + if cursor is not None: + find_args["cursor"] = {"team_id": cursor} + find_args["skip"] = 1 + + teams = await prisma_client.db.litellm_teamtable.find_many(**find_args) + + if not teams: + break + + teams_updated += await _compute_and_batch_updates( + prisma_client, teams, permissions_to_add + ) + + cursor = teams[-1].team_id + + if len(teams) < BATCH_SIZE: + break + + return teams_updated + + @router.get( "/team/daily/activity", response_model=SpendAnalyticsPaginatedResponse, diff --git a/litellm/types/proxy/management_endpoints/team_endpoints.py b/litellm/types/proxy/management_endpoints/team_endpoints.py index 77816fa78cc..e2ae20f4f47 100644 --- a/litellm/types/proxy/management_endpoints/team_endpoints.py +++ b/litellm/types/proxy/management_endpoints/team_endpoints.py @@ -3,6 +3,7 @@ from typing import Any, Dict, List, Optional, Union from pydantic import BaseModel from litellm.proxy._types import ( + KeyManagementRoutes, LiteLLM_DeletedTeamTable, LiteLLM_TeamMembership, LiteLLM_TeamTable, @@ -43,6 +44,38 @@ class UpdateTeamMemberPermissionsRequest(BaseModel): team_member_permissions: List[str] +class BulkUpdateTeamMemberPermissionsRequest(BaseModel): + """Request to bulk-update team member permissions across teams.""" + + permissions: List[KeyManagementRoutes] + """Permissions to append to the target teams (duplicates are skipped).""" + + team_ids: Optional[List[str]] = None + """Specific team IDs to update. Required unless apply_to_all_teams is True.""" + + apply_to_all_teams: bool = False + """When True, update all teams. Mutually exclusive with team_ids.""" + + +class BulkUpdateTeamMemberPermissionsResponse(BaseModel): + """Response for bulk team member permissions update.""" + + message: str + teams_updated: int + permissions_appended: Optional[List[str]] = None + + +class TeamListItem(LiteLLM_TeamTable): + """A team item in the paginated list response, enriched with computed fields.""" + + members_count: int = 0 + # Resources inherited from access groups (separate from direct assignments) + access_group_models: Optional[List[str]] = None + access_group_mcp_server_ids: Optional[List[str]] = None + access_group_agent_ids: Optional[List[str]] = None + + + class TeamListResponse(BaseModel): """Response to get the list of teams""" diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_default_params.py b/tests/test_litellm/proxy/management_endpoints/test_team_default_params.py index 7fc7cb8aae2..709ce9e6f71 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_default_params.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_default_params.py @@ -8,6 +8,7 @@ import sys from unittest.mock import AsyncMock, MagicMock, patch import pytest +from fastapi import HTTPException sys.path.insert( 0, os.path.abspath("../../../") @@ -485,3 +486,327 @@ class TestSafeDbOverrides: from litellm.constants import LITELLM_SETTINGS_SAFE_DB_OVERRIDES assert "default_internal_user_params" in LITELLM_SETTINGS_SAFE_DB_OVERRIDES + + +# --------------------------------------------------------------------------- +# POST /team/permissions/bulk_update +# --------------------------------------------------------------------------- + + +class TestBulkUpdateTeamMemberPermissions: + """Tests for the bulk_update_team_member_permissions endpoint.""" + + def _make_team(self, team_id: str, permissions: list): + """Create a mock team object.""" + team = MagicMock() + team.team_id = team_id + team.team_member_permissions = permissions + return team + + def _admin_key_dict(self): + from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + + return UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN.value, + api_key="sk-1234", + ) + + def _non_admin_key_dict(self): + from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + + return UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER.value, + api_key="sk-user", + ) + + # --- apply_to_all_teams tests --- + + @pytest.mark.asyncio + async def test_all_teams_appends_preserving_existing(self, monkeypatch): + """apply_to_all_teams: permissions are merged, not overwritten.""" + from litellm.proxy.management_endpoints.team_endpoints import ( + bulk_update_team_member_permissions, + ) + from litellm.types.proxy.management_endpoints.team_endpoints import ( + BulkUpdateTeamMemberPermissionsRequest, + ) + + team_a = self._make_team("team-a", ["/key/generate"]) + team_b = self._make_team("team-b", ["/key/delete", "/key/update"]) + + mock_batcher = MagicMock() + mock_batcher.commit = AsyncMock(return_value=None) + + mock_prisma = MagicMock() + mock_prisma.db.litellm_teamtable.find_many = AsyncMock(return_value=[team_a, team_b]) + mock_prisma.db.batch_ = MagicMock(return_value=mock_batcher) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) + + data = BulkUpdateTeamMemberPermissionsRequest( + permissions=["/team/daily/activity"], apply_to_all_teams=True + ) + result = await bulk_update_team_member_permissions(data=data, user_api_key_dict=self._admin_key_dict()) + + assert result["teams_updated"] == 2 + calls = mock_batcher.litellm_teamtable.update.call_args_list + assert len(calls) == 2 + + team_a_call = [c for c in calls if c.kwargs["where"]["team_id"] == "team-a"][0] + assert "/key/generate" in team_a_call.kwargs["data"]["team_member_permissions"] + assert "/team/daily/activity" in team_a_call.kwargs["data"]["team_member_permissions"] + + team_b_call = [c for c in calls if c.kwargs["where"]["team_id"] == "team-b"][0] + assert "/key/delete" in team_b_call.kwargs["data"]["team_member_permissions"] + assert "/key/update" in team_b_call.kwargs["data"]["team_member_permissions"] + + @pytest.mark.asyncio + async def test_all_teams_skips_teams_that_already_have_permission(self, monkeypatch): + """apply_to_all_teams: teams that already have the permission are skipped.""" + from litellm.proxy.management_endpoints.team_endpoints import ( + bulk_update_team_member_permissions, + ) + from litellm.types.proxy.management_endpoints.team_endpoints import ( + BulkUpdateTeamMemberPermissionsRequest, + ) + + team_has = self._make_team("team-has", ["/team/daily/activity", "/key/update"]) + team_missing = self._make_team("team-missing", ["/key/generate"]) + + mock_batcher = MagicMock() + mock_batcher.commit = AsyncMock(return_value=None) + + mock_prisma = MagicMock() + mock_prisma.db.litellm_teamtable.find_many = AsyncMock(return_value=[team_has, team_missing]) + mock_prisma.db.batch_ = MagicMock(return_value=mock_batcher) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) + + data = BulkUpdateTeamMemberPermissionsRequest( + permissions=["/team/daily/activity"], apply_to_all_teams=True + ) + result = await bulk_update_team_member_permissions(data=data, user_api_key_dict=self._admin_key_dict()) + + assert result["teams_updated"] == 1 + calls = mock_batcher.litellm_teamtable.update.call_args_list + assert len(calls) == 1 + assert calls[0].kwargs["where"]["team_id"] == "team-missing" + + @pytest.mark.asyncio + async def test_all_teams_pagination(self, monkeypatch): + """apply_to_all_teams: cursor-based pagination processes multiple pages.""" + from litellm.proxy.management_endpoints.team_endpoints import ( + bulk_update_team_member_permissions, + ) + from litellm.types.proxy.management_endpoints.team_endpoints import ( + BulkUpdateTeamMemberPermissionsRequest, + ) + + page1 = [self._make_team(f"team-{i}", []) for i in range(500)] + page2 = [self._make_team(f"team-{i}", []) for i in range(500, 502)] + + mock_batcher = MagicMock() + mock_batcher.commit = AsyncMock(return_value=None) + + mock_prisma = MagicMock() + mock_prisma.db.litellm_teamtable.find_many = AsyncMock(side_effect=[page1, page2]) + mock_prisma.db.batch_ = MagicMock(return_value=mock_batcher) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) + + data = BulkUpdateTeamMemberPermissionsRequest( + permissions=["/team/daily/activity"], apply_to_all_teams=True + ) + result = await bulk_update_team_member_permissions(data=data, user_api_key_dict=self._admin_key_dict()) + + assert result["teams_updated"] == 502 + find_calls = mock_prisma.db.litellm_teamtable.find_many.call_args_list + assert len(find_calls) == 2 + assert find_calls[1].kwargs["cursor"] == {"team_id": "team-499"} + assert mock_batcher.commit.call_count == 2 + + # --- team_ids tests --- + + @pytest.mark.asyncio + async def test_team_ids_updates_only_specified_teams(self, monkeypatch): + """team_ids: only the specified teams are fetched and updated.""" + from litellm.proxy.management_endpoints.team_endpoints import ( + bulk_update_team_member_permissions, + ) + from litellm.types.proxy.management_endpoints.team_endpoints import ( + BulkUpdateTeamMemberPermissionsRequest, + ) + + team_a = self._make_team("team-a", ["/key/generate"]) + team_b = self._make_team("team-b", ["/key/delete"]) + + mock_batcher = MagicMock() + mock_batcher.commit = AsyncMock(return_value=None) + + mock_prisma = MagicMock() + mock_prisma.db.litellm_teamtable.find_many = AsyncMock(return_value=[team_a, team_b]) + mock_prisma.db.batch_ = MagicMock(return_value=mock_batcher) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) + + data = BulkUpdateTeamMemberPermissionsRequest( + permissions=["/team/daily/activity"], team_ids=["team-a", "team-b"] + ) + result = await bulk_update_team_member_permissions(data=data, user_api_key_dict=self._admin_key_dict()) + + assert result["teams_updated"] == 2 + + # Verify find_many was called with the team_ids filter + find_call = mock_prisma.db.litellm_teamtable.find_many.call_args + assert find_call.kwargs["where"] == {"team_id": {"in": ["team-a", "team-b"]}} + + @pytest.mark.asyncio + async def test_team_ids_skips_teams_that_already_have_permission(self, monkeypatch): + """team_ids: teams that already have the permission are skipped.""" + from litellm.proxy.management_endpoints.team_endpoints import ( + bulk_update_team_member_permissions, + ) + from litellm.types.proxy.management_endpoints.team_endpoints import ( + BulkUpdateTeamMemberPermissionsRequest, + ) + + team_has = self._make_team("team-has", ["/team/daily/activity"]) + team_missing = self._make_team("team-missing", []) + + mock_batcher = MagicMock() + mock_batcher.commit = AsyncMock(return_value=None) + + mock_prisma = MagicMock() + mock_prisma.db.litellm_teamtable.find_many = AsyncMock(return_value=[team_has, team_missing]) + mock_prisma.db.batch_ = MagicMock(return_value=mock_batcher) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) + + data = BulkUpdateTeamMemberPermissionsRequest( + permissions=["/team/daily/activity"], team_ids=["team-has", "team-missing"] + ) + result = await bulk_update_team_member_permissions(data=data, user_api_key_dict=self._admin_key_dict()) + + assert result["teams_updated"] == 1 + calls = mock_batcher.litellm_teamtable.update.call_args_list + assert calls[0].kwargs["where"]["team_id"] == "team-missing" + + @pytest.mark.asyncio + async def test_team_ids_returns_404_for_missing_teams(self, monkeypatch): + """If any provided team_ids don't exist, return 404.""" + from litellm.proxy.management_endpoints.team_endpoints import ( + bulk_update_team_member_permissions, + ) + from litellm.types.proxy.management_endpoints.team_endpoints import ( + BulkUpdateTeamMemberPermissionsRequest, + ) + + team_a = self._make_team("team-a", ["/key/generate"]) + + mock_prisma = MagicMock() + # Only team-a exists, team-b does not + mock_prisma.db.litellm_teamtable.find_many = AsyncMock(return_value=[team_a]) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) + + data = BulkUpdateTeamMemberPermissionsRequest( + permissions=["/team/daily/activity"], team_ids=["team-a", "team-b"] + ) + + with pytest.raises(HTTPException) as exc_info: + await bulk_update_team_member_permissions(data=data, user_api_key_dict=self._admin_key_dict()) + + assert exc_info.value.status_code == 404 + assert "team-b" in str(exc_info.value.detail) + + # --- validation tests --- + + @pytest.mark.asyncio + async def test_rejects_when_no_team_ids_and_no_apply_all(self, monkeypatch): + """Must provide team_ids or set apply_to_all_teams=True.""" + from litellm.proxy.management_endpoints.team_endpoints import ( + bulk_update_team_member_permissions, + ) + from litellm.types.proxy.management_endpoints.team_endpoints import ( + BulkUpdateTeamMemberPermissionsRequest, + ) + + mock_prisma = MagicMock() + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) + + data = BulkUpdateTeamMemberPermissionsRequest(permissions=["/team/daily/activity"]) + + with pytest.raises(HTTPException) as exc_info: + await bulk_update_team_member_permissions(data=data, user_api_key_dict=self._admin_key_dict()) + + assert exc_info.value.status_code == 400 + + @pytest.mark.asyncio + async def test_rejects_when_both_team_ids_and_apply_all(self, monkeypatch): + """Cannot set both team_ids and apply_to_all_teams.""" + from litellm.proxy.management_endpoints.team_endpoints import ( + bulk_update_team_member_permissions, + ) + from litellm.types.proxy.management_endpoints.team_endpoints import ( + BulkUpdateTeamMemberPermissionsRequest, + ) + + mock_prisma = MagicMock() + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) + + data = BulkUpdateTeamMemberPermissionsRequest( + permissions=["/team/daily/activity"], + team_ids=["team-a"], + apply_to_all_teams=True, + ) + + with pytest.raises(HTTPException) as exc_info: + await bulk_update_team_member_permissions(data=data, user_api_key_dict=self._admin_key_dict()) + + assert exc_info.value.status_code == 400 + + @pytest.mark.asyncio + async def test_empty_permissions_list_is_noop(self, monkeypatch): + """Passing an empty permissions list returns immediately with 0 updated.""" + from litellm.proxy.management_endpoints.team_endpoints import ( + bulk_update_team_member_permissions, + ) + from litellm.types.proxy.management_endpoints.team_endpoints import ( + BulkUpdateTeamMemberPermissionsRequest, + ) + + mock_prisma = MagicMock() + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) + + data = BulkUpdateTeamMemberPermissionsRequest(permissions=[]) + result = await bulk_update_team_member_permissions(data=data, user_api_key_dict=self._admin_key_dict()) + + assert result["teams_updated"] == 0 + mock_prisma.db.litellm_teamtable.find_many.assert_not_called() + + @pytest.mark.asyncio + async def test_non_admin_gets_403(self, monkeypatch): + """Non-admin users are rejected with 403.""" + from litellm.proxy.management_endpoints.team_endpoints import ( + bulk_update_team_member_permissions, + ) + from litellm.types.proxy.management_endpoints.team_endpoints import ( + BulkUpdateTeamMemberPermissionsRequest, + ) + + mock_prisma = MagicMock() + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) + + data = BulkUpdateTeamMemberPermissionsRequest( + permissions=["/team/daily/activity"], apply_to_all_teams=True + ) + + with pytest.raises(HTTPException) as exc_info: + await bulk_update_team_member_permissions(data=data, user_api_key_dict=self._non_admin_key_dict()) + + assert exc_info.value.status_code == 403 + + def test_invalid_permission_rejected_by_pydantic(self): + """Invalid permission strings are rejected at the type level by Pydantic.""" + from pydantic import ValidationError + + from litellm.types.proxy.management_endpoints.team_endpoints import ( + BulkUpdateTeamMemberPermissionsRequest, + ) + + with pytest.raises(ValidationError): + BulkUpdateTeamMemberPermissionsRequest(permissions=["/not/a/real/permission"]) diff --git a/tests/test_litellm/proxy/prompts/test_prompt_endpoints_crud.py b/tests/test_litellm/proxy/prompts/test_prompt_endpoints_crud.py index 2c5bc1bf87d..0c09be99aeb 100644 --- a/tests/test_litellm/proxy/prompts/test_prompt_endpoints_crud.py +++ b/tests/test_litellm/proxy/prompts/test_prompt_endpoints_crud.py @@ -148,7 +148,10 @@ async def test_get_prompt_info_by_base_id(): ) # Mock In-Memory Registry + # Patch prisma_client to None to avoid leaking state from other tests with patch( + "litellm.proxy.proxy_server.prisma_client", None + ), patch( "litellm.proxy.prompts.prompt_registry.IN_MEMORY_PROMPT_REGISTRY" ) as mock_registry: # Setup mocks behavior