diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 0d34974fbe..8a8e703831 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -3978,11 +3978,13 @@ async def _batch_resolve_access_group_resources( def _convert_teams_to_response_models( teams: list, use_deleted_table: bool, + keys_count_by_team: Optional[Dict[str, int]] = None, ) -> List[Union[TeamListItem, LiteLLM_TeamTable, LiteLLM_DeletedTeamTable]]: """Convert raw Prisma team rows to response models.""" team_list: List[ Union[TeamListItem, LiteLLM_TeamTable, LiteLLM_DeletedTeamTable] ] = [] + counts = keys_count_by_team or {} for team in teams: try: team_dict = team.model_dump() @@ -3997,10 +3999,45 @@ def _convert_teams_to_response_models( members_with_roles = [] team_dict["members_with_roles"] = members_with_roles members_count = len(members_with_roles) - team_list.append(TeamListItem(**team_dict, members_count=members_count)) + keys_count = counts.get(team_dict.get("team_id") or "", 0) + team_list.append( + TeamListItem( + **team_dict, + members_count=members_count, + keys_count=keys_count, + ) + ) return team_list +async def _get_keys_count_by_team( + prisma_client: Any, + teams: list, +) -> Dict[str, int]: + """Aggregate virtual-key counts per team for the given page of teams. + + Runs a single GROUP BY against LiteLLM_VerificationToken. The IN clause is + bounded by page_size and uses the existing @@index([team_id]), so this is + one DB round-trip per page. Returns an empty map when the page has no teams. + """ + page_team_ids = [ + getattr(t, "team_id", None) for t in teams if getattr(t, "team_id", None) + ] + if not page_team_ids: + return {} + + grouped = await prisma_client.db.litellm_verificationtoken.group_by( + by=["team_id"], + where={"team_id": {"in": page_team_ids}}, + count={"team_id": True}, + ) + return { + row["team_id"]: row.get("_count", {}).get("team_id", 0) + for row in grouped + if row.get("team_id") + } + + async def _enforce_list_team_v2_access( user_api_key_dict: UserAPIKeyAuth, user_id: Optional[str], @@ -4228,8 +4265,16 @@ async def list_team_v2( # Calculate total pages total_pages = -(-total_count // page_size) # Ceiling division - # Convert Prisma models to response models with members_count - team_list = _convert_teams_to_response_models(teams, use_deleted_table) + # Aggregate virtual-key counts per team for the current page. The deleted + # table does not carry keys_count, so it is skipped. + keys_count_by_team: Dict[str, int] = {} + if not use_deleted_table: + keys_count_by_team = await _get_keys_count_by_team(prisma_client, teams) + + # Convert Prisma models to response models with members_count and keys_count + team_list = _convert_teams_to_response_models( + teams, use_deleted_table, keys_count_by_team=keys_count_by_team + ) # Resolve resources inherited from access groups (single batch query) if not use_deleted_table: diff --git a/litellm/types/proxy/management_endpoints/team_endpoints.py b/litellm/types/proxy/management_endpoints/team_endpoints.py index cb27fd5230..0e55553587 100644 --- a/litellm/types/proxy/management_endpoints/team_endpoints.py +++ b/litellm/types/proxy/management_endpoints/team_endpoints.py @@ -69,6 +69,7 @@ class TeamListItem(LiteLLM_TeamTable): """A team item in the paginated list response, enriched with computed fields.""" members_count: int = 0 + keys_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 diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py index 13bb39c35c..d580f1f770 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -2832,6 +2832,7 @@ async def test_list_team_v2_security_check_non_admin_user_own_teams(): ] mock_db.litellm_teamtable.find_many = AsyncMock(return_value=mock_teams) mock_db.litellm_teamtable.count = AsyncMock(return_value=2) + mock_db.litellm_verificationtoken.group_by = AsyncMock(return_value=[]) with patch( "litellm.proxy.management_endpoints.team_endpoints.get_user_object", @@ -2888,6 +2889,7 @@ async def test_list_team_v2_security_check_admin_user(): ] mock_db.litellm_teamtable.find_many = AsyncMock(return_value=mock_teams) mock_db.litellm_teamtable.count = AsyncMock(return_value=2) + mock_db.litellm_verificationtoken.group_by = AsyncMock(return_value=[]) # Should NOT raise an exception result = await list_team_v2( @@ -3036,6 +3038,7 @@ async def test_list_team_v2_org_admin_sees_org_teams(): } mock_db.litellm_teamtable.find_many = AsyncMock(return_value=[mock_team]) mock_db.litellm_teamtable.count = AsyncMock(return_value=1) + mock_db.litellm_verificationtoken.group_by = AsyncMock(return_value=[]) result = await list_team_v2( http_request=mock_request, @@ -3211,6 +3214,7 @@ async def test_list_team_v2_org_admin_with_user_id_returns_user_teams(): } mock_db.litellm_teamtable.find_many = AsyncMock(return_value=[mock_team]) mock_db.litellm_teamtable.count = AsyncMock(return_value=1) + mock_db.litellm_verificationtoken.group_by = AsyncMock(return_value=[]) result = await list_team_v2( http_request=mock_request, @@ -3390,6 +3394,163 @@ async def test_list_team_v2_search_composes_with_user_id_filter(): assert where["team_id"] == {"in": ["team_a", "team_b"]} +@pytest.mark.asyncio +async def test_list_team_v2_populates_keys_count(): + """ + Test that list_team_v2 returns a keys_count per team derived from a single + batched group_by against LiteLLM_VerificationToken. + """ + from unittest.mock import AsyncMock, Mock, patch + + from fastapi import Request + + from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + from litellm.proxy.management_endpoints.team_endpoints import list_team_v2 + + mock_request = Mock(spec=Request) + mock_user_api_key_dict_admin = UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, + user_id="admin_user_123", + ) + + with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma_client: + mock_db = Mock() + mock_prisma_client.db = mock_db + + team_a = Mock() + team_a.team_id = "team_a" + team_a.model_dump = lambda: { + "team_id": "team_a", + "team_alias": "Team A", + "members_with_roles": [{"user_id": "u1", "role": "user"}], + } + team_b = Mock() + team_b.team_id = "team_b" + team_b.model_dump = lambda: { + "team_id": "team_b", + "team_alias": "Team B", + "members_with_roles": [], + } + + mock_db.litellm_teamtable.find_many = AsyncMock(return_value=[team_a, team_b]) + mock_db.litellm_teamtable.count = AsyncMock(return_value=2) + mock_db.litellm_verificationtoken.group_by = AsyncMock( + return_value=[ + {"team_id": "team_a", "_count": {"team_id": 3}}, + # team_b intentionally absent → expect 0 + ] + ) + + result = await list_team_v2( + http_request=mock_request, + user_id=None, + user_api_key_dict=mock_user_api_key_dict_admin, + page=1, + page_size=10, + status=None, + ) + + assert result["total"] == 2 + by_id = {t.team_id: t for t in result["teams"]} + assert by_id["team_a"].keys_count == 3 + assert by_id["team_b"].keys_count == 0 + + # The aggregate is one batched query, filtered by the page's team IDs. + group_by_kwargs = mock_db.litellm_verificationtoken.group_by.call_args.kwargs + assert group_by_kwargs["by"] == ["team_id"] + assert group_by_kwargs["where"] == {"team_id": {"in": ["team_a", "team_b"]}} + assert group_by_kwargs["count"] == {"team_id": True} + + +@pytest.mark.asyncio +async def test_list_team_v2_keys_count_skipped_for_empty_page(): + """ + When the page has no teams, the keys-count group_by must not be issued. + """ + from unittest.mock import AsyncMock, Mock, patch + + from fastapi import Request + + from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + from litellm.proxy.management_endpoints.team_endpoints import list_team_v2 + + mock_request = Mock(spec=Request) + mock_user_api_key_dict_admin = UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, + user_id="admin_user_123", + ) + + with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma_client: + mock_db = Mock() + mock_prisma_client.db = mock_db + + mock_db.litellm_teamtable.find_many = AsyncMock(return_value=[]) + mock_db.litellm_teamtable.count = AsyncMock(return_value=0) + mock_db.litellm_verificationtoken.group_by = AsyncMock(return_value=[]) + + result = await list_team_v2( + http_request=mock_request, + user_id=None, + user_api_key_dict=mock_user_api_key_dict_admin, + page=1, + page_size=10, + status=None, + ) + + assert result["total"] == 0 + assert result["teams"] == [] + mock_db.litellm_verificationtoken.group_by.assert_not_called() + + +@pytest.mark.asyncio +async def test_list_team_v2_keys_count_skipped_for_deleted_status(): + """ + The deleted-table branch returns LiteLLM_DeletedTeamTable items, which do + not carry keys_count — group_by must not be issued. + """ + from unittest.mock import AsyncMock, Mock, patch + + from fastapi import Request + + from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + from litellm.proxy.management_endpoints.team_endpoints import list_team_v2 + + mock_request = Mock(spec=Request) + mock_user_api_key_dict_admin = UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, + user_id="admin_user_123", + ) + + with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma_client: + mock_db = Mock() + mock_prisma_client.db = mock_db + + mock_deleted = Mock() + mock_deleted.team_id = "team_d" + mock_deleted.model_dump = lambda: { + "team_id": "team_d", + "team_alias": "Deleted Team", + } + + mock_db.litellm_deletedteamtable.find_many = AsyncMock( + return_value=[mock_deleted] + ) + mock_db.litellm_deletedteamtable.count = AsyncMock(return_value=1) + mock_db.litellm_verificationtoken.group_by = AsyncMock(return_value=[]) + + result = await list_team_v2( + http_request=mock_request, + user_id=None, + user_api_key_dict=mock_user_api_key_dict_admin, + page=1, + page_size=10, + status="deleted", + ) + + assert result["total"] == 1 + mock_db.litellm_verificationtoken.group_by.assert_not_called() + + @pytest.mark.asyncio async def test_team_member_delete_cleans_membership(mock_db_client, mock_admin_auth): """ diff --git a/ui/litellm-dashboard/src/components/OldTeams.test.tsx b/ui/litellm-dashboard/src/components/OldTeams.test.tsx index 4b89820bad..b8707c1a33 100644 --- a/ui/litellm-dashboard/src/components/OldTeams.test.tsx +++ b/ui/litellm-dashboard/src/components/OldTeams.test.tsx @@ -1018,3 +1018,83 @@ describe("OldTeams - organization alias display", () => { }); }); }); + +describe("OldTeams - Resources column keys badge", () => { + beforeEach(() => { + vi.clearAllMocks(); + mockUseOrganizations.mockReturnValue({ data: [] }); + }); + + it("renders keys_count from the v2 payload in the Resources badge", async () => { + const { container } = renderWithQueryClient( + , + ); + + await waitFor(() => { + expect(screen.getByText("Team With Keys")).toBeInTheDocument(); + }); + const cyanTag = container.querySelector(".ant-tag-cyan"); + expect(cyanTag).not.toBeNull(); + expect(cyanTag?.textContent).toContain("3"); + }); + + it("falls back to keys.length when keys_count is absent", async () => { + const { container } = renderWithQueryClient( + , + ); + + await waitFor(() => { + expect(screen.getByText("Legacy Team")).toBeInTheDocument(); + }); + const cyanTag = container.querySelector(".ant-tag-cyan"); + expect(cyanTag).not.toBeNull(); + expect(cyanTag?.textContent).toContain("2"); + }); +}); diff --git a/ui/litellm-dashboard/src/components/OldTeams.tsx b/ui/litellm-dashboard/src/components/OldTeams.tsx index da00ad911b..8f9e5a75c2 100644 --- a/ui/litellm-dashboard/src/components/OldTeams.tsx +++ b/ui/litellm-dashboard/src/components/OldTeams.tsx @@ -105,6 +105,7 @@ interface TeamInfo { interface PerTeamInfo { keys: KeyResponse[]; + keys_count: number; team_info: TeamInfo; } @@ -364,6 +365,7 @@ const Teams: React.FC = ({ (acc, team) => { acc[team.team_id] = { keys: team.keys || [], + keys_count: team.keys_count ?? team.keys?.length ?? 0, team_info: { members_with_roles: team.members_with_roles || [], }, @@ -745,7 +747,7 @@ const Teams: React.FC = ({ render: (_: unknown, record: Team) => { const memberCount = perTeamInfo?.[record.team_id]?.team_info?.members_with_roles?.length ?? 0; const modelCount = record.models?.length ?? 0; - const keyCount = perTeamInfo?.[record.team_id]?.keys?.length ?? 0; + const keyCount = perTeamInfo?.[record.team_id]?.keys_count ?? 0; return ( @@ -977,17 +979,23 @@ const Teams: React.FC = ({ { + const deleteKeyCount = + teamToDelete?.keys_count ?? teamToDelete?.keys?.length ?? 0; + return deleteKeyCount === 0 ? undefined - : `Warning: This team has ${teamToDelete?.keys?.length} keys associated with it. Deleting the team will also delete all associated keys. This action is irreversible.` - } + : `Warning: This team has ${deleteKeyCount} keys associated with it. Deleting the team will also delete all associated keys. This action is irreversible.`; + })()} message="Are you sure you want to delete this team and all its keys? This action cannot be undone." resourceInformationTitle="Team Information" resourceInformation={[ { label: "Team ID", value: teamToDelete?.team_id, code: true }, { label: "Team Name", value: teamToDelete?.team_alias }, - { label: "Keys", value: teamToDelete?.keys?.length }, + { + label: "Keys", + value: + teamToDelete?.keys_count ?? teamToDelete?.keys?.length ?? 0, + }, { label: "Members", value: teamToDelete?.members_with_roles?.length }, ]} requiredConfirmation={teamToDelete?.team_alias} diff --git a/ui/litellm-dashboard/src/components/key_team_helpers/key_list.tsx b/ui/litellm-dashboard/src/components/key_team_helpers/key_list.tsx index 6b3c65aaf7..60568da48e 100644 --- a/ui/litellm-dashboard/src/components/key_team_helpers/key_list.tsx +++ b/ui/litellm-dashboard/src/components/key_team_helpers/key_list.tsx @@ -13,6 +13,7 @@ export interface Team { organization_id: string; created_at: string; keys: KeyResponse[]; + keys_count?: number; members_with_roles: Member[]; spend: number; access_group_ids?: string[];