From 6578133bb7b768438510ef1259de6a66110e228c Mon Sep 17 00:00:00 2001 From: Cole McIntosh <82463175+colesmcintosh@users.noreply.github.com> Date: Fri, 27 Jun 2025 22:29:49 -0600 Subject: [PATCH] Fix user-team association issues in LiteLLM proxy (#12082) * 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. --- litellm/proxy/_types.py | 40 ++- .../management_endpoints/team_endpoints.py | 338 +++++++++++------- .../test_team_endpoints.py | 246 +++++++++++++ 3 files changed, 491 insertions(+), 133 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index d4d852031da..bc0332b33a5 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -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): diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 3d4ae4e945f..7dec4b78e6a 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -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 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 9aee01ed07d..67e4966f985 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -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