diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index d0afb62e8bc..2c31b59b795 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -1771,6 +1771,116 @@ async def _add_team_members_to_team( return updated_team, updated_users, updated_team_memberships +async def _validate_and_populate_member_user_info( + member: Member, + prisma_client: PrismaClient, +) -> Member: + """ + 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) + 2. If only user_email is provided, populate user_id from DB + 3. If only user_id is provided, populate user_email from DB + 4. If only user_id is provided and doesn't exist, throw error + 5. If user_email and user_id mismatch, throw error + + Returns a Member with both user_email and user_id populated. + """ + if member.user_email is None and member.user_id is None: + raise HTTPException( + status_code=400, + 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", + ) + + 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 + 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 + return member + + # Case 2: Only user_email provided - populate user_id from DB + if member.user_email is not None and member.user_id is None: + user_by_email = await prisma_client.db.litellm_usertable.find_first( + where={"user_email": {"equals": member.user_email, "mode": "insensitive"}} + ) + + if user_by_email is None: + # User doesn't exist yet - this is fine, will be created later + return member + + # Check for multiple users with same email + users_by_email = await prisma_client.get_data( + key_val={"user_email": member.user_email}, + table_name="user", + query_type="find_all", + ) + + if users_by_email and 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." + }, + ) + + # Populate user_id + member.user_id = user_by_email.user_id + return member + + # Case 3: Only user_id provided - populate user_email from DB + if member.user_id is not None and member.user_email is None: + user_by_id = await prisma_client.db.litellm_usertable.find_unique( + where={"user_id": member.user_id} + ) + + if user_by_id is None: + raise HTTPException( + status_code=404, + detail={ + "error": f"User with user_id '{member.user_id}' not found in database" + }, + ) + + # Populate user_email + member.user_email = user_by_id.user_email + return member + + return member + @router.post( "/team/member_add", tags=["team management"], @@ -1846,6 +1956,19 @@ async def team_member_add( complete_team_data=complete_team_data, ) + # Validate and populate user_email/user_id for members before processing + if isinstance(data.member, Member): + await _validate_and_populate_member_user_info( + member=data.member, + prisma_client=prisma_client, + ) + elif isinstance(data.member, List): + for m in data.member: + await _validate_and_populate_member_user_info( + member=m, + prisma_client=prisma_client, + ) + updated_team, updated_users, updated_team_memberships = ( await _add_team_members_to_team( data=data, 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 a97ed93a85c..a5864576906 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -36,6 +36,7 @@ from litellm.proxy.management_endpoints.team_endpoints import ( _persist_deleted_team_records, _save_deleted_team_records, _transform_teams_to_deleted_records, + _validate_and_populate_member_user_info, delete_team, router, team_member_add_duplication_check, @@ -5466,3 +5467,122 @@ async def test_get_team_daily_activity_team_admin_sees_all_spend(mock_db_client) ) and mock_db_client.db.litellm_verificationtoken.find_many.called: # If it was called, that's unexpected for admin users assert False, "API keys should not be fetched for team admin users" + + +@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") + + # 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, + ) + + # Verify result matches input (both already provided and match) + assert result.user_email == "test@example.com" + assert result.user_id == "user-123" + + # 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 +async def test_validate_and_populate_member_user_info_only_email_provided(): + """ + Test _validate_and_populate_member_user_info when only user_email is provided. + Should populate user_id from database. + """ + # Create member with only user_email + member = Member(user_email="test@example.com", user_id=None, role="user") + + # Mock prisma client + mock_prisma_client = MagicMock() + + # Mock user object from find_first + mock_user_find_first = MagicMock() + mock_user_find_first.user_id = "user-456" + mock_user_find_first.user_email = "test@example.com" + + # Mock find_first to return the user + mock_prisma_client.db.litellm_usertable.find_first = AsyncMock( + return_value=mock_user_find_first + ) + + # Mock get_data to return single user (no duplicates) + mock_prisma_client.get_data = AsyncMock(return_value=[mock_user_find_first]) + + # Call the function + result = await _validate_and_populate_member_user_info( + member=member, + prisma_client=mock_prisma_client, + ) + + # Verify user_id was populated + assert result.user_email == "test@example.com" + assert result.user_id == "user-456" + + # Verify find_first was called with correct parameters + mock_prisma_client.db.litellm_usertable.find_first.assert_called_once_with( + where={"user_email": {"equals": "test@example.com", "mode": "insensitive"}} + ) + + # Verify get_data was called to check for duplicates + 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 +async def test_validate_and_populate_member_user_info_only_user_id_not_found(): + """ + Test _validate_and_populate_member_user_info when only user_id is provided + but the user doesn't exist in the database. Should raise HTTPException. + """ + # Create member with only user_id + member = Member(user_email=None, user_id="nonexistent-user", role="user") + + # Mock prisma client + mock_prisma_client = MagicMock() + + # Mock find_unique to return None (user not found) + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None) + + # Call the function and expect HTTPException + with pytest.raises(HTTPException) as exc_info: + await _validate_and_populate_member_user_info( + member=member, + prisma_client=mock_prisma_client, + ) + + # Verify the exception details + assert exc_info.value.status_code == 404 + assert "not found" in exc_info.value.detail["error"].lower() + assert "nonexistent-user" in exc_info.value.detail["error"] + + # Verify find_unique was called with correct parameters + mock_prisma_client.db.litellm_usertable.find_unique.assert_called_once_with( + where={"user_id": "nonexistent-user"} + )