diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 82e22bb5bbf..a8e545a8551 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -2640,51 +2640,63 @@ async def _process_team_members( return updated_users, updated_team_memberships +def _resolve_member_identity(member: Member, updated_users: Sequence[LiteLLM_UserTable]) -> Member: + """Return ``member`` with whichever of ``user_id`` / ``user_email`` the caller left out filled in. + + The roster entry is a snapshot, so whatever is missing here is missing for good. + Resolution runs both ways off the user rows the add just touched: added by email + -> stamp the user_id, added by user_id -> stamp the email. A value the caller + supplied is never overwritten. + """ + resolved_user_id: Final = member.user_id or next( + ( + user.user_id + for user in updated_users + if member.user_email is not None and user.user_email == member.user_email + ), + None, + ) + resolved_user_email: Final = member.user_email or next( + ( + user.user_email + for user in updated_users + if resolved_user_id is not None and user.user_id == resolved_user_id and user.user_email is not None + ), + None, + ) + return member.model_copy( + update={ # mutable-ok: pydantic update payload + "user_id": resolved_user_id, + "user_email": resolved_user_email, + } + ) + + +def _member_already_in_team(member: Member, complete_team_data: LiteLLM_TeamTable) -> bool: + return any( + (member.user_id is not None and existing_member.user_id == member.user_id) + or (member.user_email is not None and existing_member.user_email == member.user_email) + for existing_member in complete_team_data.members_with_roles + ) + + async def _update_team_members_list( data: TeamMemberAddRequest, complete_team_data: LiteLLM_TeamTable, updated_users: list[LiteLLM_UserTable], ) -> None: """Update the team's members_with_roles list.""" - if isinstance(data.member, Member): - new_member: Final = data.member.model_copy() + requested_members: Final[Sequence[Member]] = ( + (data.member,) if isinstance(data.member, Member) else tuple(data.member) + ) + resolved_members: Final = tuple(_resolve_member_identity(m, updated_users) for m in requested_members) - # get user id - if new_member.user_id is None and new_member.user_email is not None: - for user in updated_users: - if user.user_email is not None and user.user_email == new_member.user_email: - new_member.user_id = user.user_id - - # Check if member already exists in team before adding - member_already_exists = False - for existing_member in complete_team_data.members_with_roles: - if (new_member.user_id is not None and existing_member.user_id == new_member.user_id) or ( - new_member.user_email is not None and existing_member.user_email == new_member.user_email - ): - member_already_exists = True - break - - if not member_already_exists: - complete_team_data.members_with_roles.append(new_member) - - elif isinstance(data.member, list): - for nm in data.member: - if nm.user_id is None and nm.user_email is not None: - for user in updated_users: - if user.user_email is not None and user.user_email == nm.user_email: - nm.user_id = user.user_id - - # Check if member already exists in team before adding - member_already_exists = False - for existing_member in complete_team_data.members_with_roles: - if (nm.user_id is not None and existing_member.user_id == nm.user_id) or ( - nm.user_email is not None and existing_member.user_email == nm.user_email - ): - member_already_exists = True - break - - if not member_already_exists: - complete_team_data.members_with_roles.append(nm) + # extend() consumes the generator as it appends, so a member already added by this + # same call is seen by the next _member_already_in_team check - the batch dedupes + # against itself exactly as the append-one-at-a-time loop this replaced did. + complete_team_data.members_with_roles.extend( # rebind-ok: this helper's contract is to grow the caller's roster in place + m for m in resolved_members if not _member_already_in_team(m, complete_team_data) + ) async def _add_team_members_to_team( @@ -4086,6 +4098,39 @@ async def _add_team_member_budget_table( return team_info_response_object +async def _hydrate_member_emails( + prisma_client: PrismaClient, + members: Sequence[Member], +) -> tuple[Member, ...]: + """Fill in ``user_email`` for roster entries that were stored without one. + + ``members_with_roles`` is a denormalized snapshot written at add-time, so an entry + stored with ``user_email=None`` keeps that null even once the user row has an email. + Look the missing ones up in ``LiteLLM_UserTable`` (one indexed query) and fill them + in. A stored email is never overwritten - the snapshot stays the source of truth + wherever it has a value. + """ + missing_user_ids: Final = frozenset(m.user_id for m in members if not m.user_email and m.user_id is not None) + if not missing_user_ids: + return tuple(members) + + user_rows: Final[Sequence[LiteLLM_UserTable]] = await _user_db(prisma_client).find_many( + where={ # mutable-ok: Prisma query filters are dict-shaped + "user_id": { # mutable-ok: Prisma query filters are dict-shaped + "in": sorted(missing_user_ids) + } + } + ) + email_by_user_id: Final = MappingProxyType({u.user_id: u.user_email for u in user_rows if u.user_email}) + + return tuple( + m.model_copy(update={"user_email": email_by_user_id[m.user_id]}) # mutable-ok: pydantic update payload + if not m.user_email and m.user_id in email_by_user_id + else m + for m in members + ) + + async def _resolve_team_access_group_resources( _team_info: TeamInfoResponseObjectTeamTable, ) -> TeamInfoResponseObjectTeamTable: @@ -4221,9 +4266,22 @@ async def team_info( # Resolve resources inherited from access groups resolved_team_info: Final = await _resolve_team_access_group_resources(_team_info) + # Fill in emails the add-time roster snapshot never captured + hydrated_members: Final = await _hydrate_member_emails( + prisma_client=prisma_client, + members=resolved_team_info.members_with_roles, + ) + hydrated_team_info: Final = resolved_team_info.model_copy( + update={ # mutable-ok: pydantic update payload + # list(), not the tuple: model_copy skips validation, so the field has + # to be handed the list[Member] the response model declares. + "members_with_roles": list(hydrated_members) # mutable-ok: declared list[Member] + } + ) + response_object: Final = TeamInfoResponseObject( team_id=team_id, - team_info=resolved_team_info, + team_info=hydrated_team_info, keys=keys, team_memberships=returned_tm, ) 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 0767288d0bc..e39b09ae073 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -10253,6 +10253,65 @@ async def test_team_info_returns_model_aliases(): assert litellm_model_table.model_aliases == {"gpt-4o": "gpt-4o-team-1"} +@pytest.mark.asyncio +async def test_team_info_hydrates_member_emails_from_the_user_table(): + """/team/info must fill in emails missing from the members_with_roles snapshot. + + members_with_roles is written at add-time, so a member added by user_id alone + carries user_email=None forever. Without this join the Admin UI's member table + shows "-" for a user that has an email on their user row. A stored email is left + exactly as-is. + """ + from fastapi import Request + + from litellm.proxy.management_endpoints import team_endpoints + + team_row = LiteLLM_TeamTable( + team_id="team-1", + members_with_roles=[ + Member(user_id="no-email-on-roster", role="admin"), + Member(user_id="already-stored", user_email="stored@example.com", role="user"), + ], + ) + + mock_prisma = MagicMock() + mock_prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=team_row) + mock_prisma.get_data = AsyncMock(return_value=[]) + + find_many = AsyncMock( + return_value=[ + LiteLLM_UserTable( + user_id="no-email-on-roster", + user_email="real@example.com", + max_budget=None, + spend=0.0, + models=[], + ) + ] + ) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), + patch.object(team_endpoints, "get_all_team_memberships", AsyncMock(return_value=[])), + patch.object(team_endpoints, "UserRepository") as repo, + ): + repo.return_value.table.find_many = find_many + + response = await team_endpoints.team_info( + http_request=MagicMock(spec=Request), + team_id="team-1", + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), + ) + + members = response["team_info"].members_with_roles + assert [(m.user_id, m.user_email) for m in members] == [ + ("no-email-on-roster", "real@example.com"), + ("already-stored", "stored@example.com"), + ] + # only the member actually missing an email is looked up + assert find_many.await_args.kwargs["where"] == {"user_id": {"in": ["no-email-on-roster"]}} + + @pytest.mark.asyncio async def test_update_model_table_clears_aliases_with_empty_map(): """``model_aliases={}`` on /team/update must persist an empty map (json.dumps({})) @@ -11369,6 +11428,141 @@ async def test_resolve_existing_member_user_ids_skips_the_query_when_no_user_ids repo.return_value.table.find_many.assert_not_awaited() +def _user_row(user_id: str, user_email: str | None) -> LiteLLM_UserTable: + return LiteLLM_UserTable( + user_id=user_id, user_email=user_email, max_budget=None, spend=0.0, models=[] + ) + + +@pytest.mark.asyncio +async def test_hydrate_member_emails_fills_in_emails_the_roster_snapshot_never_captured(): + """A member added by user_id alone has user_email=None on the stored roster entry. + + /team/info has to fill it in from the user row, or the UI renders "-" for a user + that plainly has an email. + """ + from litellm.proxy.management_endpoints.team_endpoints import _hydrate_member_emails + + find_many = AsyncMock(return_value=[_user_row("by-id", "found@example.com")]) + + with patch("litellm.proxy.management_endpoints.team_endpoints.UserRepository") as repo: + repo.return_value.table.find_many = find_many + + hydrated = await _hydrate_member_emails( + prisma_client=MagicMock(), + members=[Member(user_id="by-id", role="admin")], + ) + + assert [(m.user_id, m.user_email, m.role) for m in hydrated] == [("by-id", "found@example.com", "admin")] + find_many.assert_awaited_once() + assert find_many.await_args.kwargs["where"] == {"user_id": {"in": ["by-id"]}} + + +@pytest.mark.asyncio +async def test_hydrate_member_emails_never_overwrites_a_stored_email(): + """The snapshot wins wherever it has a value - hydration only fills blanks. + + Overwriting would be a real behavior change to /team/info; filling a null is not. + """ + from litellm.proxy.management_endpoints.team_endpoints import _hydrate_member_emails + + find_many = AsyncMock(return_value=[_user_row("has-email", "current@example.com")]) + + with patch("litellm.proxy.management_endpoints.team_endpoints.UserRepository") as repo: + repo.return_value.table.find_many = find_many + + hydrated = await _hydrate_member_emails( + prisma_client=MagicMock(), + members=[Member(user_id="has-email", user_email="stored@example.com", role="user")], + ) + + assert hydrated[0].user_email == "stored@example.com" + # nothing was missing, so no round-trip either + find_many.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_hydrate_member_emails_leaves_members_alone_when_the_user_row_has_no_email(): + """A user row with no email leaves the member as-is rather than inventing one.""" + from litellm.proxy.management_endpoints.team_endpoints import _hydrate_member_emails + + with patch("litellm.proxy.management_endpoints.team_endpoints.UserRepository") as repo: + repo.return_value.table.find_many = AsyncMock(return_value=[_user_row("no-email", None)]) + + hydrated = await _hydrate_member_emails( + prisma_client=MagicMock(), + members=[Member(user_id="no-email", role="user"), Member(user_email="e@example.com", role="user")], + ) + + assert [m.user_email for m in hydrated] == [None, "e@example.com"] + + +@pytest.mark.asyncio +async def test_hydrate_member_emails_skips_the_query_when_every_member_has_one(): + """No blanks means /team/info pays for no extra query.""" + from litellm.proxy.management_endpoints.team_endpoints import _hydrate_member_emails + + with patch("litellm.proxy.management_endpoints.team_endpoints.UserRepository") as repo: + repo.return_value.table.find_many = AsyncMock() + + hydrated = await _hydrate_member_emails( + prisma_client=MagicMock(), + members=[Member(user_id="a", user_email="a@example.com", role="user")], + ) + + assert hydrated[0].user_email == "a@example.com" + repo.return_value.table.find_many.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_update_team_members_list_stamps_email_for_a_member_added_by_user_id(): + """Identity resolution runs both ways, so new roster entries stop being born blank. + + Previously only user_id was backfilled (from email); a member added by user_id + was written with user_email=None forever. + """ + from litellm.proxy.management_endpoints.team_endpoints import ( + _update_team_members_list, + ) + + mock_team = MagicMock(spec=LiteLLM_TeamTable) + mock_team.members_with_roles = [] + + await _update_team_members_list( + data=TeamMemberAddRequest(team_id="test-team-123", member=Member(user_id="new-user-123", role="user")), + complete_team_data=mock_team, + updated_users=[_user_row("new-user-123", "new@example.com")], + ) + + assert len(mock_team.members_with_roles) == 1 + assert mock_team.members_with_roles[0].user_email == "new@example.com" + + +@pytest.mark.asyncio +async def test_update_team_members_list_stamps_email_for_each_member_in_a_bulk_add(): + """Same both-ways resolution for the list branch.""" + from litellm.proxy.management_endpoints.team_endpoints import ( + _update_team_members_list, + ) + + mock_team = MagicMock(spec=LiteLLM_TeamTable) + mock_team.members_with_roles = [] + + await _update_team_members_list( + data=TeamMemberAddRequest( + team_id="test-team-123", + member=[Member(user_id="u1", role="user"), Member(user_email="u2@example.com", role="admin")], + ), + complete_team_data=mock_team, + updated_users=[_user_row("u1", "u1@example.com"), _user_row("u2", "u2@example.com")], + ) + + assert [(m.user_id, m.user_email) for m in mock_team.members_with_roles] == [ + ("u1", "u1@example.com"), + ("u2", "u2@example.com"), + ] + + def test_pre_existing_user_ids_counts_ids_filled_in_by_member_resolution(): """An id the member-resolution step filled in came from a matched row, so it pre-existed.