refactor(teams): address review on caller membership enrichment

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

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(

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()
 
 

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(
       <TeamMemberSpendBudgetCell
         teamSpend={600}
         teamMaxBudget={1000}
@@ -42,7 +42,6 @@ describe("TeamMemberSpendBudgetCell", () => {
     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(
       <TeamMemberSpendBudgetCell
         teamSpend={10}
         teamMaxBudget={100}
         callerMembership={{ spend: 150, max_budget: 100, budget_reset_at: null }}
       />,
     );
-    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", () => {

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 (
     <div className="flex flex-col gap-1">

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[];