From f6804333b899806414a985468a2d87f3be147d17 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Fri, 5 Dec 2025 15:08:13 -0800 Subject: [PATCH 1/3] New user route user_id conflict handling --- .../internal_user_endpoints.py | 85 ++++++++++++++----- .../test_internal_user_endpoints.py | 70 ++++++++++++++- 2 files changed, 130 insertions(+), 25 deletions(-) diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index 7c93c8424ab..1850ffa2560 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -101,35 +101,75 @@ def _update_internal_new_user_params(data_json: dict, data: NewUserRequest) -> d return data_json +async def _check_duplicate_user_field( + field_name: str, + field_value: Optional[str], + prisma_client: Any, + *, + case_insensitive: bool = False, + label: Optional[str] = None, +) -> None: + """ + Helper function to check if a field already exists in the user table. + + Args: + field_name (str): Database field name to check. + field_value (Optional[str]): Value to check for duplicates. + prisma_client (Any): Database client instance. + case_insensitive (bool): Whether to use case-insensitive comparison. + label (Optional[str]): Human readable label for error messages. + + Raises: + Exception: If database is not connected. + HTTPException: If a user with the given field value already exists. + """ + if field_value: + if prisma_client is None: + raise Exception("Database not connected") + + value = field_value.strip() + where_clause = {field_name: {"equals": value}} + if case_insensitive: + where_clause[field_name]["mode"] = "insensitive" + + existing_user = await prisma_client.db.litellm_usertable.find_first( + where=where_clause + ) + + if existing_user is not None: + existing_value = getattr(existing_user, field_name, value) + error_label = label or field_name + raise HTTPException( + status_code=409, + detail={"error": f"User with {error_label} {existing_value} already exists"}, + ) + + async def _check_duplicate_user_email( user_email: Optional[str], prisma_client: Any ) -> None: """ Helper function to check if a user email already exists in the database. - - Args: - user_email (Optional[str]): Email to check - prisma_client (Any): Database client instance - - Raises: - Exception: If database is not connected - HTTPException: If user with email already exists """ - if user_email: - if prisma_client is None: - raise Exception("Database not connected") + await _check_duplicate_user_field( + field_name="user_email", + field_value=user_email, + prisma_client=prisma_client, + case_insensitive=True, + label="email", + ) - existing_user = await prisma_client.db.litellm_usertable.find_first( - 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 {existing_user.user_email} already exists" - }, - ) +async def _check_duplicate_user_id(user_id: Optional[str], prisma_client: Any) -> None: + """ + Helper function to check if a user id already exists in the database. + """ + await _check_duplicate_user_field( + field_name="user_id", + field_value=user_id, + prisma_client=prisma_client, + label="id", + ) async def _add_user_to_organizations( @@ -361,7 +401,8 @@ async def new_user( status_code=500, detail=CommonProxyErrors.db_not_connected_error.value, ) - # Check for duplicate email + # Check for duplicate user_id or email + await _check_duplicate_user_id(data.user_id, prisma_client) await _check_duplicate_user_email(data.user_email, prisma_client) # Check if license is over limit 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 d30cce067a0..33f2a75fac6 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 @@ -261,14 +261,21 @@ async def test_new_user_license_over_limit(mocker): mock_prisma_client.db.litellm_usertable.count = mock_count - # Mock check_duplicate_user_email to pass + # Mock duplicate checks to pass async def mock_check_duplicate_user_email(*args, **kwargs): return None # No duplicate found + async def mock_check_duplicate_user_id(*args, **kwargs): + return None # No duplicate found + mocker.patch( "litellm.proxy.management_endpoints.internal_user_endpoints._check_duplicate_user_email", mock_check_duplicate_user_email, ) + mocker.patch( + "litellm.proxy.management_endpoints.internal_user_endpoints._check_duplicate_user_id", + mock_check_duplicate_user_id, + ) # Mock the license check to return True (over limit) mock_license_check = mocker.MagicMock() @@ -449,14 +456,21 @@ async def test_new_user_default_teams_flow(mocker): mock_prisma_client.db.litellm_usertable.count = mock_count - # Mock check_duplicate_user_email to pass + # Mock duplicate checks to pass async def mock_check_duplicate_user_email(*args, **kwargs): return None # No duplicate found + async def mock_check_duplicate_user_id(*args, **kwargs): + return None # No duplicate found + mocker.patch( "litellm.proxy.management_endpoints.internal_user_endpoints._check_duplicate_user_email", mock_check_duplicate_user_email, ) + mocker.patch( + "litellm.proxy.management_endpoints.internal_user_endpoints._check_duplicate_user_id", + mock_check_duplicate_user_id, + ) # Mock the license check to return False (under limit) mock_license_check = mocker.MagicMock() @@ -737,7 +751,7 @@ async def test_check_duplicate_user_email_case_insensitive(mocker): 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 exc_info.value.status_code == 409 assert "User with email User@Example.com already exists" in str( exc_info.value.detail ) @@ -770,6 +784,56 @@ async def test_check_duplicate_user_email_case_insensitive(mocker): ) # Should not raise exception +@pytest.mark.asyncio +async def test_check_duplicate_user_id(mocker): + """ + Test that _check_duplicate_user_id detects duplicates and does not use case insensitive matching. + """ + from fastapi import HTTPException + + from litellm.proxy.management_endpoints.internal_user_endpoints import ( + _check_duplicate_user_id, + ) + + mock_prisma_client = mocker.MagicMock() + + # Duplicate user_id should raise + mock_existing_user = mocker.MagicMock() + mock_existing_user.user_id = "existing-user-id" + + async def mock_find_first_duplicate(*args, **kwargs): + where_clause = kwargs.get("where", {}) + user_id_clause = where_clause.get("user_id", {}) + assert user_id_clause.get("equals") == "existing-user-id" + assert "mode" not in user_id_clause + return mock_existing_user + + mock_prisma_client.db.litellm_usertable.find_first = mock_find_first_duplicate + + with pytest.raises(HTTPException) as exc_info: + await _check_duplicate_user_id("existing-user-id", mock_prisma_client) + + assert exc_info.value.status_code == 409 + assert "User with id existing-user-id already exists" in str( + exc_info.value.detail + ) + + # No duplicate should pass + async def mock_find_first_no_duplicate(*args, **kwargs): + where_clause = kwargs.get("where", {}) + user_id_clause = where_clause.get("user_id", {}) + assert user_id_clause.get("equals") == "new-user-id" + assert "mode" not in user_id_clause + return None + + mock_prisma_client.db.litellm_usertable.find_first = mock_find_first_no_duplicate + + await _check_duplicate_user_id("new-user-id", mock_prisma_client) + + # None user_id should no-op + await _check_duplicate_user_id(None, mock_prisma_client) + + def test_process_keys_for_user_info_filters_dashboard_keys(monkeypatch): """ Test that _process_keys_for_user_info filters out keys with team_id='litellm-dashboard' From a2bf6e169cdbe9db41a9068ae033e102fb07e777 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Fri, 5 Dec 2025 16:07:51 -0800 Subject: [PATCH 2/3] fixing tests --- tests/proxy_admin_ui_tests/test_key_management.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/proxy_admin_ui_tests/test_key_management.py b/tests/proxy_admin_ui_tests/test_key_management.py index 589a394cbbc..126718af848 100644 --- a/tests/proxy_admin_ui_tests/test_key_management.py +++ b/tests/proxy_admin_ui_tests/test_key_management.py @@ -341,7 +341,7 @@ async def test_get_users(prisma_client): # Create some test users test_users = [ NewUserRequest( - user_id=f"test_user_{i}", + user_id=f"test_user_{i}_{uuid.uuid4()}", user_role=( LitellmUserRoles.INTERNAL_USER.value if i % 2 == 0 From a2e977b8e66358aacb45fd2b395cf5b694d64521 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Sat, 6 Dec 2025 22:28:23 -0800 Subject: [PATCH 3/3] Fixing tests --- tests/test_team.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/tests/test_team.py b/tests/test_team.py index 06a2e7a3648..4ef9b8d9195 100644 --- a/tests/test_team.py +++ b/tests/test_team.py @@ -683,26 +683,28 @@ async def test_team_alias(): @pytest.mark.asyncio async def test_users_in_team_budget(): """ - - Create Team - Create User + - Create Team with User - Add User to team with budget = 0.0000001 - Make Call 1 -> pass - Make Call 2 -> fail """ get_user = f"krrish_{time.time()}@berri.ai" async with aiohttp.ClientSession() as session: - team = await new_team(session, 0, user_id=get_user) - print("New team=", team) + # Create user first to avoid user_id collision when creating team key_gen = await new_user( session, 0, user_id=get_user, budget=10, budget_duration="5s", - team_id=team["team_id"], models=["fake-openai-endpoint"], ) key = key_gen["key"] + + # Create team with the user (user already exists, so it will just add them) + team = await new_team(session, 0, user_id=get_user) + print("New team=", team) # update user to have budget = 0.0000001 await update_member(