From b77c9f5de2b15aa93c540a77d87d2d84ea3d7ba5 Mon Sep 17 00:00:00 2001 From: Krish Dholakia Date: Thu, 17 Jul 2025 22:31:44 -0700 Subject: [PATCH] =?UTF-8?q?fix(team=5Fendpoints.py):=20ensure=20user=20id?= =?UTF-8?q?=20correctly=20added=20when=20new=20team=20=E2=80=A6=20(#12719)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix(team_endpoints.py): ensure user id correctly added when new team created with user email as member Fixes issue where user not correctly added to team on /team/new * fix(internal_user_endpoints.py): make user email validation check case insensitive Fixes issue where uppercase email was added even when lowercase email existed * test: update test --- litellm/proxy/_types.py | 5 +- .../internal_user_endpoints.py | 6 +- .../management_endpoints/team_endpoints.py | 93 ++++++++++++------- .../test_internal_user_endpoints.py | 79 ++++++++++++++++ 4 files changed, 147 insertions(+), 36 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 8c6ba7172f..44f9da39f1 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -2540,7 +2540,10 @@ class MemberAddRequest(LiteLLMPydanticObjectBase): member_data = data.get("member") if isinstance(member_data, list): # If member is a list of dictionaries, convert each dictionary to a Member object - members = [Member(**item) for item in member_data] + members = [ + Member(**item) if isinstance(item, dict) else item + for item in member_data + ] # Replace member_data with the list of Member objects data["member"] = members elif isinstance(member_data, dict): diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index bd94d46ec4..808033b070 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -122,13 +122,15 @@ async def _check_duplicate_user_email( raise Exception("Database not connected") existing_user = await prisma_client.db.litellm_usertable.find_first( - where={"user_email": user_email.strip()} + where={"user_email": {"equals": user_email.strip(), "mode": "insensitive"}} ) if existing_user is not None: raise HTTPException( status_code=400, - detail={"error": f"User with email {user_email} already exists"}, + detail={ + "error": f"User with email {existing_user.user_email} already exists" + }, ) diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index b48d063b70..a5cd2321fa 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -406,7 +406,6 @@ async def new_team( # noqa: PLR0915 _model_id = model_dict.id - ## Handle Object Permission - MCP, Vector Stores etc. object_permission_id = await _set_object_permission( data=data, @@ -447,27 +446,34 @@ async def new_team( # noqa: PLR0915 budget_duration=complete_team_data.budget_duration, ) + ## Add Team Member Budget Table + members_with_roles: List[Member] = [] + if complete_team_data.members_with_roles is not None: + members_with_roles = complete_team_data.members_with_roles + complete_team_data.members_with_roles = [] + complete_team_data_dict = complete_team_data.model_dump(exclude_none=True) complete_team_data_dict = prisma_client.jsonify_team_object( db_data=complete_team_data_dict ) + team_row: LiteLLM_TeamTable = await prisma_client.db.litellm_teamtable.create( data=complete_team_data_dict, include={"litellm_model_table": True}, # type: ignore ) ## ADD TEAM ID TO USER TABLE ## - for user in complete_team_data.members_with_roles: - ## add team id to user row ## - await prisma_client.update_data( - user_id=user.user_id, - data={"user_id": user.user_id, "teams": [team_row.team_id]}, - update_key_values_custom_query={ - "teams": { - "push ": [team_row.team_id], - } - }, - ) + team_member_add_request = TeamMemberAddRequest( + team_id=data.team_id, + member=members_with_roles, + ) + await _add_team_members_to_team( + data=team_member_add_request, + complete_team_data=team_row, + prisma_client=prisma_client, + user_api_key_dict=user_api_key_dict, + litellm_proxy_admin_name=litellm_proxy_admin_name, + ) # Enterprise Feature - Audit Logging. Enable with litellm.store_audit_logs = True if litellm.store_audit_logs is True: @@ -1131,6 +1137,40 @@ async def _update_team_members_list( complete_team_data.members_with_roles.append(nm) +async def _add_team_members_to_team( + data: TeamMemberAddRequest, + complete_team_data: LiteLLM_TeamTable, + prisma_client: PrismaClient, + user_api_key_dict: UserAPIKeyAuth, + litellm_proxy_admin_name: str, +) -> Tuple[LiteLLM_TeamTable, List[LiteLLM_UserTable], List[LiteLLM_TeamMembership]]: + """Add team members to the team.""" + # 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, + ) + + # 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] + updated_team = await prisma_client.db.litellm_teamtable.update( + where={"team_id": data.team_id}, + data={"members_with_roles": json.dumps(_db_team_members)}, # type: ignore + ) + + return updated_team, updated_users, updated_team_memberships + + @router.post( "/team/member_add", tags=["team management"], @@ -1206,27 +1246,14 @@ async def team_member_add( complete_team_data=complete_team_data, ) - # 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, - ) - - # 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] - updated_team = await prisma_client.db.litellm_teamtable.update( - where={"team_id": data.team_id}, - data={"members_with_roles": json.dumps(_db_team_members)}, # type: ignore + updated_team, updated_users, updated_team_memberships = ( + await _add_team_members_to_team( + 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, + ) ) # Check if updated_team is None diff --git a/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py index 1ce5d24c8c..266056bcdd 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py @@ -611,3 +611,82 @@ def test_update_internal_new_user_params_internal_user_role(): else: if hasattr(litellm, "default_internal_user_params"): delattr(litellm, "default_internal_user_params") + + +@pytest.mark.asyncio +async def test_check_duplicate_user_email_case_insensitive(mocker): + """ + Test that _check_duplicate_user_email performs case insensitive email matching. + + This ensures that emails like 'User@Example.com' and 'user@example.com' + are treated as the same user, preventing duplicate accounts. + """ + from fastapi import HTTPException + + from litellm.proxy.management_endpoints.internal_user_endpoints import ( + _check_duplicate_user_email, + ) + + # Mock the prisma client + mock_prisma_client = mocker.MagicMock() + + # Test Case 1: Duplicate found with different case + # Mock existing user with uppercase email + mock_existing_user = mocker.MagicMock() + mock_existing_user.user_email = "User@Example.com" + + async def mock_find_first_duplicate(*args, **kwargs): + # Verify that the query uses case insensitive matching + where_clause = kwargs.get("where", {}) + user_email_clause = where_clause.get("user_email", {}) + + # Check that the query structure is correct for case insensitive search + assert ( + "equals" in user_email_clause + ), "Query should use 'equals' for case insensitive search" + assert ( + user_email_clause.get("mode") == "insensitive" + ), "Query should use 'insensitive' mode" + assert ( + user_email_clause.get("equals") == "user@example.com" + ), "Query should search for the provided email" + + return mock_existing_user # Return existing user to simulate duplicate + + mock_prisma_client.db.litellm_usertable.find_first = mock_find_first_duplicate + + # Should raise HTTPException when duplicate is found + with pytest.raises(HTTPException) as exc_info: + await _check_duplicate_user_email("user@example.com", mock_prisma_client) + + assert exc_info.value.status_code == 400 + assert "User with email User@Example.com already exists" in str( + exc_info.value.detail + ) + + # Test Case 2: No duplicate found + async def mock_find_first_no_duplicate(*args, **kwargs): + # Verify the query structure again + where_clause = kwargs.get("where", {}) + user_email_clause = where_clause.get("user_email", {}) + + assert "equals" in user_email_clause + assert user_email_clause.get("mode") == "insensitive" + assert user_email_clause.get("equals") == "newuser@example.com" + + return None # No existing user found + + mock_prisma_client.db.litellm_usertable.find_first = mock_find_first_no_duplicate + + # Should not raise any exception when no duplicate is found + try: + await _check_duplicate_user_email("newuser@example.com", mock_prisma_client) + # If we reach here, no exception was raised (which is expected) + assert True + except Exception as e: + pytest.fail(f"Should not raise exception when no duplicate found, but got: {e}") + + # Test Case 3: None email should not cause issues + await _check_duplicate_user_email( + None, mock_prisma_client + ) # Should not raise exception