mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-24 05:45:11 +00:00
Merge 54c72de33c into b578dcf278
This commit is contained in:
Hyunseung Jung
GitHub
commit
02f0fec87e
6 changed files with 601 additions and 133 deletions
|
|
@ -2898,48 +2898,61 @@ def team_call_validation_checks(
|
|||
raise HTTPException(status_code=400, detail={"error": str(e)})
|
||||
|
||||
|
||||
def _requested_members(data: TeamMemberAddRequest) -> tuple[Member, ...]:
|
||||
return tuple(data.member) if isinstance(data.member, list) else (data.member,)
|
||||
|
||||
|
||||
def _first_repeated(values: Sequence[str | None]) -> str | None:
|
||||
return next((value for index, value in enumerate(values) if value is not None and value in values[:index]), None)
|
||||
|
||||
|
||||
def _validate_no_repeated_user_id(members: Sequence[Member]) -> None:
|
||||
repeated_user_id: Final = _first_repeated(tuple(member.user_id for member in members))
|
||||
if repeated_user_id is not None:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail={"error": f"{repeated_user_id} appears more than once in this request."},
|
||||
)
|
||||
|
||||
|
||||
def _validate_member_request_shape(members: Sequence[Member]) -> None:
|
||||
if len(members) == 0:
|
||||
raise HTTPException(status_code=400, detail={"error": "At least one member is required."})
|
||||
_validate_no_repeated_user_id(members)
|
||||
|
||||
|
||||
def _validate_resolved_members_are_distinct(members: Sequence[Member]) -> None:
|
||||
_validate_no_repeated_user_id(members)
|
||||
repeated_email: Final = _first_repeated(tuple(member.user_email for member in members))
|
||||
if repeated_email is not None:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail={"error": f"Several requested users share the email {repeated_email}. Add only one of them."},
|
||||
)
|
||||
|
||||
|
||||
def team_member_add_duplication_check(
|
||||
data: TeamMemberAddRequest,
|
||||
existing_team_row: LiteLLM_TeamTable,
|
||||
):
|
||||
"""
|
||||
Check if a member already exists in the team.
|
||||
This check is done BEFORE we create/fetch the user, so it only prevents
|
||||
obvious duplicates where both user_id and user_email match exactly.
|
||||
"""
|
||||
members: Final = _requested_members(data)
|
||||
invalid_team_members: Final = tuple(
|
||||
member for member in members if _member_already_in_team(member, existing_team_row)
|
||||
)
|
||||
|
||||
invalid_team_members: Final = []
|
||||
|
||||
def _check_member_duplication(member: Member):
|
||||
if member.user_id is not None:
|
||||
for existing_member in existing_team_row.members_with_roles:
|
||||
if existing_member.user_id == member.user_id:
|
||||
invalid_team_members.append(member)
|
||||
|
||||
# Check by user_email if provided
|
||||
if member.user_email is not None:
|
||||
for existing_member in existing_team_row.members_with_roles:
|
||||
if existing_member.user_email == member.user_email:
|
||||
invalid_team_members.append(member)
|
||||
|
||||
# First, populate the invalid_team_members list by checking for duplicates
|
||||
if isinstance(data.member, Member):
|
||||
_check_member_duplication(data.member)
|
||||
elif isinstance(data.member, list):
|
||||
for m in data.member:
|
||||
_check_member_duplication(m)
|
||||
|
||||
# Then check the populated list and raise exceptions if needed
|
||||
if isinstance(data.member, list) and len(invalid_team_members) == len(data.member):
|
||||
if isinstance(data.member, list) and members and len(invalid_team_members) == len(members):
|
||||
raise ProxyException(
|
||||
message=f"All users are already in team. Existing members={existing_team_row.members_with_roles}",
|
||||
message="All requested users are already members of this team.",
|
||||
type=ProxyErrorTypes.team_member_already_in_team,
|
||||
param="member",
|
||||
code="400",
|
||||
)
|
||||
elif isinstance(data.member, Member) and len(invalid_team_members) == 1:
|
||||
elif isinstance(data.member, Member) and (existing_member := _matching_team_member(data.member, existing_team_row)):
|
||||
raise ProxyException(
|
||||
message=f"User already in team. Member: user_id={data.member.user_id}, user_email={data.member.user_email}. Existing members={existing_team_row.members_with_roles}",
|
||||
message=_already_in_team_message(data.member, existing_member),
|
||||
type=ProxyErrorTypes.team_member_already_in_team,
|
||||
param="member",
|
||||
code="400",
|
||||
|
|
@ -3002,7 +3015,7 @@ async def _validate_team_member_add_permissions(
|
|||
# Available-team self-join: caller may add only themselves, only as a
|
||||
# standard user. Enforce that here so the bypass cannot be used as a
|
||||
# privilege-escalation or cross-user-injection primitive.
|
||||
members: Final = data.member if isinstance(data.member, list) else [data.member]
|
||||
members: Final = _requested_members(data)
|
||||
caller_user_id: Final = getattr(user_api_key_dict, "user_id", None)
|
||||
for member in members:
|
||||
if getattr(member, "role", "user") != "user":
|
||||
|
|
@ -3054,9 +3067,7 @@ async def _process_team_members(
|
|||
if member_allowed_models is None and team_default_member_models:
|
||||
member_allowed_models = team_default_member_models
|
||||
|
||||
requested_members: Final[Sequence[Member]] = (
|
||||
(data.member,) if isinstance(data.member, Member) else tuple(data.member)
|
||||
)
|
||||
requested_members: Final = _requested_members(data)
|
||||
for m in requested_members:
|
||||
if _member_already_in_team(m, complete_team_data):
|
||||
continue
|
||||
|
|
@ -3117,12 +3128,27 @@ def _resolve_member_identity(member: Member, updated_users: Sequence[LiteLLM_Use
|
|||
)
|
||||
|
||||
|
||||
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
|
||||
def _matching_team_member(member: Member, complete_team_data: LiteLLM_TeamTable) -> Member | None:
|
||||
roster: Final = complete_team_data.members_with_roles
|
||||
same_user: Final = next(
|
||||
(existing for existing in roster if member.user_id is not None and existing.user_id == member.user_id), None
|
||||
)
|
||||
if same_user is not None:
|
||||
return same_user
|
||||
return next(
|
||||
(existing for existing in roster if member.user_email is not None and existing.user_email == member.user_email),
|
||||
None,
|
||||
)
|
||||
|
||||
|
||||
def _member_already_in_team(member: Member, complete_team_data: LiteLLM_TeamTable) -> bool:
|
||||
return _matching_team_member(member, complete_team_data) is not None
|
||||
|
||||
|
||||
def _already_in_team_message(member: Member, existing_member: Member) -> str:
|
||||
if member.user_id is not None and existing_member.user_id == member.user_id:
|
||||
return f"{member.user_id} is already a member of this team."
|
||||
return f"A member of this team already uses {member.user_email}."
|
||||
|
||||
|
||||
async def _update_team_members_list(
|
||||
|
|
@ -3131,9 +3157,7 @@ async def _update_team_members_list(
|
|||
updated_users: list[LiteLLM_UserTable],
|
||||
) -> None:
|
||||
"""Update the team's members_with_roles list."""
|
||||
requested_members: Final[Sequence[Member]] = (
|
||||
(data.member,) if isinstance(data.member, Member) else tuple(data.member)
|
||||
)
|
||||
requested_members: Final = _requested_members(data)
|
||||
resolved_members: Final = tuple(_resolve_member_identity(m, updated_users) for m in requested_members)
|
||||
|
||||
# extend() consumes the generator as it appends, so a member already added by this
|
||||
|
|
@ -3437,7 +3461,8 @@ async def _validate_and_populate_member_user_info(
|
|||
Validate and populate user_email/user_id for a member.
|
||||
|
||||
Logic:
|
||||
1. If both user_email and user_id are provided, verify they belong to the same user (use user_email as source of truth)
|
||||
1. If both user_email and user_id are provided, the user_id decides: an existing user with an email on record
|
||||
must have that email, and any other user_id must not take an email another user already has
|
||||
2. If only user_email is provided, populate user_id from DB
|
||||
3. If only user_id is provided, populate user_email from DB (if user exists)
|
||||
4. If only user_id is provided and doesn't exist, allow it to pass with user_email as None (will be upserted later)
|
||||
|
|
@ -3451,41 +3476,23 @@ async def _validate_and_populate_member_user_info(
|
|||
detail={"error": "Either user_id or user_email must be provided"},
|
||||
)
|
||||
|
||||
# Case 1: Both user_email and user_id provided - verify they match
|
||||
if member.user_email is not None and member.user_id is not None:
|
||||
# Use user_email as source of truth
|
||||
# Check for multiple users with same email first
|
||||
users_by_email = await prisma_client.get_data(
|
||||
key_val={"user_email": member.user_email},
|
||||
table_name="user",
|
||||
query_type="find_all",
|
||||
mismatch_error: Final = HTTPException(
|
||||
status_code=400,
|
||||
detail={
|
||||
"error": (
|
||||
f"user_email '{member.user_email}' and user_id '{member.user_id}' do not belong to the same user."
|
||||
)
|
||||
},
|
||||
)
|
||||
|
||||
if users_by_email is None or (isinstance(users_by_email, list) and len(users_by_email) == 0):
|
||||
# User doesn't exist yet - this is fine, will be created later
|
||||
existing_user: Final = await _user_db(prisma_client).find_unique(where={"user_id": member.user_id})
|
||||
if existing_user is not None and existing_user.user_email is not None:
|
||||
if existing_user.user_email != member.user_email:
|
||||
raise mismatch_error
|
||||
return member
|
||||
|
||||
if isinstance(users_by_email, list) and len(users_by_email) > 1:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail={
|
||||
"error": f"Multiple users found with email '{member.user_email}'. Please use 'user_id' instead."
|
||||
},
|
||||
)
|
||||
|
||||
# Get the single user
|
||||
user_by_email = users_by_email[0]
|
||||
|
||||
# Verify the user_id matches
|
||||
if user_by_email.user_id != member.user_id:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail={
|
||||
"error": f"user_email '{member.user_email}' and user_id '{member.user_id}' do not belong to the same user."
|
||||
},
|
||||
)
|
||||
|
||||
# Both match, return as is
|
||||
if await _user_db(prisma_client).find_first(where={"user_email": member.user_email}) is not None:
|
||||
raise mismatch_error
|
||||
return member
|
||||
|
||||
# Case 2: Only user_email provided - populate user_id from DB
|
||||
|
|
@ -3598,19 +3605,20 @@ async def team_member_add(
|
|||
|
||||
complete_team_data: Final = LiteLLM_TeamTable.model_validate(existing_team_row.model_dump())
|
||||
|
||||
team_member_add_duplication_check(
|
||||
data=data,
|
||||
existing_team_row=complete_team_data,
|
||||
)
|
||||
|
||||
# Validate permissions
|
||||
await _validate_team_member_add_permissions(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
complete_team_data=complete_team_data,
|
||||
data=data,
|
||||
)
|
||||
|
||||
requested_members: Final = tuple(data.member) if isinstance(data.member, list) else (data.member,)
|
||||
requested_members: Final = _requested_members(data)
|
||||
_validate_member_request_shape(requested_members)
|
||||
if isinstance(data.member, Member):
|
||||
team_member_add_duplication_check(
|
||||
data=data,
|
||||
existing_team_row=complete_team_data,
|
||||
)
|
||||
|
||||
caller_supplied_user_ids = frozenset(member.user_id for member in requested_members if member.user_id is not None)
|
||||
existing_user_ids: Final = await _resolve_existing_member_user_ids(
|
||||
members=requested_members,
|
||||
|
|
@ -3636,6 +3644,12 @@ async def team_member_add(
|
|||
prisma_client=prisma_client,
|
||||
)
|
||||
|
||||
_validate_resolved_members_are_distinct(requested_members)
|
||||
team_member_add_duplication_check(
|
||||
data=data,
|
||||
existing_team_row=complete_team_data,
|
||||
)
|
||||
|
||||
pre_existing_user_ids: Final = _pre_existing_user_ids(
|
||||
members=requested_members,
|
||||
caller_supplied_user_ids=caller_supplied_user_ids,
|
||||
|
|
|
|||
|
|
@ -1180,8 +1180,68 @@ def test_team_member_add_duplication_check_raises_proxy_exception():
|
|||
assert exc_info.value.type == ProxyErrorTypes.team_member_already_in_team
|
||||
assert exc_info.value.param == "member"
|
||||
assert exc_info.value.code == "400"
|
||||
assert "existing-user-id" in str(exc_info.value.message)
|
||||
assert "already in team" in str(exc_info.value.message)
|
||||
assert exc_info.value.message == "existing-user-id is already a member of this team."
|
||||
|
||||
|
||||
def _team_with_roster(*members: Member) -> LiteLLM_TeamTable:
|
||||
return LiteLLM_TeamTable(team_id="test-team-123", members_with_roles=list(members))
|
||||
|
||||
|
||||
_DUP_USER_1: Final = Member(user_id="dup-user-1", user_email="shared@example.com", role="user")
|
||||
_DUP_USER_2: Final = Member(user_id="dup-user-2", user_email="shared@example.com", role="user")
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("roster", "member", "expected_message"),
|
||||
[
|
||||
((_DUP_USER_1,), _DUP_USER_1, "dup-user-1 is already a member of this team."),
|
||||
((_DUP_USER_1,), Member(user_id="dup-user-1", role="user"), "dup-user-1 is already a member of this team."),
|
||||
(
|
||||
(_DUP_USER_1,),
|
||||
Member(user_email="shared@example.com", role="user"),
|
||||
"A member of this team already uses shared@example.com.",
|
||||
),
|
||||
((_DUP_USER_1,), _DUP_USER_2, "A member of this team already uses shared@example.com."),
|
||||
((_DUP_USER_1, _DUP_USER_2), _DUP_USER_2, "dup-user-2 is already a member of this team."),
|
||||
((_DUP_USER_2, _DUP_USER_1), _DUP_USER_2, "dup-user-2 is already a member of this team."),
|
||||
(
|
||||
(Member(user_email="invited@example.com", role="user"),),
|
||||
Member(user_email="invited@example.com", role="user"),
|
||||
"A member of this team already uses invited@example.com.",
|
||||
),
|
||||
],
|
||||
ids=[
|
||||
"id_and_email",
|
||||
"id_only",
|
||||
"email_only",
|
||||
"other_user_same_email",
|
||||
"same_id_after_shared_email_member",
|
||||
"same_id_before_shared_email_member",
|
||||
"email_only_roster_entry",
|
||||
],
|
||||
)
|
||||
def test_team_member_add_duplication_check_names_only_the_requested_member(
|
||||
roster: tuple[Member, ...], member: Member, expected_message: str
|
||||
):
|
||||
data: Final = TeamMemberAddRequest(team_id="test-team-123", member=member)
|
||||
|
||||
with pytest.raises(ProxyException) as exc_info:
|
||||
team_member_add_duplication_check(data=data, existing_team_row=_team_with_roster(*roster))
|
||||
|
||||
assert exc_info.value.type == ProxyErrorTypes.team_member_already_in_team
|
||||
assert exc_info.value.message == expected_message
|
||||
|
||||
|
||||
def test_team_member_add_duplication_check_bulk_message_omits_roster():
|
||||
data: Final = TeamMemberAddRequest(
|
||||
team_id="test-team-123",
|
||||
member=[Member(user_id="dup-user-1", role="user"), Member(user_email="shared@example.com", role="user")],
|
||||
)
|
||||
|
||||
with pytest.raises(ProxyException) as exc_info:
|
||||
team_member_add_duplication_check(data=data, existing_team_row=_team_with_roster(_DUP_USER_1))
|
||||
|
||||
assert exc_info.value.message == "All requested users are already members of this team."
|
||||
|
||||
|
||||
def test_team_member_add_duplication_check_allows_new_member():
|
||||
|
|
@ -10920,42 +10980,93 @@ async def test_get_team_daily_activity_team_admin_sees_all_spend(mock_db_client)
|
|||
pytest.fail("API keys should not be fetched for team admin users")
|
||||
|
||||
|
||||
def _first_email_owner(rows: tuple[MagicMock, ...], where: dict[str, object]) -> MagicMock | None:
|
||||
criterion: Final = where["user_email"]
|
||||
if isinstance(criterion, dict):
|
||||
return next((row for row in rows if (row.user_email or "").lower() == criterion["equals"].lower()), None)
|
||||
return next((row for row in rows if row.user_email == criterion), None)
|
||||
|
||||
|
||||
def _prisma_with_users(*users: tuple[str, str | None], team: LiteLLM_TeamTable | None = None) -> MagicMock:
|
||||
rows: Final = tuple(MagicMock(user_id=user_id, user_email=user_email) for user_id, user_email in users)
|
||||
return MagicMock(
|
||||
**{
|
||||
"db.litellm_usertable.find_unique": AsyncMock(
|
||||
side_effect=lambda where: next((row for row in rows if row.user_id == where["user_id"]), None)
|
||||
),
|
||||
"db.litellm_usertable.find_first": AsyncMock(side_effect=lambda where: _first_email_owner(rows, where)),
|
||||
"db.litellm_usertable.find_many": AsyncMock(
|
||||
side_effect=lambda where: [row for row in rows if row.user_id in where["user_id"]["in"]]
|
||||
),
|
||||
"get_data": AsyncMock(
|
||||
side_effect=lambda key_val, table_name, query_type: [
|
||||
row for row in rows if row.user_email == key_val["user_email"]
|
||||
]
|
||||
),
|
||||
"writer_db.litellm_teamtable.find_unique": AsyncMock(return_value=team),
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
_USER_DIRECTORY: Final = (
|
||||
("dup-user-1", "shared@example.com"),
|
||||
("dup-user-2", "shared@example.com"),
|
||||
("solo-user", "solo@example.com"),
|
||||
("no-email-user", None),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_validate_and_populate_member_user_info_both_provided_match():
|
||||
"""
|
||||
Test _validate_and_populate_member_user_info when both user_email and user_id
|
||||
are provided and they match the same user in the database.
|
||||
"""
|
||||
# Create member with both user_email and user_id
|
||||
member = Member(user_email="test@example.com", user_id="user-123", role="user")
|
||||
@pytest.mark.parametrize(
|
||||
("user_id", "user_email"),
|
||||
[
|
||||
("dup-user-2", "shared@example.com"),
|
||||
("solo-user", "solo@example.com"),
|
||||
("brand-new-user", "unused@example.com"),
|
||||
("no-email-user", "first-email@example.com"),
|
||||
],
|
||||
ids=["one_of_shared_email_owners", "single_owner", "new_user_with_unused_email", "existing_user_without_email"],
|
||||
)
|
||||
async def test_validate_and_populate_member_user_info_accepts_owner_of_the_email(user_id: str, user_email: str):
|
||||
member: Final = Member(user_id=user_id, user_email=user_email, role="user")
|
||||
|
||||
# Mock prisma client
|
||||
mock_prisma_client = MagicMock()
|
||||
|
||||
# Mock user object that matches both email and user_id
|
||||
mock_user = MagicMock()
|
||||
mock_user.user_id = "user-123"
|
||||
mock_user.user_email = "test@example.com"
|
||||
|
||||
# Mock get_data to return single user matching email
|
||||
mock_prisma_client.get_data = AsyncMock(return_value=[mock_user])
|
||||
|
||||
# Call the function
|
||||
result = await _validate_and_populate_member_user_info(
|
||||
member=member,
|
||||
prisma_client=mock_prisma_client,
|
||||
result: Final = await _validate_and_populate_member_user_info(
|
||||
member=member, prisma_client=_prisma_with_users(*_USER_DIRECTORY)
|
||||
)
|
||||
|
||||
# Verify result matches input (both already provided and match)
|
||||
assert result.user_email == "test@example.com"
|
||||
assert result.user_id == "user-123"
|
||||
assert (result.user_id, result.user_email) == (user_id, user_email)
|
||||
|
||||
# Verify get_data was called with correct parameters
|
||||
mock_prisma_client.get_data.assert_called_once_with(
|
||||
key_val={"user_email": "test@example.com"},
|
||||
table_name="user",
|
||||
query_type="find_all",
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("user_id", "user_email"),
|
||||
[
|
||||
("brand-new-user", "shared@example.com"),
|
||||
("solo-user", "shared@example.com"),
|
||||
("dup-user-1", "unused@example.com"),
|
||||
("brand-new-user", "solo@example.com"),
|
||||
("no-email-user", "solo@example.com"),
|
||||
],
|
||||
ids=[
|
||||
"new_user_id_with_shared_email",
|
||||
"existing_user_id_with_shared_email",
|
||||
"existing_user_id_with_unused_email",
|
||||
"new_user_id_with_taken_email",
|
||||
"existing_user_without_email_with_taken_email",
|
||||
],
|
||||
)
|
||||
async def test_validate_and_populate_member_user_info_rejects_non_owner_of_the_email(
|
||||
user_id: str, user_email: str
|
||||
):
|
||||
member: Final = Member(user_id=user_id, user_email=user_email, role="user")
|
||||
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await _validate_and_populate_member_user_info(member=member, prisma_client=_prisma_with_users(*_USER_DIRECTORY))
|
||||
|
||||
assert exc_info.value.status_code == 400
|
||||
assert exc_info.value.detail == {
|
||||
"error": f"user_email '{user_email}' and user_id '{user_id}' do not belong to the same user."
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -14296,6 +14407,201 @@ async def test_team_member_add_evicts_the_cached_team_roster(monkeypatch):
|
|||
assert cache.get_cache(key="team_alias:roster-evict") is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_team_member_add_rejects_non_admin_before_revealing_existing_members(monkeypatch):
|
||||
from litellm.proxy._types import TeamMemberAddRequest
|
||||
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
|
||||
from litellm.proxy.management_endpoints.team_endpoints import team_member_add
|
||||
|
||||
team_row = LiteLLM_TeamTable(
|
||||
team_id="team-private",
|
||||
members_with_roles=[Member(user_id="dup-user-1", user_email="shared@example.com", role="user")],
|
||||
)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", AsyncMock())
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", UserApiKeyCache())
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", None)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.litellm_proxy_admin_name", "default_user_id")
|
||||
|
||||
with (
|
||||
patch(
|
||||
"litellm.proxy.management_endpoints.team_endpoints.get_team_object",
|
||||
new_callable=AsyncMock,
|
||||
return_value=team_row,
|
||||
),
|
||||
pytest.raises(HTTPException) as exc_info,
|
||||
):
|
||||
await team_member_add(
|
||||
data=TeamMemberAddRequest(
|
||||
team_id="team-private", member=Member(user_email="shared@example.com", role="user")
|
||||
),
|
||||
user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, user_id="outsider"),
|
||||
)
|
||||
|
||||
assert exc_info.value.status_code == 403
|
||||
assert "dup-user-1" not in str(exc_info.value.detail)
|
||||
|
||||
|
||||
_WRITE_METHODS: Final = frozenset({"create", "create_many", "upsert", "update", "update_many", "delete", "tx"})
|
||||
_MEMBER_ADD_USERS: Final = (
|
||||
("dup-user-1", "shared@example.com"),
|
||||
("dup-user-2", "shared@example.com"),
|
||||
("dup-user-3", "shared@example.com"),
|
||||
("solo-user", "solo@example.com"),
|
||||
)
|
||||
|
||||
|
||||
def _write_calls(prisma: MagicMock) -> tuple[str, ...]:
|
||||
return tuple(name for name, _, _ in prisma.mock_calls if name.rsplit(".", 1)[-1] in _WRITE_METHODS)
|
||||
|
||||
|
||||
def _rejection(error: Exception) -> tuple[str, int | str, str]:
|
||||
if isinstance(error, ProxyException):
|
||||
assert (error.type, error.param) == (ProxyErrorTypes.team_member_already_in_team, "member"), error
|
||||
return ("already_in_team", error.code, error.message)
|
||||
assert isinstance(error, HTTPException), error
|
||||
return ("http", error.status_code, error.detail["error"])
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("roster", "member", "expected"),
|
||||
[
|
||||
(
|
||||
(_DUP_USER_1,),
|
||||
Member(user_id="dup-user-2", role="user"),
|
||||
("already_in_team", "400", "A member of this team already uses shared@example.com."),
|
||||
),
|
||||
(
|
||||
(_DUP_USER_1, _DUP_USER_2),
|
||||
_DUP_USER_2,
|
||||
("already_in_team", "400", "dup-user-2 is already a member of this team."),
|
||||
),
|
||||
(
|
||||
(),
|
||||
Member(user_email="shared@example.com", role="user"),
|
||||
("http", 400, "Multiple users found with email 'shared@example.com'. Please use 'user_id' instead."),
|
||||
),
|
||||
(
|
||||
(),
|
||||
Member(user_id="solo-user", user_email="shared@example.com", role="user"),
|
||||
("http", 400, "user_email 'shared@example.com' and user_id 'solo-user' do not belong to the same user."),
|
||||
),
|
||||
((), [], ("http", 400, "At least one member is required.")),
|
||||
(
|
||||
(),
|
||||
[_DUP_USER_2, Member(user_id="dup-user-3", user_email="shared@example.com", role="user")],
|
||||
("http", 400, "Several requested users share the email shared@example.com. Add only one of them."),
|
||||
),
|
||||
(
|
||||
(),
|
||||
[Member(user_id="dup-user-2", role="user"), Member(user_id="dup-user-3", role="user")],
|
||||
("http", 400, "Several requested users share the email shared@example.com. Add only one of them."),
|
||||
),
|
||||
(
|
||||
(),
|
||||
[Member(user_id="solo-user", role="user"), Member(user_id="solo-user", role="admin")],
|
||||
("http", 400, "solo-user appears more than once in this request."),
|
||||
),
|
||||
(
|
||||
(),
|
||||
[
|
||||
Member(user_id="solo-user", user_email="solo@example.com", role="user"),
|
||||
Member(user_email="solo@example.com", role="user"),
|
||||
],
|
||||
("http", 400, "solo-user appears more than once in this request."),
|
||||
),
|
||||
(
|
||||
(Member(user_id="solo-user", user_email="solo@example.com", role="user"),),
|
||||
[Member(user_id="solo-user", role="user"), Member(user_email="solo@example.com", role="user")],
|
||||
("http", 400, "solo-user appears more than once in this request."),
|
||||
),
|
||||
(
|
||||
(Member(user_id="solo-user", role="user"),),
|
||||
[Member(user_email="solo@example.com", role="user")],
|
||||
("already_in_team", "400", "All requested users are already members of this team."),
|
||||
),
|
||||
],
|
||||
ids=[
|
||||
"row11_user_id_whose_email_a_member_uses",
|
||||
"row8_same_user_already_in_team_after_shared_email_member",
|
||||
"row4_email_with_several_owners",
|
||||
"row6_user_id_that_does_not_own_the_email",
|
||||
"row15_empty_list",
|
||||
"row18_two_users_sharing_an_email",
|
||||
"row18_two_user_ids_that_resolve_to_one_email",
|
||||
"row19_same_user_id_twice",
|
||||
"row19_same_user_by_id_and_email_then_email",
|
||||
"row19_member_already_in_team_by_id_and_by_email",
|
||||
"row16_all_members_already_in_team_after_resolution",
|
||||
],
|
||||
)
|
||||
async def test_team_member_add_rejects_without_writing(
|
||||
monkeypatch,
|
||||
roster: tuple[Member, ...],
|
||||
member: Member | list[Member],
|
||||
expected: tuple[str, int | str, str],
|
||||
):
|
||||
from litellm.proxy._types import TeamMemberAddRequest
|
||||
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
|
||||
from litellm.proxy.management_endpoints.team_endpoints import team_member_add
|
||||
|
||||
team_id: Final = f"team-{uuid.uuid4().hex}"
|
||||
prisma: Final = _prisma_with_users(
|
||||
*_MEMBER_ADD_USERS, team=LiteLLM_TeamTable(team_id=team_id, members_with_roles=list(roster))
|
||||
)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", prisma)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", UserApiKeyCache())
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", None)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.litellm_proxy_admin_name", "default_user_id")
|
||||
|
||||
with pytest.raises((HTTPException, ProxyException)) as exc_info:
|
||||
await team_member_add(
|
||||
data=TeamMemberAddRequest(team_id=team_id, member=member),
|
||||
user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin-1"),
|
||||
)
|
||||
|
||||
assert _rejection(exc_info.value) == expected
|
||||
assert _write_calls(prisma) == ()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_team_member_add_self_join_cannot_claim_another_users_email(monkeypatch):
|
||||
from litellm.proxy._types import TeamMemberAddRequest
|
||||
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
|
||||
from litellm.proxy.management_endpoints.team_endpoints import team_member_add
|
||||
|
||||
team_id: Final = f"team-{uuid.uuid4().hex}"
|
||||
prisma: Final = _prisma_with_users(
|
||||
("no-email-user", None),
|
||||
("victim-user", "victim@example.com"),
|
||||
team=LiteLLM_TeamTable(team_id=team_id, members_with_roles=[]),
|
||||
)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", prisma)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", UserApiKeyCache())
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", None)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.litellm_proxy_admin_name", "default_user_id")
|
||||
monkeypatch.setattr("litellm.default_internal_user_params", {"available_teams": [team_id]})
|
||||
|
||||
with pytest.raises((HTTPException, ProxyException)) as exc_info:
|
||||
await team_member_add(
|
||||
data=TeamMemberAddRequest(
|
||||
team_id=team_id,
|
||||
member=Member(user_id="no-email-user", user_email="victim@example.com", role="user"),
|
||||
),
|
||||
user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, user_id="no-email-user"),
|
||||
)
|
||||
|
||||
assert _rejection(exc_info.value) == (
|
||||
"http",
|
||||
400,
|
||||
"user_email 'victim@example.com' and user_id 'no-email-user' do not belong to the same user.",
|
||||
)
|
||||
assert _write_calls(prisma) == ()
|
||||
|
||||
|
||||
class _RecordingAuditLogger(CustomLogger):
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
|
|
|
|||
|
|
@ -205,6 +205,70 @@ describe("UserSearchModal submit payload", () => {
|
|||
});
|
||||
});
|
||||
|
||||
describe("UserSearchModal users sharing an email", () => {
|
||||
beforeEach(() => {
|
||||
vi.mocked(userFilterUICall).mockReset();
|
||||
vi.mocked(userFilterUICall).mockResolvedValue([
|
||||
{ user_id: "u-first", user_email: "shared@example.com" },
|
||||
{ user_id: "u-second", user_email: "shared@example.com" },
|
||||
] as never);
|
||||
});
|
||||
|
||||
it("lists the shared email once and submits the user ID picked from its owners", async () => {
|
||||
const user = userEvent.setup();
|
||||
const onSubmit = vi.fn();
|
||||
render(<UserSearchModal isVisible onCancel={vi.fn()} onSubmit={onSubmit} accessToken="sk-test" />);
|
||||
const save = screen.getByRole("button", { name: /add member/i });
|
||||
|
||||
const emailInput = getEmailSearchInput();
|
||||
await user.click(emailInput);
|
||||
await user.type(emailInput, "shared");
|
||||
await waitFor(() => expect(userFilterUICall).toHaveBeenCalled(), { timeout: 3000 });
|
||||
const emailOptions = await screen.findAllByRole("option");
|
||||
expect(emailOptions.map((option) => option.textContent)).toEqual(["shared@example.com"]);
|
||||
await user.click(emailOptions[0]);
|
||||
|
||||
expect(screen.getByLabelText("User ID")).toHaveValue("");
|
||||
expect(screen.getByText(/multiple users share this email/i)).toBeInTheDocument();
|
||||
expect(save).toBeDisabled();
|
||||
|
||||
await user.click(screen.getByLabelText("User ID"));
|
||||
const idOptions = await screen.findAllByRole("option");
|
||||
expect(idOptions.map((option) => option.textContent)).toEqual(["u-first", "u-second"]);
|
||||
await user.click(screen.getByRole("option", { name: "u-second" }));
|
||||
await user.click(save);
|
||||
|
||||
await waitFor(() => expect(onSubmit).toHaveBeenCalledTimes(1));
|
||||
expect(onSubmit.mock.calls[0][0]).toStrictEqual({
|
||||
user_email: "shared@example.com",
|
||||
user_id: "u-second",
|
||||
role: "user",
|
||||
});
|
||||
expect(userFilterUICall).toHaveBeenCalledTimes(1);
|
||||
});
|
||||
|
||||
it("lists the shared email's user IDs again when the ID dropdown reopens after a search", async () => {
|
||||
const user = userEvent.setup();
|
||||
render(<UserSearchModal isVisible onCancel={vi.fn()} onSubmit={vi.fn()} accessToken="sk-test" />);
|
||||
|
||||
const emailInput = getEmailSearchInput();
|
||||
await user.click(emailInput);
|
||||
await user.type(emailInput, "shared");
|
||||
await user.click((await screen.findAllByRole("option"))[0]);
|
||||
|
||||
vi.mocked(userFilterUICall).mockResolvedValue([] as never);
|
||||
const idInput = screen.getByLabelText("User ID");
|
||||
await user.click(idInput);
|
||||
await user.type(idInput, "zzz");
|
||||
await waitFor(() => expect(userFilterUICall).toHaveBeenCalledTimes(2), { timeout: 3000 });
|
||||
await user.keyboard("{Escape}");
|
||||
await user.click(idInput);
|
||||
|
||||
const idOptions = await screen.findAllByRole("option", {}, { timeout: 3000 });
|
||||
expect(idOptions.map((option) => option.textContent)).toEqual(["u-first", "u-second"]);
|
||||
});
|
||||
});
|
||||
|
||||
describe("UserSearchModal search lifecycle", () => {
|
||||
const directory = [
|
||||
{ user_id: "u-jones", user_email: "alice.jones@example.com" },
|
||||
|
|
|
|||
|
|
@ -18,12 +18,21 @@ interface User {
|
|||
role?: string;
|
||||
}
|
||||
|
||||
type SearchField = "user_email" | "user_id";
|
||||
|
||||
interface UserOption {
|
||||
label: string;
|
||||
value: string;
|
||||
user: User | null;
|
||||
users: User[];
|
||||
}
|
||||
|
||||
const groupUsersByField = (users: User[], fieldName: SearchField): UserOption[] =>
|
||||
[...new Set(users.map((user) => `${user[fieldName]}`))].map((value) => ({
|
||||
label: value,
|
||||
value,
|
||||
users: users.filter((user) => `${user[fieldName]}` === value),
|
||||
}));
|
||||
|
||||
interface Role {
|
||||
label: string;
|
||||
value: string;
|
||||
|
|
@ -70,11 +79,15 @@ const UserSearchModal: React.FC<UserSearchModalProps> = ({
|
|||
const selectedUserEmail = form.watch("user_email");
|
||||
const [userOptions, setUserOptions] = useState<UserOption[]>([]);
|
||||
const [loading, setLoading] = useState<boolean>(false);
|
||||
const [selectedField, setSelectedField] = useState<"user_email" | "user_id">("user_email");
|
||||
const [selectedField, setSelectedField] = useState<SearchField>("user_email");
|
||||
const [searchQuery, setSearchQuery] = useState("");
|
||||
const [usersSharingEmail, setUsersSharingEmail] = useState<User[]>([]);
|
||||
const [isSubmitting, setIsSubmitting] = useState(false);
|
||||
const latestSearchRef = useRef(0);
|
||||
const needsUserIdPick = usersSharingEmail.length > 0 && !selectedUserId;
|
||||
const hasPickedUser = Boolean(selectedUserId || selectedUserEmail) && !needsUserIdPick;
|
||||
|
||||
const fetchUsers = async (searchText: string, fieldName: "user_email" | "user_id"): Promise<void> => {
|
||||
const fetchUsers = async (searchText: string, fieldName: SearchField): Promise<void> => {
|
||||
const searchId = latestSearchRef.current + 1;
|
||||
latestSearchRef.current = searchId;
|
||||
const isLatestSearch = (): boolean => searchId === latestSearchRef.current;
|
||||
|
|
@ -99,12 +112,7 @@ const UserSearchModal: React.FC<UserSearchModalProps> = ({
|
|||
if (!isLatestSearch()) return;
|
||||
|
||||
const data: User[] = response;
|
||||
const options: UserOption[] = data.map((user) => ({
|
||||
label: fieldName === "user_email" ? `${user.user_email}` : `${user.user_id}`,
|
||||
value: fieldName === "user_email" ? user.user_email : user.user_id,
|
||||
user,
|
||||
}));
|
||||
setUserOptions(options);
|
||||
setUserOptions(groupUsersByField(data, fieldName));
|
||||
} catch (error) {
|
||||
console.error("Error fetching users:", error);
|
||||
} finally {
|
||||
|
|
@ -112,15 +120,23 @@ const UserSearchModal: React.FC<UserSearchModalProps> = ({
|
|||
}
|
||||
};
|
||||
|
||||
const handleSearch = (value: string, fieldName: "user_email" | "user_id"): void => {
|
||||
const handleSearch = (value: string, fieldName: SearchField): void => {
|
||||
setSelectedField(fieldName);
|
||||
setSearchQuery(value);
|
||||
void fetchUsers(value, fieldName);
|
||||
};
|
||||
|
||||
const handleSelect = (option: UserOption | null): void => {
|
||||
if (option?.user == null) return;
|
||||
form.setValue("user_email", option.user.user_email);
|
||||
form.setValue("user_id", option.user.user_id);
|
||||
const handleSelect = (option: UserOption | null, fieldName: SearchField): void => {
|
||||
const [firstUser, ...otherUsers] = option?.users ?? [];
|
||||
if (firstUser === undefined) return;
|
||||
const isAmbiguous = otherUsers.length > 0;
|
||||
form.setValue("user_email", firstUser.user_email);
|
||||
form.setValue("user_id", isAmbiguous ? null : firstUser.user_id);
|
||||
if (fieldName === "user_email") {
|
||||
setUsersSharingEmail(isAmbiguous ? option?.users ?? [] : []);
|
||||
return;
|
||||
}
|
||||
setUsersSharingEmail((users) => (users.some((user) => user.user_id === firstUser.user_id) ? users : []));
|
||||
};
|
||||
|
||||
const handleSubmit = async (values: FormValues): Promise<void> => {
|
||||
|
|
@ -135,6 +151,8 @@ const UserSearchModal: React.FC<UserSearchModalProps> = ({
|
|||
const handleClose = (): void => {
|
||||
form.reset(emptyValues);
|
||||
setUserOptions([]);
|
||||
setSearchQuery("");
|
||||
setUsersSharingEmail([]);
|
||||
onCancel();
|
||||
};
|
||||
|
||||
|
|
@ -143,7 +161,7 @@ const UserSearchModal: React.FC<UserSearchModalProps> = ({
|
|||
};
|
||||
|
||||
const renderUserSearch = (
|
||||
fieldName: "user_email" | "user_id",
|
||||
fieldName: SearchField,
|
||||
placeholder: string,
|
||||
controlProps: {
|
||||
id: string;
|
||||
|
|
@ -152,15 +170,17 @@ const UserSearchModal: React.FC<UserSearchModalProps> = ({
|
|||
},
|
||||
testId?: string,
|
||||
) => {
|
||||
const items = selectedField === fieldName ? userOptions : [];
|
||||
const sharedEmailOptions = fieldName === "user_id" ? groupUsersByField(usersSharingEmail, "user_id") : [];
|
||||
const items = selectedField === fieldName && searchQuery !== "" ? userOptions : sharedEmailOptions;
|
||||
const handleValueChange = (value: string | null) => {
|
||||
if (value === null) {
|
||||
form.setValue("user_email", null);
|
||||
form.setValue("user_id", null);
|
||||
setUsersSharingEmail([]);
|
||||
return;
|
||||
}
|
||||
controlProps.onChange(value);
|
||||
handleSelect(items.find((option) => option.value === value) ?? null);
|
||||
handleSelect(items.find((option) => option.value === value) ?? null, fieldName);
|
||||
};
|
||||
return (
|
||||
<div data-testid={testId} onKeyDown={swallowEnter}>
|
||||
|
|
@ -205,7 +225,12 @@ const UserSearchModal: React.FC<UserSearchModalProps> = ({
|
|||
|
||||
<div className="text-center">OR</div>
|
||||
|
||||
<FormField control={form.control} name="user_id" label="User ID">
|
||||
<FormField
|
||||
control={form.control}
|
||||
name="user_id"
|
||||
label="User ID"
|
||||
description={needsUserIdPick ? "Multiple users share this email. Pick the user ID to add." : undefined}
|
||||
>
|
||||
{({ id, value, onChange }) => renderUserSearch("user_id", "Search by user ID", { id, value, onChange })}
|
||||
</FormField>
|
||||
|
||||
|
|
@ -238,7 +263,7 @@ const UserSearchModal: React.FC<UserSearchModalProps> = ({
|
|||
</FieldGroup>
|
||||
|
||||
<div className="mt-4 text-right">
|
||||
<Button type="submit" disabled={isSubmitting || (!selectedUserId && !selectedUserEmail)}>
|
||||
<Button type="submit" disabled={isSubmitting || !hasPickedUser}>
|
||||
{isSubmitting ? <UiLoadingSpinner className="size-4" /> : <UserPlus />}
|
||||
{isSubmitting ? "Adding..." : "Add Member"}
|
||||
</Button>
|
||||
|
|
|
|||
|
|
@ -118,6 +118,50 @@ describe("loginCall - storeLoginToken integration", () => {
|
|||
});
|
||||
});
|
||||
|
||||
describe("teamMemberAddCall error message", () => {
|
||||
afterEach(() => {
|
||||
vi.restoreAllMocks();
|
||||
});
|
||||
|
||||
const rejectWith = (body: string) => {
|
||||
vi.spyOn(global, "fetch").mockResolvedValue(new Response(body, { status: 400 }));
|
||||
};
|
||||
|
||||
const member: Networking.Member = { user_id: "dup-user-1", role: "user" };
|
||||
const addMember = () => Networking.teamMemberAddCall("token", "team-1", member);
|
||||
|
||||
it.each([
|
||||
["the proxy error envelope", { error: { message: "dup-user-1 is already a member of this team." } }],
|
||||
["an HTTPException detail", { detail: { error: "dup-user-1 is already a member of this team." } }],
|
||||
["a plain message field", { message: "dup-user-1 is already a member of this team." }],
|
||||
["a plain error string", { error: "dup-user-1 is already a member of this team." }],
|
||||
["a plain detail string", { detail: "dup-user-1 is already a member of this team." }],
|
||||
[
|
||||
"a lower-priority field when a higher one is not a string",
|
||||
{ error: { message: 5 }, detail: { error: "dup-user-1 is already a member of this team." } },
|
||||
],
|
||||
])("surfaces the backend message from %s", async (_shape, body) => {
|
||||
rejectWith(JSON.stringify(body));
|
||||
|
||||
await expect(addMember()).rejects.toHaveProperty("message", "dup-user-1 is already a member of this team.");
|
||||
});
|
||||
|
||||
it.each([
|
||||
["is not JSON", "<html>bad gateway</html>"],
|
||||
["is JSON null", "null"],
|
||||
["is an empty object", "{}"],
|
||||
["is an object without a known message field", JSON.stringify({ status: 502 })],
|
||||
["is an array", JSON.stringify(["unknown"])],
|
||||
["has a non-string message", JSON.stringify({ message: { reason: "unknown" } })],
|
||||
["has an empty error message", JSON.stringify({ error: { message: "" } })],
|
||||
["has a whitespace-only message", JSON.stringify({ message: " " })],
|
||||
])("falls back to a generic message when the body %s", async (_case, body) => {
|
||||
rejectWith(body);
|
||||
|
||||
await expect(addMember()).rejects.toHaveProperty("message", "Failed to add team member");
|
||||
});
|
||||
});
|
||||
|
||||
describe("exchangeLoginCode - storeLoginToken integration", () => {
|
||||
const originalFetch = global.fetch;
|
||||
|
||||
|
|
|
|||
|
|
@ -2894,6 +2894,22 @@ export interface Member {
|
|||
model_max_budget?: StoredModelMaxBudget | null;
|
||||
}
|
||||
|
||||
const nonEmptyString = (value: unknown): string | undefined =>
|
||||
typeof value === "string" && value.trim() !== "" ? value : undefined;
|
||||
|
||||
const objectField = (value: unknown, key: string): unknown =>
|
||||
value !== null && typeof value === "object" && !Array.isArray(value)
|
||||
? (value as Record<string, unknown>)[key]
|
||||
: undefined;
|
||||
|
||||
const teamMemberAddErrorMessage = (body: unknown): string =>
|
||||
nonEmptyString(objectField(objectField(body, "error"), "message")) ??
|
||||
nonEmptyString(objectField(body, "error")) ??
|
||||
nonEmptyString(objectField(body, "message")) ??
|
||||
nonEmptyString(objectField(body, "detail")) ??
|
||||
nonEmptyString(objectField(objectField(body, "detail"), "error")) ??
|
||||
"Failed to add team member";
|
||||
|
||||
export const teamMemberAddCall = async (accessToken: string, teamId: string, formValues: Member) => {
|
||||
try {
|
||||
const url = proxyBaseUrl ? `${proxyBaseUrl}/team/member_add` : `/team/member_add`;
|
||||
|
|
@ -2921,8 +2937,7 @@ export const teamMemberAddCall = async (accessToken: string, teamId: string, for
|
|||
console.warn("Failed to parse error body as JSON:", errorText);
|
||||
}
|
||||
|
||||
const rawMessage = parsedError?.detail?.error || "Failed to add team member";
|
||||
const err = new Error(rawMessage);
|
||||
const err = new Error(teamMemberAddErrorMessage(parsedError));
|
||||
(err as any).raw = parsedError;
|
||||
throw err;
|
||||
}
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue