scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-08 03:08:45 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 115

Fix user-team association issues in LiteLLM proxy (#12082)

Browse source 
* Fix user-team association issues in LiteLLM proxy

- Update list_team function to properly filter teams using user's teams array instead of only checking members_with_roles field
- Add Field descriptions and docstring to TeamMemberAddRequest and related models for better Swagger/OpenAPI documentation
- Maintain backward compatibility with fallback to members_with_roles if user lookup fails

This ensures users created with teams parameter appear correctly in team views and improves API documentation.

* Fix duplicate member checking in team_member_add endpoint

- Enhanced team_member_add_duplication_check to check both user_id and user_email
- Added additional duplicate prevention logic after user creation/lookup
- Fixed issue where users added by email could be duplicated in teams
- Added logging for debugging duplicate detection

This addresses the bug where adding the same user by email multiple times would create duplicate entries in the team's members_with_roles array.

* Improve duplicate member prevention in team_member_add

- Enhanced early duplicate check to handle both user_id and user_email
- Added late-stage duplicate prevention after user lookup/creation
- Fixed issue where users could be added multiple times by email
- Cleaned up debug logging

Note: There's still an edge case where the duplicate prevention may not work
correctly in all scenarios. This needs further investigation and testing.

* Refactor team_member_add endpoint for improved member management

- Split team_member_add functionality into smaller, dedicated functions for permission validation, member processing, and team member list updates.
- Enhanced permission checks to ensure only authorized users can add members.
- Streamlined member addition logic to reduce redundancy and improve readability.
- Maintained existing functionality while improving code structure and maintainability.

* Add tests for team_member_add helper functions

- Add test for _validate_team_member_add_permissions with admin user
- Add test for _validate_team_member_add_permissions with non-admin user
- Add test for _process_team_members with single member
- Add test for _process_team_members with multiple members
- Add test for _update_team_members_list with new member
- Add test for _update_team_members_list duplicate prevention

These tests ensure the refactored helper functions work correctly
after fixing the PLR0915 linting error.
This commit is contained in:
Cole McIntosh 2025-06-27 22:29:49 -06:00 • committed by GitHub
parent 1c68c24358
commit 6578133bb7
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 491 additions and 133 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

40
litellm/proxy/_types.py
View file

@ -1075,8 +1075,14 @@ class DeleteCustomerRequest(LiteLLMPydanticObjectBase):
class MemberBase(LiteLLMPydanticObjectBase):
user_id: Optional[str] = None
user_email: Optional[str] = None
user_id: Optional[str] = Field(
default=None,
description="The unique ID of the user to add. Either user_id or user_email must be provided"
)
user_email: Optional[str] = Field(
default=None,
description="The email address of the user to add. Either user_id or user_email must be provided"
)
@model_validator(mode="before")
@classmethod
@ -1092,7 +1098,9 @@ class Member(MemberBase):
role: Literal[
"admin",
"user",
]
] = Field(
description="The role of the user within the team. 'admin' users can manage team settings and members, 'user' is a regular team member"
)
class OrgMember(MemberBase):
@ -2527,7 +2535,9 @@ class LiteLLM_TeamMembership(LiteLLMPydanticObjectBase):
class MemberAddRequest(LiteLLMPydanticObjectBase):
member: Union[List[Member], Member]
member: Union[List[Member], Member] = Field(
description="Member object or list of member objects to add. Each member must include either user_id or user_email, and a role"
)
def __init__(self, **data):
member_data = data.get("member")
@ -2597,8 +2607,26 @@ class MemberUpdateResponse(LiteLLMPydanticObjectBase):
# Team Member Requests
class TeamMemberAddRequest(MemberAddRequest):
team_id: str
max_budget_in_team: Optional[float] = None # Users max budget within the team
"""
Request body for adding members to a team.
Example:
```json
{
"team_id": "45e3e396-ee08-4a61-a88e-16b3ce7e0849",
"member": {
"role": "user",
"user_id": "user123"
},
"max_budget_in_team": 100.0
}
```
"""
team_id: str = Field(description="The ID of the team to add the member to")
max_budget_in_team: Optional[float] = Field(
default=None,
description="Maximum budget allocated to this user within the team. If not set, user has unlimited budget within team limits"
)
class TeamMemberDeleteRequest(MemberDeleteRequest):

