From 66a396123985657b536e99a96833ee2823e9ef94 Mon Sep 17 00:00:00 2001 From: jesus Date: Mon, 12 Oct 2026 21:42:40 +0000 Subject: [PATCH] refactor(teams): address review on caller membership enrichment Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../management_endpoints/team_endpoints.py | 112 +++++++++--------- .../test_team_endpoints.py | 106 ++++++++--------- .../TeamMemberSpendBudgetCell.test.tsx | 10 +- .../TeamsPage/TeamMemberSpendBudgetCell.tsx | 9 +- ...st.tsx => TeamsTable.integration.test.tsx} | 0 .../components/key_team_helpers/key_list.tsx | 7 +- 6 files changed, 115 insertions(+), 129 deletions(-) rename ui/litellm-dashboard/src/components/TeamsPage/{TeamsTable.test.tsx => TeamsTable.integration.test.tsx} (100%) diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 3da231b1551..0cdc5c60ef2 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -5227,16 +5227,15 @@ async def team_member_me( user_api_key_cache=user_api_key_cache, ) - caller_user_email: Final = user_api_key_dict.user_email - member_role: str | None = None - for m in team_table.members_with_roles: - # Match by user_id when present, else fall back to email — members - # added by email may have user_id=None on the stored entry. - if (m.user_id is not None and m.user_id == caller_user_id) or ( - m.user_email is not None and caller_user_email is not None and m.user_email == caller_user_email - ): - member_role = m.role - break + caller_identity: Final = _TeamListCallerIdentity( + user_id=caller_user_id, + user_email=user_api_key_dict.user_email, + ) + matched_member: Final = next( + (member for member in team_table.members_with_roles if _team_member_matches_caller(member, caller_identity)), + None, + ) + member_role: Final = matched_member.role if matched_member is not None else None if member_role is None: # Caller is not a member of this team. Even proxy admins get 404 here — @@ -5695,22 +5694,31 @@ def _team_list_budget_model(budget_row: _TeamListBudgetRow | None) -> LiteLLM_Bu ) -def _caller_in_roster(team: TeamListItem, caller_user_id: str, caller_user_email: str | None) -> bool: - return any( - (member.user_id is not None and member.user_id == caller_user_id) - or (member.user_email is not None and caller_user_email is not None and member.user_email == caller_user_email) - for member in team.members_with_roles +@dataclass(frozen=True, slots=True) +class _TeamListCallerIdentity: + user_id: str | None + user_email: str | None + + +def _team_member_matches_caller(member: Member, caller_identity: _TeamListCallerIdentity) -> bool: + return (member.user_id is not None and member.user_id == caller_identity.user_id) or ( + member.user_email is not None + and caller_identity.user_email is not None + and member.user_email == caller_identity.user_email ) +def _caller_in_roster(team: TeamListItem, caller_identity: _TeamListCallerIdentity) -> bool: + return any(_team_member_matches_caller(member, caller_identity) for member in team.members_with_roles) + + def _team_member_default_budget_id( team: TeamListItem, - caller_user_id: str, - caller_user_email: str | None, + caller_identity: _TeamListCallerIdentity, memberships_by_team: Mapping[str, _TeamListMembershipRow], ) -> str | None: membership_row: Final = memberships_by_team.get(team.team_id) - if membership_row is None and not _caller_in_roster(team, caller_user_id, caller_user_email): + if membership_row is None and not _caller_in_roster(team, caller_identity): return None member_budget: Final = _team_list_budget_model( membership_row.litellm_budget_table if membership_row is not None else None @@ -5728,20 +5736,17 @@ def _team_member_default_budget_id( def _team_list_caller_membership( team: TeamListItem, - caller_user_id: str, - caller_user_email: str | None, + caller_identity: _TeamListCallerIdentity, memberships_by_team: Mapping[str, _TeamListMembershipRow], + default_budget_id: str | None, budgets_by_id: Mapping[str, _TeamListBudgetRow], now: datetime, ) -> TeamListCallerMembership | None: membership_row: Final = memberships_by_team.get(team.team_id) - if membership_row is None and not _caller_in_roster(team, caller_user_id, caller_user_email): + if membership_row is None and not _caller_in_roster(team, caller_identity): return None member_budget_row: Final = membership_row.litellm_budget_table if membership_row is not None else None member_budget: Final = _team_list_budget_model(member_budget_row) - default_budget_id: Final = _team_member_default_budget_id( - team, caller_user_id, caller_user_email, memberships_by_team - ) default_budget_row: Final = budgets_by_id.get(default_budget_id) if default_budget_id is not None else None team_default_budget: Final = _team_list_budget_model(default_budget_row) max_budget: Final = resolve_effective_team_member_budget( @@ -5762,10 +5767,10 @@ def _team_list_caller_membership( async def _enrich_team_list_with_caller_membership( prisma_client: PrismaClient, teams: Sequence[TeamListItem], - caller_user_id: str | None, - caller_user_email: str | None, + caller_identity: _TeamListCallerIdentity, now: datetime, ) -> tuple[TeamListItem, ...]: + caller_user_id: Final = caller_identity.user_id if caller_user_id is None: return tuple(team.model_copy(update={"caller_membership": None}) for team in teams) if not teams: @@ -5778,17 +5783,11 @@ async def _enrich_team_list_with_caller_membership( include={"litellm_budget_table": True}, ) memberships_by_team: Final = {row.team_id: row for row in membership_rows} + default_budget_id_by_team: Final = MappingProxyType( + {team.team_id: _team_member_default_budget_id(team, caller_identity, memberships_by_team) for team in teams} + ) default_budget_ids: Final = tuple( - dict.fromkeys( - budget_id - for team in teams - if ( - budget_id := _team_member_default_budget_id( - team, caller_user_id, caller_user_email, memberships_by_team - ) - ) - is not None - ) + dict.fromkeys(budget_id for budget_id in default_budget_id_by_team.values() if budget_id is not None) ) default_budget_rows: Final[Sequence[_TeamListBudgetRow]] = ( await _budget_db(prisma_client).find_many( @@ -5804,9 +5803,9 @@ async def _enrich_team_list_with_caller_membership( update={ "caller_membership": _team_list_caller_membership( team=team, - caller_user_id=caller_user_id, - caller_user_email=caller_user_email, + caller_identity=caller_identity, memberships_by_team=memberships_by_team, + default_budget_id=default_budget_id_by_team.get(team.team_id), budgets_by_id=budgets_by_id, now=now, ) @@ -5929,7 +5928,7 @@ async def list_team_v2( sort_order: str = fastapi.Query(default="asc", description="Sort order ('asc' or 'desc')"), status: str | None = fastapi.Query(default=None, description="Filter by status (e.g. 'deleted')"), user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), -) -> Mapping[str, int | Sequence[TeamListItem | LiteLLM_TeamTable | LiteLLM_DeletedTeamTable]]: +) -> TeamListResponse: """ Get a paginated list of teams with filtering and sorting options. @@ -6004,13 +6003,13 @@ async def list_team_v2( ) if where_conditions is None: - return { - "teams": [], - "total": 0, - "page": page, - "page_size": page_size, - "total_pages": 0, - } + return TeamListResponse( + teams=[], + total=0, + page=page, + page_size=page_size, + total_pages=0, + ) # Build order_by conditions valid_sort_columns: Final = ["team_id", "team_alias", "created_at"] @@ -6076,24 +6075,27 @@ async def list_team_v2( {a for group in team_groups for a in (group.access_agent_ids or [])} ) + caller_identity: Final = _TeamListCallerIdentity( + user_id=user_api_key_dict.user_id, + user_email=user_api_key_dict.user_email, + ) team_list_with_caller_membership: Final = ( await _enrich_team_list_with_caller_membership( prisma_client=prisma_client, teams=tuple(team for team in team_list if isinstance(team, TeamListItem)), - caller_user_id=user_api_key_dict.user_id, - caller_user_email=user_api_key_dict.user_email, + caller_identity=caller_identity, now=get_utc_datetime(), ) if not use_deleted_table and team_list else team_list ) - return { - "teams": team_list_with_caller_membership, - "total": total_count, - "page": page, - "page_size": page_size, - "total_pages": total_pages, - } + return TeamListResponse( + teams=list(team_list_with_caller_membership), + total=total_count, + page=page, + page_size=page_size, + total_pages=total_pages, + ) async def _authorize_and_filter_teams( diff --git a/tests/unit/proxy/management_endpoints/test_team_endpoints.py b/tests/unit/proxy/management_endpoints/test_team_endpoints.py index 330f8d8310e..51366948a60 100644 --- a/tests/unit/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_team_endpoints.py @@ -50,6 +50,7 @@ from litellm.proxy.management_endpoints.team_admin_field_permissions import ( from litellm.proxy.management_endpoints.team_endpoints import ( _STRIP_DELETED_TEAM_FROM_USERS_SQL, _enrich_team_list_with_caller_membership, + _TeamListCallerIdentity, GetTeamMemberPermissionsResponse, UpdateTeamMemberPermissionsRequest, _build_team_list_where_conditions, @@ -3851,9 +3852,8 @@ async def test_list_team_v2_security_check_non_admin_user_own_teams(): ) # Should return results without error - assert "teams" in result - assert "total" in result - assert result["total"] == 2 + assert len(result.teams) == 2 + assert result.total == 2 @pytest.mark.asyncio @@ -3903,9 +3903,8 @@ async def test_list_team_v2_security_check_admin_user(): ) # Should return results without error - assert "teams" in result - assert "total" in result - assert result["total"] == 2 + assert len(result.teams) == 2 + assert result.total == 2 @pytest.mark.asyncio @@ -3971,10 +3970,8 @@ async def test_list_team_v2_with_status_deleted(): mock_db.litellm_teamtable.count.assert_not_called() # Should return results without error - assert "teams" in result - assert "total" in result - assert result["total"] == 2 - assert len(result["teams"]) == 2 + assert len(result.teams) == 2 + assert result.total == 2 @pytest.mark.asyncio @@ -4045,8 +4042,8 @@ async def test_list_team_v2_includes_litellm_model_table(): status=None, ) - assert result["teams"][0].litellm_model_table is not None - assert result["teams"][0].litellm_model_table.model_aliases == {"my-fast-model": "fake-model"} + assert result.teams[0].litellm_model_table is not None + assert result.teams[0].litellm_model_table.model_aliases == {"my-fast-model": "fake-model"} mock_db.litellm_deletedteamtable.find_many = AsyncMock( side_effect=lambda **kw: [_team_row("team_2", kw.get("include"))] @@ -4145,9 +4142,9 @@ async def test_list_team_v2_org_admin_sees_org_teams(): status=None, ) - assert result["total"] == 1 - assert len(result["teams"]) == 1 - assert result["teams"][0].members_count == 1 + assert result.total == 1 + assert len(result.teams) == 1 + assert result.teams[0].members_count == 1 # Verify org-scoped where clause where = mock_db.litellm_teamtable.find_many.call_args.kwargs["where"] @@ -4247,8 +4244,8 @@ async def test_list_team_v2_org_admin_own_user_id_sees_all_org_teams(): status=None, ) - assert result["total"] == 2 - assert len(result["teams"]) == 2 + assert result.total == 2 + assert len(result.teams) == 2 # Verify the where clause scopes by org OR own membership — no # top-level team_id filter that would hide org teams they aren't in @@ -4386,8 +4383,8 @@ async def test_list_team_v2_org_admin_own_query_keeps_memberships_in_other_orgs( sort_order="asc", status=None, ) - assert result["total"] == len(result["teams"]) - return [t.team_id for t in result["teams"]] + assert result.total == len(result.teams) + return [t.team_id for t in result.teams] own_view = ["team_in_org_A", "other_team_in_org_A", "team_in_org_B"] assert await list_teams("org_admin_user") == own_view @@ -4632,7 +4629,7 @@ async def test_list_team_v2_org_admin_with_user_id_returns_user_teams(): status=None, ) - assert result["total"] == 1 + assert result.total == 1 # Verify the where clause filters by user's teams AND org scope where = mock_db.litellm_teamtable.find_many.call_args.kwargs["where"] @@ -4902,8 +4899,8 @@ async def test_list_team_v2_populates_keys_count(): status=None, ) - assert result["total"] == 2 - by_id = {t.team_id: t for t in result["teams"]} + 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. @@ -4943,6 +4940,13 @@ def _team_list_member( ) +def _team_list_caller_identity( + user_id: str | None = "caller", + user_email: str | None = None, +) -> _TeamListCallerIdentity: + return _TeamListCallerIdentity(user_id=user_id, user_email=user_email) + + def _budget_table( budget_id: str, max_budget: float | None, @@ -4975,8 +4979,7 @@ async def test_team_list_caller_membership_uses_member_spend_budget_and_reset_da teams: Final = await _enrich_team_list_with_caller_membership( prisma_client, [_team_list_member("team-a")], - "caller", - None, + _team_list_caller_identity(), datetime(2026, 10, 1, tzinfo=timezone.utc), ) @@ -4999,8 +5002,7 @@ async def test_team_list_caller_membership_includes_roster_members_without_membe teams: Final = await _enrich_team_list_with_caller_membership( prisma_client, [_team_list_member("team-a", default_budget_id="team-budget")], - "caller", - None, + _team_list_caller_identity(), datetime(2026, 10, 1, tzinfo=timezone.utc), ) @@ -5029,8 +5031,7 @@ async def test_team_list_caller_membership_matches_email_only_roster_member() -> teams: Final = await _enrich_team_list_with_caller_membership( prisma_client, [team], - "caller", - "caller@x.com", + _team_list_caller_identity(user_email="caller@x.com"), datetime(2026, 10, 1, tzinfo=timezone.utc), ) @@ -5040,7 +5041,7 @@ async def test_team_list_caller_membership_matches_email_only_roster_member() -> budget_reset_at=None, ) membership_query.assert_awaited_once() - budget_query.assert_awaited_once_with(where={"budget_id": {"in": ("team-budget",)}}) + budget_query.assert_awaited_once() @pytest.mark.parametrize("caller_user_email", ("other@x.com", None)) @@ -5061,8 +5062,7 @@ async def test_team_list_caller_membership_ignores_email_only_roster_member_with teams: Final = await _enrich_team_list_with_caller_membership( prisma_client, [team], - "caller", - caller_user_email, + _team_list_caller_identity(user_email=caller_user_email), datetime(2026, 10, 1, tzinfo=timezone.utc), ) @@ -5083,8 +5083,7 @@ async def test_team_list_caller_membership_uses_team_default_without_member_rese teams: Final = await _enrich_team_list_with_caller_membership( prisma_client, [_team_list_member("team-a", default_budget_id="team-budget")], - "caller", - None, + _team_list_caller_identity(), datetime(2026, 10, 1, tzinfo=timezone.utc), ) @@ -5093,7 +5092,7 @@ async def test_team_list_caller_membership_uses_team_default_without_member_rese max_budget=50, budget_reset_at=None, ) - budget_query.assert_awaited_once_with(where={"budget_id": {"in": ("team-budget",)}}) + budget_query.assert_awaited_once() @pytest.mark.asyncio @@ -5110,8 +5109,7 @@ async def test_team_list_caller_membership_uses_reset_date_from_member_linked_de teams: Final = await _enrich_team_list_with_caller_membership( prisma_client, [_team_list_member("team-a", default_budget_id="team-budget")], - "caller", - None, + _team_list_caller_identity(), datetime(2026, 10, 1, tzinfo=timezone.utc), ) @@ -5134,8 +5132,7 @@ async def test_team_list_caller_membership_ignores_zero_team_default_budget() -> teams: Final = await _enrich_team_list_with_caller_membership( prisma_client, [_team_list_member("team-a", default_budget_id="team-budget")], - "caller", - None, + _team_list_caller_identity(), datetime(2026, 10, 1, tzinfo=timezone.utc), ) @@ -5171,8 +5168,7 @@ async def test_team_list_caller_membership_applies_only_active_temporary_increas teams: Final = await _enrich_team_list_with_caller_membership( prisma_client, [_team_list_member("team-a")], - "caller", - None, + _team_list_caller_identity(), datetime(2030, 1, 1, tzinfo=timezone.utc), ) @@ -5213,8 +5209,7 @@ async def test_team_list_caller_membership_uses_private_row_reset_with_inherited teams: Final = await _enrich_team_list_with_caller_membership( prisma_client, [_team_list_member("team-a", default_budget_id="team-budget")], - "caller", - None, + _team_list_caller_identity(), datetime(2030, 1, 1, tzinfo=timezone.utc), ) @@ -5237,8 +5232,7 @@ async def test_team_list_caller_membership_is_null_for_nonmember_admins() -> Non teams: Final = await _enrich_team_list_with_caller_membership( prisma_client, [team], - "caller", - None, + _team_list_caller_identity(), datetime(2026, 10, 1, tzinfo=timezone.utc), ) @@ -5254,8 +5248,7 @@ async def test_team_list_caller_membership_skips_queries_without_caller_user_id( teams: Final = await _enrich_team_list_with_caller_membership( prisma_client, [_team_list_member("team-a")], - None, - None, + _team_list_caller_identity(user_id=None), datetime(2026, 10, 1, tzinfo=timezone.utc), ) @@ -5281,17 +5274,16 @@ async def test_team_list_caller_membership_batches_page_membership_and_budget_qu _team_list_member("team-a", default_budget_id="budget-a"), _team_list_member("team-b", default_budget_id="budget-b"), ], - "caller", - None, + _team_list_caller_identity(), datetime(2026, 10, 1, tzinfo=timezone.utc), ) - assert [team.caller_membership.max_budget for team in teams if team.caller_membership is not None] == [50, 75] - membership_query.assert_awaited_once_with( - where={"team_id": {"in": ("team-a", "team-b")}, "user_id": "caller"}, - include={"litellm_budget_table": True}, - ) - budget_query.assert_awaited_once_with(where={"budget_id": {"in": ("budget-a", "budget-b")}}) + assert [team.caller_membership for team in teams] == [ + TeamListCallerMembership(spend=0.0, max_budget=50, budget_reset_at=None), + TeamListCallerMembership(spend=0.0, max_budget=75, budget_reset_at=None), + ] + membership_query.assert_awaited_once() + budget_query.assert_awaited_once() @pytest.mark.asyncio @@ -5329,8 +5321,8 @@ async def test_list_team_v2_keys_count_skipped_for_empty_page(): status=None, ) - assert result["total"] == 0 - assert result["teams"] == [] + assert result.total == 0 + assert result.teams == [] mock_db.litellm_verificationtoken.group_by.assert_not_called() @@ -5379,7 +5371,7 @@ async def test_list_team_v2_keys_count_skipped_for_deleted_status(): status="deleted", ) - assert result["total"] == 1 + assert result.total == 1 mock_db.litellm_verificationtoken.group_by.assert_not_called() diff --git a/ui/litellm-dashboard/src/components/TeamsPage/TeamMemberSpendBudgetCell.test.tsx b/ui/litellm-dashboard/src/components/TeamsPage/TeamMemberSpendBudgetCell.test.tsx index bfc47f740ab..e5ff85460d0 100644 --- a/ui/litellm-dashboard/src/components/TeamsPage/TeamMemberSpendBudgetCell.test.tsx +++ b/ui/litellm-dashboard/src/components/TeamsPage/TeamMemberSpendBudgetCell.test.tsx @@ -29,7 +29,7 @@ describe("TeamMemberSpendBudgetCell", () => { }); it("shows the team budget and Unlimited for a member without a budget", () => { - const { container } = render( + render( { expect(screen.getByText("Member")).toBeInTheDocument(); expect(screen.getByText("/ Unlimited")).toBeInTheDocument(); expect(screen.getAllByRole("meter")).toHaveLength(1); - expect(container.querySelectorAll('[data-slot="meter-indicator"]')).toHaveLength(1); }); it("omits the reset text when no member reset date exists", () => { @@ -65,17 +64,16 @@ describe("TeamMemberSpendBudgetCell", () => { expect(screen.getAllByRole("meter")).toHaveLength(1); }); - it("uses the over tone when member spend exceeds the budget", () => { - const { container } = render( + it("exposes member spend above its budget through the meter", () => { + render( , ); - const indicators = container.querySelectorAll('[data-slot="meter-indicator"]'); - expect(indicators[1]).toHaveClass("bg-destructive"); + expect(screen.getAllByRole("meter")[1]).toHaveAttribute("aria-valuetext", "Member $150.00 of $100.00"); }); it("does not render meters for zero or unlimited budgets", () => { diff --git a/ui/litellm-dashboard/src/components/TeamsPage/TeamMemberSpendBudgetCell.tsx b/ui/litellm-dashboard/src/components/TeamsPage/TeamMemberSpendBudgetCell.tsx index d060ae350ad..45289fd3ffc 100644 --- a/ui/litellm-dashboard/src/components/TeamsPage/TeamMemberSpendBudgetCell.tsx +++ b/ui/litellm-dashboard/src/components/TeamsPage/TeamMemberSpendBudgetCell.tsx @@ -2,14 +2,11 @@ import { Meter, MeterIndicator, MeterTrack } from "@/components/shared/Meter"; import { getSpendBudgetMeterTone } from "@/components/shared/table_cells/spend_budget_cell"; +import type { components } from "@/lib/http/schema"; import { formatBudgetResetUtc } from "@/utils/budgetUtils"; import { formatNumberWithCommas, getSpendString } from "@/utils/dataUtils"; -interface CallerMembership { - spend: number; - max_budget: number | null; - budget_reset_at: string | null; -} +type CallerMembership = components["schemas"]["TeamListCallerMembership"]; interface TeamMemberSpendBudgetCellProps { teamSpend: number | null | undefined; @@ -29,7 +26,7 @@ function BudgetLine({ label, spend, maxBudget, resetAt }: BudgetLineProps) { const budget = typeof maxBudget === "number" && Number.isFinite(maxBudget) ? maxBudget : null; const spendText = spendValue === 0 ? "$0.00" : getSpendString(spendValue, 2); const budgetText = budget === null ? "Unlimited" : `$${formatNumberWithCommas(budget, 2)}`; - const resetText = label === "Member" ? formatBudgetResetUtc(resetAt) : null; + const resetText = formatBudgetResetUtc(resetAt); return (
diff --git a/ui/litellm-dashboard/src/components/TeamsPage/TeamsTable.test.tsx b/ui/litellm-dashboard/src/components/TeamsPage/TeamsTable.integration.test.tsx similarity index 100% rename from ui/litellm-dashboard/src/components/TeamsPage/TeamsTable.test.tsx rename to ui/litellm-dashboard/src/components/TeamsPage/TeamsTable.integration.test.tsx 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 09b08e1c933..11dd4f2c2e1 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 @@ -1,4 +1,5 @@ import { Member } from "../networking"; +import type { components } from "@/lib/http/schema"; import type { ObjectPermission } from "../object_permission_types"; import type { ModelBudgetUsage, ModelMaxBudget } from "./ModelMaxBudgetEditor"; @@ -21,11 +22,7 @@ export interface Team { keys_count?: number; members_count?: number; members_with_roles: Member[]; - caller_membership?: { - spend: number; - max_budget: number | null; - budget_reset_at: string | null; - } | null; + caller_membership?: components["schemas"]["TeamListCallerMembership"] | null; team_member_permissions?: string[] | null; spend: number; access_group_ids?: string[];