fix(proxy): do not claim rows a concurrent request inserted when /user/bulk_new create_many fails

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

litellm/proxy/management_helpers/bulk_user_creation.py

@@ -31,6 +31,7 @@ from litellm.proxy._types import (
 from litellm.proxy.auth.auth_checks import invalidate_team_member_spend_state
 from litellm.proxy.auth.litellm_license import LicenseCheck
 from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time
+from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler
 from litellm.proxy.hooks.user_management_event_hooks import UserManagementEventHooks
 from litellm.proxy.management_endpoints.common_utils import (
     _is_user_org_admin_for_team,  # pyright: ignore[reportPrivateUsage]  # same team-admin check /user/new uses
@@ -433,23 +434,38 @@ async def _insert_users(
     try:
         await table.create_many(data=payloads)
         return tuple(prepared), ()
-    except Exception:  # noqa: BLE001  # fall back to per-row inserts so the failing row can be identified
+    except Exception as exc:  # noqa: BLE001  # fall back to per-row inserts so the failing row can be identified
         verbose_proxy_logger.warning("/user/bulk_new: create_many failed, retrying rows individually", exc_info=True)
-    landed_rows: Final = await table.find_many(
-        where={"user_id": {"in": [payload["user_id"] for payload in payloads]}}  # mutable-ok: Prisma filter
-    )
+        outcome_unknown: Final = PrismaDBExceptionHandler.is_database_infrastructure_error(exc)
+    requested: Final = frozenset(payload["user_id"] for payload in payloads)
+    landed_rows: Final = await table.find_many(where={"user_id": {"in": list(requested)}})  # mutable-ok: Prisma filter
     landed: Final = frozenset(row.user_id for row in landed_rows)
+    # create_many is one INSERT: after a lost response the full set is ours, any partial set belongs to another request
+    if outcome_unknown and landed == requested:
+        return tuple(prepared), ()
+    taken: Final = tuple(user for user in prepared if user.row.user_id in landed)
     retried: Final = tuple(user for user in prepared if user.row.user_id not in landed)
     outcomes: Final = await _bounded(
         BULK_NEW_USER_CONCURRENCY, tuple(table.create(data=_user_create_payload(user)) for user in retried)
     )
     failed: Final = MappingProxyType(
         {
-            user.row.user_id: _RowFailure(
-                user.pending.index, user.pending.user_id, user.row.user_email, _error_message(outcome)
-            )
-            for user, outcome in zip(retried, outcomes, strict=True)
-            if isinstance(outcome, BaseException)
+            **{
+                user.row.user_id: _RowFailure(
+                    user.pending.index,
+                    user.pending.user_id,
+                    user.row.user_email,
+                    f"User id={user.row.user_id} already exists",
+                )
+                for user in taken
+            },
+            **{
+                user.row.user_id: _RowFailure(
+                    user.pending.index, user.pending.user_id, user.row.user_email, _error_message(outcome)
+                )
+                for user, outcome in zip(retried, outcomes, strict=True)
+                if isinstance(outcome, BaseException)
+            },
         }
     )
     return (

tests/test_litellm/proxy/management_helpers/test_bulk_user_creation.py

@@ -2,8 +2,10 @@ import json
 from contextlib import asynccontextmanager
 from typing import Final
 
+import httpx
 import pytest
 from fastapi import HTTPException
+from prisma.errors import UniqueViolationError
 from pydantic import BaseModel, ConfigDict, ValidationError
 
 from litellm.caching.caching import DualCache
@@ -31,10 +33,16 @@ class _UserRow(BaseModel):
 class _UserTable:
     """Enough of the Prisma user table for the bulk path: set lookups, one create_many and per-row fallbacks."""
 
-    def __init__(self, fail_ids: frozenset[str] = frozenset(), commit_then_drop: bool = False) -> None:
+    def __init__(
+        self,
+        fail_ids: frozenset[str] = frozenset(),
+        commit_then_drop: bool = False,
+        raced_ids: frozenset[str] = frozenset(),
+    ) -> None:
         self.rows: dict[str, _UserRow] = {}
         self.fail_ids = fail_ids
         self.commit_then_drop = commit_then_drop
+        self.raced_ids = raced_ids
         self.create_many_calls = 0
 
     async def count(self, where: object = None) -> int:
@@ -59,10 +67,15 @@ class _UserTable:
         rows = [_UserRow.model_validate(d) for d in data]
         if any(row.user_id in self.fail_ids for row in rows):
             raise RuntimeError("batch insert failed")
+        raced = [row.user_id for row in rows if row.user_id in self.raced_ids]
+        if raced:
+            for user_id in raced:
+                self.rows[user_id] = _UserRow(user_id=user_id, user_email=f"{user_id}@other-request.example")
+            raise UniqueViolationError({}, message="Unique constraint failed on the fields: (`user_id`)")
         for row in rows:
             self.rows[row.user_id] = row
         if self.commit_then_drop:
-            raise ConnectionError("connection reset after commit")
+            raise httpx.ReadError("connection reset after commit")
         return len(rows)
 
     async def update(self, where: dict[str, str], data: dict[str, object]) -> _UserRow:
@@ -114,9 +127,13 @@ class _Tx:
 
 class _Db:
     def __init__(
-        self, teams: list[LiteLLM_TeamTable], fail_ids: frozenset[str] = frozenset(), commit_then_drop: bool = False
+        self,
+        teams: list[LiteLLM_TeamTable],
+        fail_ids: frozenset[str] = frozenset(),
+        commit_then_drop: bool = False,
+        raced_ids: frozenset[str] = frozenset(),
     ) -> None:
-        self.litellm_usertable = _UserTable(fail_ids, commit_then_drop)
+        self.litellm_usertable = _UserTable(fail_ids, commit_then_drop, raced_ids)
         self.litellm_teamtable = _TeamTable(teams)
         self.litellm_teammembership = _MembershipTable()
 
@@ -127,8 +144,9 @@ class _FakePrisma:
         teams: list[LiteLLM_TeamTable] | None = None,
         fail_ids: frozenset[str] = frozenset(),
         commit_then_drop: bool = False,
+        raced_ids: frozenset[str] = frozenset(),
     ) -> None:
-        self.db = _Db(teams or [], fail_ids, commit_then_drop)
+        self.db = _Db(teams or [], fail_ids, commit_then_drop, raced_ids)
         self.tx_count = 0
         self.locks: list[str] = []
 
@@ -276,6 +294,17 @@ async def test_insert_that_committed_but_lost_its_response_still_counts_as_creat
     assert [m.user_id for m in prisma.db.litellm_teamtable.rows["t1"].members_with_roles] == ["u1"]
 
 
+@pytest.mark.asyncio
+async def test_user_id_taken_by_a_concurrent_request_is_not_claimed_by_this_batch():
+    prisma = _FakePrisma(teams=[_team("t1")], raced_ids=frozenset({"u1"}))
+    response = await _run(prisma, [{"user_id": "u1", "teams": ["t1"]}, {"user_id": "u2", "teams": ["t1"]}])
+
+    assert [r.success for r in response.results] == [False, True]
+    assert "User id=u1 already exists" in (response.results[0].error or "")
+    assert prisma.db.litellm_usertable.rows["u1"].user_email == "u1@other-request.example"
+    assert [m.user_id for m in prisma.db.litellm_teamtable.rows["t1"].members_with_roles] == ["u2"]
+
+
 @pytest.mark.asyncio
 async def test_team_write_failure_keeps_user_and_reports_it_on_the_row():
     prisma = _FakePrisma(teams=[_team("t1"), _team("t2")])