338
litellm/proxy/management_endpoints/team_endpoints.py
View file

@ -942,14 +942,33 @@ 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.
"""
def _check_member_duplication(member: Member):
if member.user_id in [m.user_id for m in existing_team_row.members_with_roles]:
raise ProxyException(
message=f"User={member.user_id} already in team. Existing members={existing_team_row.members_with_roles}",
type=ProxyErrorTypes.team_member_already_in_team,
param="user_id",
code="400",
)
# Check by user_id if provided
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:
raise ProxyException(
message=f"User with user_id={member.user_id} already in team. Existing members={existing_team_row.members_with_roles}",
type=ProxyErrorTypes.team_member_already_in_team,
param="user_id",
code="400",
)
# 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:
raise ProxyException(
message=f"User with user_email={member.user_email} already in team. Existing members={existing_team_row.members_with_roles}",
type=ProxyErrorTypes.team_member_already_in_team,
param="user_email",
code="400",
)
if isinstance(data.member, Member):
_check_member_duplication(data.member)
@ -958,6 +977,149 @@ def team_member_add_duplication_check(
_check_member_duplication(m)
async def _validate_team_member_add_permissions(
user_api_key_dict: UserAPIKeyAuth,
complete_team_data: LiteLLM_TeamTable,
) -> None:
"""Validate if user has permission to add members to the team."""
if (
hasattr(user_api_key_dict, "user_role")
and user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value
and not _is_user_team_admin(
user_api_key_dict=user_api_key_dict, team_obj=complete_team_data
)
and not _is_available_team(
team_id=complete_team_data.team_id,
user_api_key_dict=user_api_key_dict,
)
):
raise HTTPException(
status_code=403,
detail={
"error": "Call not allowed. User not proxy admin OR team admin. route={}, team_id={}".format(
"/team/member_add",
complete_team_data.team_id,
)
},
)
async def _process_team_members(
data: TeamMemberAddRequest,
complete_team_data: LiteLLM_TeamTable,
prisma_client: PrismaClient,
user_api_key_dict: UserAPIKeyAuth,
litellm_proxy_admin_name: str,
) -> Tuple[List[LiteLLM_UserTable], List[LiteLLM_TeamMembership]]:
"""Process and add new team members."""
updated_users: List[LiteLLM_UserTable] = []
updated_team_memberships: List[LiteLLM_TeamMembership] = []
default_team_budget_id = (
complete_team_data.metadata.get("team_member_budget_id")
if complete_team_data.metadata is not None
else None
)
if isinstance(data.member, Member):
try:
updated_user, updated_tm = await add_new_member(
new_member=data.member,
max_budget_in_team=data.max_budget_in_team,
prisma_client=prisma_client,
user_api_key_dict=user_api_key_dict,
litellm_proxy_admin_name=litellm_proxy_admin_name,
team_id=data.team_id,
default_team_budget_id=default_team_budget_id,
)
except Exception as e:
raise HTTPException(
status_code=500,
detail={
"error": "Unable to add user - {}, to team - {}, for reason - {}".format(
data.member, data.team_id, str(e)
)
},
)
updated_users.append(updated_user)
if updated_tm is not None:
updated_team_memberships.append(updated_tm)
elif isinstance(data.member, List):
for m in data.member:
try:
updated_user, updated_tm = await add_new_member(
new_member=m,
max_budget_in_team=data.max_budget_in_team,
prisma_client=prisma_client,
user_api_key_dict=user_api_key_dict,
litellm_proxy_admin_name=litellm_proxy_admin_name,
team_id=data.team_id,
default_team_budget_id=default_team_budget_id,
)
except Exception as e:
raise HTTPException(
status_code=500,
detail={
"error": "Unable to add user - {}, to team - {}, for reason - {}".format(
m, data.team_id, str(e)
)
},
)
updated_users.append(updated_user)
if updated_tm is not None:
updated_team_memberships.append(updated_tm)
return updated_users, updated_team_memberships
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 = data.member.model_copy()
# 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)
@router.post(
"/team/member_add",
tags=["team management"],
@ -1027,120 +1189,27 @@ async def team_member_add(
existing_team_row=complete_team_data,
)
## CHECK IF USER IS PROXY ADMIN OR TEAM ADMIN
# Validate permissions
await _validate_team_member_add_permissions(
user_api_key_dict=user_api_key_dict,
complete_team_data=complete_team_data,
)
if (
hasattr(user_api_key_dict, "user_role")
and user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value
and not _is_user_team_admin(
user_api_key_dict=user_api_key_dict, team_obj=complete_team_data
)
and not _is_available_team(
team_id=complete_team_data.team_id,
user_api_key_dict=user_api_key_dict,
)
):
raise HTTPException(
status_code=403,
detail={
"error": "Call not allowed. User not proxy admin OR team admin. route={}, team_id={}".format(
"/team/member_add",
complete_team_data.team_id,
)
},
)
# Process and add new members
updated_users, updated_team_memberships = await _process_team_members(
data=data,
complete_team_data=complete_team_data,
prisma_client=prisma_client,
user_api_key_dict=user_api_key_dict,
litellm_proxy_admin_name=litellm_proxy_admin_name,
)
updated_users: List[LiteLLM_UserTable] = []
updated_team_memberships: List[LiteLLM_TeamMembership] = []
## VALIDATE IF NEW MEMBER ##
if isinstance(data.member, Member):
try:
updated_user, updated_tm = await add_new_member(
new_member=data.member,
max_budget_in_team=data.max_budget_in_team,
prisma_client=prisma_client,
user_api_key_dict=user_api_key_dict,
litellm_proxy_admin_name=litellm_proxy_admin_name,
team_id=data.team_id,
default_team_budget_id=(
complete_team_data.metadata.get("team_member_budget_id")
if complete_team_data.metadata is not None
else None
),
)
except Exception as e:
raise HTTPException(
status_code=500,
detail={
"error": "Unable to add user - {}, to team - {}, for reason - {}".format(
data.member, data.team_id, str(e)
)
},
)
updated_users.append(updated_user)
if updated_tm is not None:
updated_team_memberships.append(updated_tm)
elif isinstance(data.member, List):
tasks: List = []
for m in data.member:
try:
updated_user, updated_tm = await add_new_member(
new_member=m,
max_budget_in_team=data.max_budget_in_team,
prisma_client=prisma_client,
user_api_key_dict=user_api_key_dict,
litellm_proxy_admin_name=litellm_proxy_admin_name,
team_id=data.team_id,
default_team_budget_id=(
complete_team_data.metadata.get("team_member_budget_id")
if complete_team_data.metadata is not None
else None
),
)
except Exception as e:
raise HTTPException(
status_code=500,
detail={
"error": "Unable to add user - {}, to team - {}, for reason - {}".format(
data.member, data.team_id, str(e)
)
},
)
updated_users.append(updated_user)
if updated_tm is not None:
updated_team_memberships.append(updated_tm)
await asyncio.gather(*tasks)
## ADD TO TEAM ##
if isinstance(data.member, Member):
# add to team db
new_member = data.member
# 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
complete_team_data.members_with_roles.append(new_member)
elif isinstance(data.member, List):
# add to team db
new_members = data.member
for nm in new_members:
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
complete_team_data.members_with_roles.extend(new_members)
# Update team members list
await _update_team_members_list(
data=data,
complete_team_data=complete_team_data,
updated_users=updated_users,
)
# ADD MEMBER TO TEAM
_db_team_members = [m.model_dump() for m in complete_team_data.members_with_roles]
@ -2104,15 +2173,30 @@ async def list_team(
filtered_response = []
if user_id:
for team in response:
if team.members_with_roles:
for member in team.members_with_roles:
if (
"user_id" in member
and member["user_id"] is not None
and member["user_id"] == user_id
):
# Get user object to access their teams array
try:
user_object = await prisma_client.db.litellm_usertable.find_unique(
where={"user_id": user_id}
)
if user_object and user_object.teams:
# Filter teams based on user's teams array
for team in response:
if team.team_id in user_object.teams:
filtered_response.append(team)
except Exception as e:
verbose_proxy_logger.debug(
f"Error fetching user for team filtering: {str(e)}"
)
# Fall back to checking members_with_roles if user lookup fails
for team in response:
if team.members_with_roles:
for member in team.members_with_roles:
if (
"user_id" in member
and member["user_id"] is not None
and member["user_id"] == user_id
):
filtered_response.append(team)
else:
filtered_response = response

246
tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py
View file

@ -738,3 +738,249 @@ def test_add_new_models_to_team():
"model5",
].sort()
)
@pytest.mark.asyncio
async def test_validate_team_member_add_permissions_admin():
"""
Test _validate_team_member_add_permissions allows proxy admin
"""
from litellm.proxy.management_endpoints.team_endpoints import (
_validate_team_member_add_permissions,
)
# Create admin user
admin_user = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN.value)
# Create mock team
team = MagicMock(spec=LiteLLM_TeamTable)
team.team_id = "test-team-123"
# Should not raise any exception for admin
await _validate_team_member_add_permissions(
user_api_key_dict=admin_user,
complete_team_data=team,
)
@pytest.mark.asyncio
async def test_validate_team_member_add_permissions_non_admin():
"""
Test _validate_team_member_add_permissions raises exception for non-admin non-team-admin
"""
from litellm.proxy.management_endpoints.team_endpoints import (
_validate_team_member_add_permissions,
)
# Create non-admin user
regular_user = UserAPIKeyAuth(
user_id="regular-user",
user_role=LitellmUserRoles.INTERNAL_USER.value,
team_id="different-team"
)
# Create mock team
team = MagicMock(spec=LiteLLM_TeamTable)
team.team_id = "test-team-123"
team.members_with_roles = []
# Mock the helper functions to return False
with patch(
"litellm.proxy.management_endpoints.team_endpoints._is_user_team_admin",
return_value=False
), patch(
"litellm.proxy.management_endpoints.team_endpoints._is_available_team",
return_value=False
):
# Should raise HTTPException for non-admin
with pytest.raises(HTTPException) as exc_info:
await _validate_team_member_add_permissions(
user_api_key_dict=regular_user,
complete_team_data=team,
)
assert exc_info.value.status_code == 403
assert "not proxy admin OR team admin" in str(exc_info.value.detail)
@pytest.mark.asyncio
async def test_process_team_members_single_member():
"""
Test _process_team_members with a single member
"""
from litellm.proxy.management_endpoints.team_endpoints import _process_team_members
from litellm.proxy._types import LiteLLM_UserTable, LiteLLM_TeamMembership
# Mock dependencies
mock_prisma_client = MagicMock()
mock_team = MagicMock(spec=LiteLLM_TeamTable)
mock_team.metadata = {"team_member_budget_id": "budget-123"}
# Mock user and membership objects
mock_user = MagicMock(spec=LiteLLM_UserTable)
mock_user.user_id = "new-user-123"
mock_membership = MagicMock(spec=LiteLLM_TeamMembership)
# Create request with single member
single_member = Member(user_email="new@example.com", role="user")
request_data = TeamMemberAddRequest(
team_id="test-team-123",
member=single_member,
)
with patch(
"litellm.proxy.management_endpoints.team_endpoints.add_new_member",
new_callable=AsyncMock,
return_value=(mock_user, mock_membership)
) as mock_add_member:
users, memberships = await _process_team_members(
data=request_data,
complete_team_data=mock_team,
prisma_client=mock_prisma_client,
user_api_key_dict=UserAPIKeyAuth(),
litellm_proxy_admin_name="admin",
)
# Verify results
assert len(users) == 1
assert len(memberships) == 1
assert users[0] == mock_user
assert memberships[0] == mock_membership
# Verify add_new_member was called correctly
mock_add_member.assert_called_once_with(
new_member=single_member,
max_budget_in_team=None,
prisma_client=mock_prisma_client,
user_api_key_dict=UserAPIKeyAuth(),
litellm_proxy_admin_name="admin",
team_id="test-team-123",
default_team_budget_id="budget-123",
)
@pytest.mark.asyncio
async def test_process_team_members_multiple_members():
"""
Test _process_team_members with multiple members
"""
from litellm.proxy.management_endpoints.team_endpoints import _process_team_members
from litellm.proxy._types import LiteLLM_UserTable, LiteLLM_TeamMembership
# Mock dependencies
mock_prisma_client = MagicMock()
mock_team = MagicMock(spec=LiteLLM_TeamTable)
mock_team.metadata = None
# Create multiple members as dictionaries (they will be converted to Member objects)
members = [
{"user_email": "user1@example.com", "role": "user"},
{"user_email": "user2@example.com", "role": "admin"},
]
request_data = TeamMemberAddRequest(
team_id="test-team-123",
member=members,
max_budget_in_team=100.0,
)
# Mock different users and memberships for each call
mock_users = [MagicMock(spec=LiteLLM_UserTable) for _ in range(2)]
mock_memberships = [MagicMock(spec=LiteLLM_TeamMembership) for _ in range(2)]
with patch(
"litellm.proxy.management_endpoints.team_endpoints.add_new_member",
new_callable=AsyncMock,
side_effect=[(mock_users[0], mock_memberships[0]), (mock_users[1], mock_memberships[1])]
) as mock_add_member:
users, memberships = await _process_team_members(
data=request_data,
complete_team_data=mock_team,
prisma_client=mock_prisma_client,
user_api_key_dict=UserAPIKeyAuth(),
litellm_proxy_admin_name="admin",
)
# Verify results
assert len(users) == 2
assert len(memberships) == 2
assert users == mock_users
assert memberships == mock_memberships
# Verify add_new_member was called for each member
assert mock_add_member.call_count == 2
@pytest.mark.asyncio
async def test_update_team_members_list_single_member():
"""
Test _update_team_members_list with a single member
"""
from litellm.proxy.management_endpoints.team_endpoints import _update_team_members_list
from litellm.proxy._types import LiteLLM_UserTable
# Create mock team with existing members
mock_team = MagicMock(spec=LiteLLM_TeamTable)
mock_team.members_with_roles = [
Member(user_id="existing-user", role="admin")
]
# Create new member without user_id
new_member = Member(user_email="new@example.com", role="user")
request_data = TeamMemberAddRequest(
team_id="test-team-123",
member=new_member,
)
# Create mock user with matching email
mock_user = MagicMock(spec=LiteLLM_UserTable)
mock_user.user_id = "new-user-123"
mock_user.user_email = "new@example.com"
await _update_team_members_list(
data=request_data,
complete_team_data=mock_team,
updated_users=[mock_user],
)
# Verify member was added
assert len(mock_team.members_with_roles) == 2
added_member = mock_team.members_with_roles[1]
assert added_member.user_id == "new-user-123"
assert added_member.user_email == "new@example.com"
assert added_member.role == "user"
@pytest.mark.asyncio
async def test_update_team_members_list_duplicate_prevention():
"""
Test _update_team_members_list prevents duplicate members
"""
from litellm.proxy.management_endpoints.team_endpoints import _update_team_members_list
from litellm.proxy._types import LiteLLM_UserTable
# Create mock team with existing members
mock_team = MagicMock(spec=LiteLLM_TeamTable)
mock_team.members_with_roles = [
Member(user_id="existing-user", user_email="existing@example.com", role="admin")
]
# Try to add the same member again
duplicate_member = Member(user_id="existing-user", role="user")
request_data = TeamMemberAddRequest(
team_id="test-team-123",
member=duplicate_member,
)
# Create mock user
mock_user = MagicMock(spec=LiteLLM_UserTable)
mock_user.user_id = "existing-user"
mock_user.user_email = "existing@example.com"
await _update_team_members_list(
data=request_data,
complete_team_data=mock_team,
updated_users=[mock_user],
)
# Verify member was NOT added (still only 1 member)
assert len(mock_team.members_with_roles) == 1