From 202d963f8f3aa00df4c096eee713796bf0287dd5 Mon Sep 17 00:00:00 2001 From: jesus Date: Sat, 10 Oct 2026 23:46:36 +0000 Subject: [PATCH] fix(teams): match email-only roster members for caller membership Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../management_endpoints/team_endpoints.py | 28 ++++++-- .../test_team_endpoints.py | 68 +++++++++++++++++++ 2 files changed, 92 insertions(+), 4 deletions(-) diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 01575e4e787..3da231b1551 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -5695,13 +5695,22 @@ 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 + ) + + def _team_member_default_budget_id( team: TeamListItem, caller_user_id: str, + caller_user_email: str | None, 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 any(member.user_id == caller_user_id for member in team.members_with_roles): + if membership_row is None and not _caller_in_roster(team, caller_user_id, caller_user_email): return None member_budget: Final = _team_list_budget_model( membership_row.litellm_budget_table if membership_row is not None else None @@ -5720,16 +5729,19 @@ def _team_member_default_budget_id( def _team_list_caller_membership( team: TeamListItem, caller_user_id: str, + caller_user_email: str | None, memberships_by_team: Mapping[str, _TeamListMembershipRow], 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 any(member.user_id == caller_user_id for member in team.members_with_roles): + if membership_row is None and not _caller_in_roster(team, caller_user_id, caller_user_email): 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, memberships_by_team) + 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( @@ -5751,6 +5763,7 @@ async def _enrich_team_list_with_caller_membership( prisma_client: PrismaClient, teams: Sequence[TeamListItem], caller_user_id: str | None, + caller_user_email: str | None, now: datetime, ) -> tuple[TeamListItem, ...]: if caller_user_id is None: @@ -5769,7 +5782,12 @@ async def _enrich_team_list_with_caller_membership( dict.fromkeys( budget_id for team in teams - if (budget_id := _team_member_default_budget_id(team, caller_user_id, memberships_by_team)) is not None + if ( + budget_id := _team_member_default_budget_id( + team, caller_user_id, caller_user_email, memberships_by_team + ) + ) + is not None ) ) default_budget_rows: Final[Sequence[_TeamListBudgetRow]] = ( @@ -5787,6 +5805,7 @@ async def _enrich_team_list_with_caller_membership( "caller_membership": _team_list_caller_membership( team=team, caller_user_id=caller_user_id, + caller_user_email=caller_user_email, memberships_by_team=memberships_by_team, budgets_by_id=budgets_by_id, now=now, @@ -6062,6 +6081,7 @@ async def list_team_v2( 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, now=get_utc_datetime(), ) if not use_deleted_table and team_list diff --git a/tests/unit/proxy/management_endpoints/test_team_endpoints.py b/tests/unit/proxy/management_endpoints/test_team_endpoints.py index dd3cc2d2874..330f8d8310e 100644 --- a/tests/unit/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_team_endpoints.py @@ -4976,6 +4976,7 @@ async def test_team_list_caller_membership_uses_member_spend_budget_and_reset_da prisma_client, [_team_list_member("team-a")], "caller", + None, datetime(2026, 10, 1, tzinfo=timezone.utc), ) @@ -4999,6 +5000,7 @@ async def test_team_list_caller_membership_includes_roster_members_without_membe prisma_client, [_team_list_member("team-a", default_budget_id="team-budget")], "caller", + None, datetime(2026, 10, 1, tzinfo=timezone.utc), ) @@ -5011,6 +5013,64 @@ async def test_team_list_caller_membership_includes_roster_members_without_membe budget_query.assert_awaited_once() +@pytest.mark.asyncio +async def test_team_list_caller_membership_matches_email_only_roster_member() -> None: + reset_at: Final = datetime(2026, 10, 20, tzinfo=timezone.utc) + team: Final = TeamListItem( + team_id="team-a", + metadata={"team_member_budget_id": "team-budget"}, + members_with_roles=[Member(user_id=None, user_email="caller@x.com", role="user")], + ) + prisma_client, membership_query, budget_query = _team_list_prisma( + (), + [_budget_table("team-budget", 50, reset_at)], + ) + + teams: Final = await _enrich_team_list_with_caller_membership( + prisma_client, + [team], + "caller", + "caller@x.com", + datetime(2026, 10, 1, tzinfo=timezone.utc), + ) + + assert teams[0].caller_membership == TeamListCallerMembership( + spend=0.0, + max_budget=50, + budget_reset_at=None, + ) + membership_query.assert_awaited_once() + budget_query.assert_awaited_once_with(where={"budget_id": {"in": ("team-budget",)}}) + + +@pytest.mark.parametrize("caller_user_email", ("other@x.com", None)) +@pytest.mark.asyncio +async def test_team_list_caller_membership_ignores_email_only_roster_member_without_matching_email( + caller_user_email: str | None, +) -> None: + team: Final = TeamListItem( + team_id="team-a", + metadata={"team_member_budget_id": "team-budget"}, + members_with_roles=[Member(user_id=None, user_email="caller@x.com", role="user")], + ) + prisma_client, membership_query, budget_query = _team_list_prisma( + (), + [_budget_table("team-budget", 50)], + ) + + teams: Final = await _enrich_team_list_with_caller_membership( + prisma_client, + [team], + "caller", + caller_user_email, + datetime(2026, 10, 1, tzinfo=timezone.utc), + ) + + assert teams[0].caller_membership is None + membership_query.assert_awaited_once() + budget_query.assert_not_awaited() + + @pytest.mark.asyncio async def test_team_list_caller_membership_uses_team_default_without_member_reset_date() -> None: reset_at: Final = datetime(2026, 10, 20, tzinfo=timezone.utc) @@ -5024,6 +5084,7 @@ async def test_team_list_caller_membership_uses_team_default_without_member_rese prisma_client, [_team_list_member("team-a", default_budget_id="team-budget")], "caller", + None, datetime(2026, 10, 1, tzinfo=timezone.utc), ) @@ -5050,6 +5111,7 @@ async def test_team_list_caller_membership_uses_reset_date_from_member_linked_de prisma_client, [_team_list_member("team-a", default_budget_id="team-budget")], "caller", + None, datetime(2026, 10, 1, tzinfo=timezone.utc), ) @@ -5073,6 +5135,7 @@ async def test_team_list_caller_membership_ignores_zero_team_default_budget() -> prisma_client, [_team_list_member("team-a", default_budget_id="team-budget")], "caller", + None, datetime(2026, 10, 1, tzinfo=timezone.utc), ) @@ -5109,6 +5172,7 @@ async def test_team_list_caller_membership_applies_only_active_temporary_increas prisma_client, [_team_list_member("team-a")], "caller", + None, datetime(2030, 1, 1, tzinfo=timezone.utc), ) @@ -5150,6 +5214,7 @@ async def test_team_list_caller_membership_uses_private_row_reset_with_inherited prisma_client, [_team_list_member("team-a", default_budget_id="team-budget")], "caller", + None, datetime(2030, 1, 1, tzinfo=timezone.utc), ) @@ -5173,6 +5238,7 @@ async def test_team_list_caller_membership_is_null_for_nonmember_admins() -> Non prisma_client, [team], "caller", + None, datetime(2026, 10, 1, tzinfo=timezone.utc), ) @@ -5189,6 +5255,7 @@ async def test_team_list_caller_membership_skips_queries_without_caller_user_id( prisma_client, [_team_list_member("team-a")], None, + None, datetime(2026, 10, 1, tzinfo=timezone.utc), ) @@ -5215,6 +5282,7 @@ async def test_team_list_caller_membership_batches_page_membership_and_budget_qu _team_list_member("team-b", default_budget_id="budget-b"), ], "caller", + None, datetime(2026, 10, 1, tzinfo=timezone.utc), )