mirror of
https://github.com/BerriAI/litellm.git
synced 2026-08-28 05:25:59 +00:00
fix(proxy/team): resolve member_delete cleanup by user id, not the addressed email (#36839)
/team/member_delete dropped the roster entry by matching user_email against members_with_roles, then built its user-row lookup from that same raw email instead of from the user_id the roster entry already carries. An email the user row does not literally hold matched nothing, so the team id stayed in the user's teams array and the team-membership row was left orphaned while the call still returned 200. /team/member_add resolves an email to a user case-insensitively but stores the caller's casing on the roster, so inviting "Alice@Example.com" for a row holding "alice@example.com" and removing by that same string is enough to reach it. _cleanup_members_with_roles now returns the roster entries it removed, and both the user-row update and the membership delete run against their user ids.
This commit is contained in:
parent
ab2333b6c4
commit
56b08c19d6
3 changed files with 141 additions and 50 deletions
|
|
@ -2311,7 +2311,7 @@ async def delete_user(
|
|||
fetch_all_teams = await TeamRepository(prisma_client).table.find_many(where={"team_id": {"in": user_row.teams}})
|
||||
teams_to_update = []
|
||||
for team in fetch_all_teams:
|
||||
is_member_in_team, new_team_members = _cleanup_members_with_roles(
|
||||
removed_team_members, new_team_members = _cleanup_members_with_roles(
|
||||
existing_team_row=LiteLLM_TeamTable.model_validate(team.model_dump()),
|
||||
data=TeamMemberDeleteRequest(
|
||||
team_id=team.team_id,
|
||||
|
|
@ -2319,7 +2319,7 @@ async def delete_user(
|
|||
user_email=user_row.user_email,
|
||||
),
|
||||
)
|
||||
if is_member_in_team:
|
||||
if removed_team_members:
|
||||
_db_new_team_members: list[dict] = [m.model_dump() for m in new_team_members]
|
||||
team.members_with_roles = json.dumps(_db_new_team_members)
|
||||
teams_to_update.append(team)
|
||||
|
|
|
|||
|
|
@ -3091,26 +3091,27 @@ async def team_member_add(
|
|||
)
|
||||
|
||||
|
||||
def _is_member_addressed_by(member: Member, data: TeamMemberDeleteRequest) -> bool:
|
||||
return (data.user_id is not None and member.user_id is not None and data.user_id == member.user_id) or (
|
||||
data.user_email is not None and member.user_email is not None and data.user_email == member.user_email
|
||||
)
|
||||
|
||||
|
||||
def _cleanup_members_with_roles(
|
||||
existing_team_row: LiteLLM_TeamTable,
|
||||
data: TeamMemberDeleteRequest,
|
||||
) -> tuple[bool, list[Member]]:
|
||||
"""Cleanup members_with_roles list for a team."""
|
||||
is_member_in_team = False
|
||||
new_team_members: Final[list[Member]] = []
|
||||
for m in existing_team_row.members_with_roles:
|
||||
if (
|
||||
data.user_id is not None
|
||||
and m.user_id is not None
|
||||
and data.user_id == m.user_id
|
||||
or data.user_email is not None
|
||||
and m.user_email is not None
|
||||
and data.user_email == m.user_email
|
||||
):
|
||||
is_member_in_team = True
|
||||
continue
|
||||
new_team_members.append(m)
|
||||
return is_member_in_team, new_team_members
|
||||
) -> tuple[tuple[Member, ...], list[Member]]:
|
||||
"""Split a team's members_with_roles into the entries the request addresses and the ones that stay.
|
||||
|
||||
The addressed entries are returned rather than a bare found/not-found flag because they carry the
|
||||
user_id the request may not have supplied, and every cleanup that keys off the user rather than
|
||||
off the roster has to run against that id.
|
||||
"""
|
||||
removed_team_members: Final = tuple(
|
||||
m for m in existing_team_row.members_with_roles if _is_member_addressed_by(m, data)
|
||||
)
|
||||
new_team_members: Final = [m for m in existing_team_row.members_with_roles if not _is_member_addressed_by(m, data)]
|
||||
return removed_team_members, new_team_members
|
||||
|
||||
|
||||
@router.post(
|
||||
|
|
@ -3182,12 +3183,12 @@ async def team_member_delete(
|
|||
)
|
||||
|
||||
## DELETE MEMBER FROM TEAM
|
||||
is_member_in_team, new_team_members = _cleanup_members_with_roles(
|
||||
removed_team_members, new_team_members = _cleanup_members_with_roles(
|
||||
existing_team_row=existing_team_row,
|
||||
data=data,
|
||||
)
|
||||
|
||||
if not is_member_in_team:
|
||||
if not removed_team_members:
|
||||
raise HTTPException(status_code=400, detail={"error": "User not found in team"})
|
||||
|
||||
existing_team_row.members_with_roles = new_team_members
|
||||
|
|
@ -3205,38 +3206,28 @@ async def team_member_delete(
|
|||
|
||||
## DELETE TEAM ID from USER ROW, IF EXISTS ##
|
||||
# get user row
|
||||
key_val: Final = {}
|
||||
if data.user_id is not None:
|
||||
key_val["user_id"] = data.user_id
|
||||
elif data.user_email is not None:
|
||||
key_val["user_email"] = data.user_email
|
||||
existing_user_rows: Final[Sequence[LiteLLM_UserTable] | None] = await UserRepository(prisma_client).table.find_many(
|
||||
where=key_val
|
||||
removed_user_ids: Final = frozenset(m.user_id for m in removed_team_members if m.user_id is not None)
|
||||
key_val: Final[Mapping[str, object]] = (
|
||||
{"user_id": {"in": sorted(removed_user_ids)}} if removed_user_ids else {"user_email": data.user_email}
|
||||
)
|
||||
existing_user_rows: Final[Sequence[LiteLLM_UserTable]] = await _user_db(prisma_client).find_many(where=key_val)
|
||||
|
||||
if existing_user_rows is not None and (isinstance(existing_user_rows, list) and len(existing_user_rows) > 0):
|
||||
for existing_user in existing_user_rows:
|
||||
team_list = []
|
||||
if data.team_id in existing_user.teams:
|
||||
team_list = existing_user.teams
|
||||
team_list.remove(data.team_id)
|
||||
await _user_db(prisma_client).update(
|
||||
where={
|
||||
"user_id": existing_user.user_id,
|
||||
},
|
||||
data={"teams": {"set": team_list}},
|
||||
)
|
||||
for existing_user in existing_user_rows:
|
||||
if data.team_id in existing_user.teams:
|
||||
await _user_db(prisma_client).update(
|
||||
where={
|
||||
"user_id": existing_user.user_id,
|
||||
},
|
||||
data={"teams": {"set": [team for team in existing_user.teams if team != data.team_id]}},
|
||||
)
|
||||
|
||||
# Also clean up any existing team membership rows for this user and team
|
||||
user_ids_to_delete: Final = set[str]()
|
||||
if data.user_id is not None:
|
||||
user_ids_to_delete.add(data.user_id)
|
||||
if existing_user_rows is not None and isinstance(existing_user_rows, list):
|
||||
for existing_user in existing_user_rows:
|
||||
if getattr(existing_user, "user_id", None):
|
||||
user_ids_to_delete.add(existing_user.user_id)
|
||||
user_ids_to_delete: Final = removed_user_ids.union(
|
||||
(data.user_id,) if data.user_id is not None else (),
|
||||
(user.user_id for user in existing_user_rows if user.user_id),
|
||||
)
|
||||
|
||||
for _uid in user_ids_to_delete:
|
||||
for _uid in sorted(user_ids_to_delete):
|
||||
await _team_membership_db(prisma_client).delete_many(where={"team_id": data.team_id, "user_id": _uid})
|
||||
|
||||
## DELETE KEYS CREATED BY USER FOR THIS TEAM
|
||||
|
|
@ -3248,7 +3239,7 @@ async def team_member_delete(
|
|||
# Fetch keys before deletion to persist them
|
||||
keys_to_delete: Final[list[LiteLLM_VerificationToken]] = await _tokens_db(prisma_client).find_many(
|
||||
where={
|
||||
"user_id": {"in": list(user_ids_to_delete)},
|
||||
"user_id": {"in": sorted(user_ids_to_delete)},
|
||||
"team_id": data.team_id,
|
||||
}
|
||||
)
|
||||
|
|
@ -3263,7 +3254,7 @@ async def team_member_delete(
|
|||
|
||||
await _tokens_db(prisma_client).delete_many(
|
||||
where={
|
||||
"user_id": {"in": list(user_ids_to_delete)},
|
||||
"user_id": {"in": sorted(user_ids_to_delete)},
|
||||
"team_id": data.team_id,
|
||||
}
|
||||
)
|
||||
|
|
|
|||
|
|
@ -4133,6 +4133,106 @@ async def test_team_member_delete_cleans_verification_tokens(
|
|||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"roster_email",
|
||||
["Alice@Example.com", "alice-invited-as@example.com"],
|
||||
ids=["case_variant_of_the_row_email", "email_the_row_never_carried"],
|
||||
)
|
||||
@pytest.mark.parametrize("user_row_exists", [True, False])
|
||||
@pytest.mark.asyncio
|
||||
async def test_team_member_delete_by_email_the_user_row_does_not_carry(
|
||||
user_row_exists, roster_email, mock_db_client, mock_admin_auth
|
||||
):
|
||||
"""
|
||||
Removing a member addressed by user_email drove its user-row and membership cleanup off that raw
|
||||
email instead of off the user_id the roster entry already carries, so an email the user row does
|
||||
not literally hold matched nothing and both cleanups silently no-opped behind a 200.
|
||||
|
||||
Both roster emails here are reachable over plain HTTP. /team/member_add resolves an email to a
|
||||
user case-insensitively but stores the caller's casing in members_with_roles, which produces the
|
||||
case variant; it also leaves an unmatched email on the entry when no user row carries it at all,
|
||||
which produces the second. Both converge on the same lookup, so they are parametrized inputs
|
||||
rather than separate paths, and each one has to detect the bug on its own.
|
||||
|
||||
The user table below is case-sensitive like Postgres, so only a lookup driven by the resolved
|
||||
user_id finds the row. The user_row_exists=False leg pins the second half on its own: the
|
||||
membership row has to go even when no user row is left to resolve it from.
|
||||
"""
|
||||
from litellm.proxy._types import TeamMemberDeleteRequest
|
||||
from litellm.proxy.management_endpoints.team_endpoints import team_member_delete
|
||||
|
||||
test_team_id = "team-del-email-case-123"
|
||||
test_user_id = "user-del-email-case-123"
|
||||
user_row_email = "alice@example.com"
|
||||
|
||||
mock_team_row = MagicMock()
|
||||
mock_team_row.model_dump.return_value = {
|
||||
"team_id": test_team_id,
|
||||
"members_with_roles": [
|
||||
{"user_id": test_user_id, "user_email": roster_email, "role": "user"}
|
||||
],
|
||||
"team_member_permissions": [],
|
||||
"metadata": {},
|
||||
"models": [],
|
||||
"spend": 0.0,
|
||||
}
|
||||
|
||||
mock_db_client.db.litellm_teamtable.find_unique = AsyncMock(
|
||||
return_value=mock_team_row
|
||||
)
|
||||
mock_db_client.db.litellm_teamtable.update = AsyncMock(return_value=mock_team_row)
|
||||
|
||||
mock_user_row = MagicMock()
|
||||
mock_user_row.user_id = test_user_id
|
||||
mock_user_row.user_email = user_row_email
|
||||
mock_user_row.teams = [test_team_id]
|
||||
|
||||
async def find_user_rows(where):
|
||||
if not user_row_exists:
|
||||
return []
|
||||
user_id_filter = where.get("user_id")
|
||||
if isinstance(user_id_filter, dict) and test_user_id in user_id_filter.get(
|
||||
"in", []
|
||||
):
|
||||
return [mock_user_row]
|
||||
if where.get("user_email") == user_row_email:
|
||||
return [mock_user_row]
|
||||
return []
|
||||
|
||||
mock_db_client.db.litellm_usertable.find_many = AsyncMock(
|
||||
side_effect=find_user_rows
|
||||
)
|
||||
mock_db_client.db.litellm_usertable.update = AsyncMock(return_value=MagicMock())
|
||||
|
||||
mock_db_client.db.litellm_teammembership = MagicMock()
|
||||
mock_db_client.db.litellm_teammembership.delete_many = AsyncMock(
|
||||
return_value=MagicMock()
|
||||
)
|
||||
|
||||
mock_db_client.db.litellm_verificationtoken = MagicMock()
|
||||
mock_db_client.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[])
|
||||
mock_db_client.db.litellm_verificationtoken.delete_many = AsyncMock(
|
||||
return_value=MagicMock()
|
||||
)
|
||||
|
||||
await team_member_delete(
|
||||
data=TeamMemberDeleteRequest(team_id=test_team_id, user_email=roster_email),
|
||||
user_api_key_dict=mock_admin_auth,
|
||||
)
|
||||
|
||||
if user_row_exists:
|
||||
mock_db_client.db.litellm_usertable.update.assert_awaited_once_with(
|
||||
where={"user_id": test_user_id},
|
||||
data={"teams": {"set": []}},
|
||||
)
|
||||
else:
|
||||
mock_db_client.db.litellm_usertable.update.assert_not_awaited()
|
||||
|
||||
mock_db_client.db.litellm_teammembership.delete_many.assert_awaited_once_with(
|
||||
where={"team_id": test_team_id, "user_id": test_user_id}
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_new_team_max_budget_exceeds_user_max_budget():
|
||||
"""
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue