Merge dc44926af1 into 44b85d0c5c

litellm/proxy/management_endpoints/key_management_endpoints.py

@@ -5482,7 +5482,9 @@ async def generate_key_helper_fn(
                 if query_type == "insert_data":
                     user_row = cast(  # cast-ok: table_name="user" is the insert_data branch returning the user row
                         "prisma_models.LiteLLM_UserTable | None",
-                        await prisma_client.insert_data(data=user_data, table_name="user"),
+                        await prisma_client.insert_data(
+                            data=user_data, table_name="user", allow_existing_user=request_type != "user"
+                        ),
                     )
 
                     if user_row is None:

litellm/proxy/management_endpoints/scim/scim_v2.py

@@ -1294,7 +1294,7 @@ async def _create_user_if_not_exists(user_id: str, created_via: str = "scim_grou
         if litellm.default_internal_user_params:
             default_role = litellm.default_internal_user_params.get("user_role")
 
-        new_user_request: Final = NewUserRequest(
+        new_user_request: Final = SCIMNewUserRequest(
             user_id=user_id,
             user_email=user_id,  # We don't have email from group membership
             user_alias=None,
@@ -1873,7 +1873,7 @@ async def create_user(
         admin_group: Final = await _get_scim_admin_group()
         resolved_role: Final = _resolve_scim_user_role(user.groups or [], admin_group, default_role)
 
-        new_user_request: Final = NewUserRequest(
+        new_user_request: Final = SCIMNewUserRequest(
             user_id=user_id,
             user_email=user_data["user_email"],
             user_alias=user_data["user_alias"],

litellm/proxy/utils.py

@@ -310,6 +310,7 @@ if TYPE_CHECKING:
     from prisma import models as prisma_models
     from prisma.actions import LiteLLM_DeprecatedVerificationTokenActions
     from prisma.client import TransactionManager
+    from prisma.errors import UniqueViolationError
     from prisma.models import LiteLLM_DeprecatedVerificationToken
     from prisma.types import HttpConfig, LiteLLM_VerificationTokenInclude
 
@@ -5526,6 +5527,25 @@ class _StaleReadEngine:
         return self.wrapper is current and self.generation == current.engine_generation
 
 
+def _user_unique_conflict_message(metadata: object, user_id: object) -> str:
+    match metadata:
+        case {"target": ["user_id"]}:
+            return f"User with id {user_id} already exists"
+        case {"target": ["sso_user_id"]}:
+            return "User with this sso_user_id already exists"
+        case _:
+            return "User conflicts with an existing unique value"
+
+
+def _raise_user_unique_conflict(error: "UniqueViolationError", user_id: object, *, allow_existing: bool) -> NoReturn:
+    if allow_existing:
+        raise error
+    raise HTTPException(
+        status_code=409,
+        detail={"error": _user_unique_conflict_message(error.meta, user_id)},
+    ) from error
+
+
 class PrismaClient:
     spend_log_transactions: list = []
     _spend_log_transactions_lock = asyncio.Lock()
@@ -6413,12 +6433,15 @@ class PrismaClient:
         max_tries=3,  # maximum number of retries
         max_time=10,  # maximum total time to retry for
         on_backoff=on_backoff,  # specifying the function to call on backoff
+        giveup=lambda exc: isinstance(exc, HTTPException) and exc.status_code == 409,
     )
     @log_db_metrics
     async def insert_data(
         self,
         data: Mapping[str, object],
         table_name: Literal["user", "key", "config", "spend", "team", "user_notification"],
+        *,
+        allow_existing_user: bool = True,
     ):
         """
         Add a key to the database. If it already exists, do nothing.
@@ -6454,14 +6477,20 @@ class PrismaClient:
                 verbose_proxy_logger.info("Data Inserted into Keys Table")
                 return new_verification_token
             elif table_name == "user":
+                from prisma.errors import UniqueViolationError
+
                 db_data = self.jsonify_object(data=data)
                 try:
-                    new_user_row: Final = await UserRepository(self).upsert_user(db_data)
+                    new_user_row: Final = await UserRepository(self).upsert_user(
+                        db_data, allow_existing=allow_existing_user
+                    )
                     if new_user_row is None:
                         raise HTTPException(
                             status_code=403,
                             detail="SSO for more than 5 users requires a LiteLLM Enterprise license. Existing users can still sign in.",
                         )
+                except UniqueViolationError as e:
+                    _raise_user_unique_conflict(e, db_data["user_id"], allow_existing=allow_existing_user)
                 except Exception as e:
                     if (
                         "Foreign key constraint failed on the field: `LiteLLM_UserTable_organization_id_fkey (index)`"

litellm/repositories/user_repository.py

@@ -168,7 +168,9 @@ class UserRepository(BaseRepository[LiteLLM_UserTable]):
         """
         return await _count_billable_users(self.table)
 
-    async def upsert_user(self, data: Mapping[str, object]) -> "prisma_models.LiteLLM_UserTable | None":
+    async def upsert_user(
+        self, data: Mapping[str, object], *, allow_existing: bool = True
+    ) -> "prisma_models.LiteLLM_UserTable | None":
         """Return the user row, or None when a new free SSO user would exceed capacity."""
         scope: Final = _FREE_SSO_CREATION.get()
         if (
@@ -176,6 +178,8 @@ class UserRepository(BaseRepository[LiteLLM_UserTable]):
             or (scope[0] is not None and scope[0] != data["user_id"])
             or scope[1] is not asyncio.current_task()
         ):
+            if not allow_existing:
+                return await self.table.create(data=dict(data))
             return await self.table.upsert(
                 where={"user_id": data["user_id"]}, data={"create": dict(data), "update": {}}
             )

litellm/types/proxy/management_endpoints/internal_user_endpoints.py

@@ -120,6 +120,7 @@ class BulkNewUserItem(NewUserRequest):
 
     model_config = ConfigDict(extra="forbid", protected_namespaces=())
 
+    user_id: str | None = None
     auto_create_key: bool = False
 
     @field_validator("send_invite_email")

litellm/types/proxy/management_endpoints/scim_v2.py

@@ -13,6 +13,7 @@ from pydantic import (
 from pydantic_core.core_schema import SerializerFunctionWrapHandler
 
 from litellm.types.llms.base import LiteLLMBaseModel
+from litellm.types.proxy.management_endpoints.user_requests import NewUserRequest
 
 SCIM_ENTERPRISE_USER_SCHEMA: Final = "urn:ietf:params:scim:schemas:extension:enterprise:2.0:User"
 SCIM_ENTERPRISE_METADATA_KEY: Final = "scim_enterprise"
@@ -26,6 +27,14 @@ _SCIM_SERIALIZED_FIELDS_ADAPTER: Final = TypeAdapter(
 )
 
 
+class SCIMNewUserRequest(NewUserRequest):
+    """Preserve opaque directory IDs when provisioning SCIM users and group members."""
+
+    model_config = ConfigDict(frozen=True)
+
+    user_id: str | None = None
+
+
 class LiteLLM_UserScimMetadata(LiteLLMBaseModel):
     """
     Scim metadata stored in LiteLLM_UserTable.metadata

litellm/types/proxy/management_endpoints/user_requests.py

@@ -1,7 +1,7 @@
 from collections.abc import Mapping
-from typing import Literal
+from typing import Annotated, Final, Literal
 
-from pydantic import Field, field_validator, model_validator
+from pydantic import AfterValidator, Field, field_validator, model_validator
 
 from litellm.models.user import LiteLLM_UserTable
 from litellm.types.proxy.auth.user_roles import LitellmUserRoles
@@ -15,7 +15,17 @@ class NewUserRequestTeam(LiteLLMPydanticObjectBase):
     user_role: Literal["user", "admin"] = "user"
 
 
+def _validate_custom_user_id(value: str) -> str:
+    normalized: Final = value.strip()
+    if not normalized:
+        raise ValueError("user_id must not be empty or contain only whitespace")
+    if "\x00" in normalized:
+        raise ValueError("user_id must not contain NUL characters")
+    return normalized
+
+
 class NewUserRequest(GenerateRequestBase):
+    user_id: Annotated[str, AfterValidator(_validate_custom_user_id)] | None = None
     tpm_limit: int | None = Field(default=None, ge=0)
     rpm_limit: int | None = Field(default=None, ge=0)
     max_budget: float | None = None

tests/unit/proxy/management_endpoints/scim/test_scim_v2_endpoints.py

@@ -411,10 +411,11 @@ async def test_create_user_without_groups_defers_to_default_team(mocker: MockerF
 
 
 @pytest.mark.asyncio
-async def test_create_user_with_groups_keeps_idp_teams(mocker: MockerFixture, monkeypatch):
+@pytest.mark.parametrize("user_id", ["new-user", " employee-42 "])
+async def test_create_user_with_groups_keeps_idp_teams(mocker: MockerFixture, monkeypatch, user_id: str):
     scim_user = SCIMUser(
         schemas=["urn:ietf:params:scim:schemas:core:2.0:User"],
-        userName="new-user",
+        userName=user_id,
         emails=[SCIMUserEmail(value="new@example.com")],
         groups=[SCIMUserGroup(value="idp-team")],
     )
@@ -428,10 +429,12 @@ async def test_create_user_with_groups_keeps_idp_teams(mocker: MockerFixture, mo
     await create_user(user=scim_user)
 
     assert new_user_mock.call_args.kwargs["data"].teams == ["idp-team"]
+    assert new_user_mock.call_args.kwargs["data"].user_id == user_id
 
 
 @pytest.mark.asyncio
-async def test_create_user_if_not_exists_defers_to_default_team(mocker: MockerFixture, monkeypatch):
+@pytest.mark.parametrize("user_id", ["group-user", " employee-42 "])
+async def test_create_user_if_not_exists_defers_to_default_team(mocker: MockerFixture, monkeypatch, user_id: str):
     monkeypatch.setattr(
         "litellm.default_internal_user_params",
         {"teams": [{"team_id": "default-team", "max_budget_in_team": 25}]},
@@ -439,14 +442,15 @@ async def test_create_user_if_not_exists_defers_to_default_team(mocker: MockerFi
     )
     new_user_mock = mocker.patch(  # test-quality-ok: new_user is imported inside the helper, not injectable
         "litellm.proxy.management_endpoints.internal_user_endpoints.new_user",
-        AsyncMock(return_value=NewUserResponse(user_id="group-user", key="k")),
+        AsyncMock(return_value=NewUserResponse(user_id=user_id, key="k")),
     )
 
-    created = await _create_user_if_not_exists(user_id="group-user")
+    created = await _create_user_if_not_exists(user_id=user_id)
 
     assert created is not None
     assert new_user_mock.call_args.kwargs["data"].teams is None
     assert new_user_mock.call_args.kwargs["user_api_key_dict"] == UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN)
+    assert new_user_mock.call_args.kwargs["data"].user_id == user_id
 
 
 @pytest.mark.asyncio

tests/unit/proxy/management_endpoints/test_internal_user_endpoints.py

@@ -13,6 +13,7 @@ import httpx
 import pytest
 from fastapi import HTTPException
 from fastapi.testclient import TestClient
+from pydantic import ValidationError
 from pytest_mock import MockerFixture
 
 from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
@@ -48,6 +49,51 @@ from tests.unit.proxy.management_endpoints.jwt_key_mapping_doubles import (
 
 client = TestClient(app)
 
+
+@pytest.mark.parametrize("user_id", ["", " \t\n", "user\x00id", 123, True, [], {}])
+def test_new_user_rejects_invalid_custom_id(user_id: object) -> None:
+    with pytest.raises(ValidationError) as exc:
+        NewUserRequest.model_validate({"user_id": user_id})
+    assert exc.value.errors()[0]["loc"] == ("user_id",)
+
+
+@pytest.mark.parametrize("user_id", [None, "employee-42", "社員@example.test", "  employee-42  "])
+def test_new_user_normalizes_custom_id_and_preserves_auto_generation(user_id: str | None) -> None:
+    request: Final = NewUserRequest(user_id=user_id)
+    assert request.user_id == (user_id.strip() if user_id is not None else None)
+
+
+@pytest.mark.asyncio
+async def test_new_user_returns_conflict_when_id_is_taken_after_duplicate_check(mocker: MockerFixture) -> None:
+    from prisma.errors import UniqueViolationError
+
+    from litellm.proxy import proxy_server
+    from litellm.proxy.utils import PrismaClient
+
+    table: Final = SimpleNamespace(
+        find_first=AsyncMock(return_value=None),
+        count=AsyncMock(return_value=0),
+        create=AsyncMock(side_effect=UniqueViolationError(
+            {"user_facing_error": {"meta": {"target": ["user_id"]}}},
+            message="Unique constraint failed on user_id",
+        )),
+        upsert=AsyncMock(return_value=LiteLLM_UserTable(user_id="employee-42")),
+    )
+    mocker.patch("prisma.Prisma", return_value=SimpleNamespace(litellm_usertable=table))
+    database: Final = PrismaClient(
+        database_url="postgresql://unused:unused@localhost/unused",
+        proxy_logging_obj=proxy_server.proxy_logging_obj,
+    )
+    mocker.patch("litellm.proxy.proxy_server.prisma_client", database)
+    with pytest.raises(ProxyException) as exc:
+        await new_user(
+            data=NewUserRequest(user_id="employee-42", auto_create_key=False, send_invite_email=False),
+            user_api_key_dict=UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN),
+        )
+    assert str(exc.value.code) == "409"
+    assert "employee-42" in exc.value.message
+
+
 if TYPE_CHECKING:
     from litellm.caching.redis_cache import RedisCache
     from litellm.proxy.utils import PrismaClient
@@ -4852,7 +4898,9 @@ async def test_user_new_persists_model_max_budget(monkeypatch, model_max_budget,
         models = []
 
     class _FakePrisma:
-        async def insert_data(self, data, table_name):
+        async def insert_data(
+            self, data: Mapping[str, object], table_name: str, *, allow_existing_user: bool = True
+        ) -> _FakeUserRow | SimpleNamespace:
             if table_name == "user":
                 captured["user_data"] = dict(data)
                 return _FakeUserRow()

tests/unit/proxy/management_endpoints/test_ui_sso.py

@@ -22,6 +22,7 @@ from litellm.proxy._types import (
     LiteLLM_JWTAuth,
     LiteLLM_UserTable,
     NewUserResponse,
+    SSOUserDefinedValues,
     ProxyException,
 )
 from litellm.proxy.auth.handle_jwt import JWTHandler
@@ -52,6 +53,62 @@ from litellm.types.proxy.management_endpoints.ui_sso import (
 from tests._master_key import MASTER_KEY
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    ("billable_users", "premium", "admitted"),
+    [(4, False, True), (5, False, False), (26, False, False), (26, True, True)],
+)
+async def test_sso_auto_creation_enforces_free_user_cap_through_new_user(
+    monkeypatch: pytest.MonkeyPatch, billable_users: int, premium: bool, admitted: bool
+) -> None:
+    from litellm.proxy import proxy_server
+    from litellm.proxy.utils import PrismaClient
+
+    row: Final = LiteLLM_UserTable(user_id="new-sso-user", user_email="new-sso-user@example.test")
+    table: Final = SimpleNamespace(
+        find_first=AsyncMock(return_value=None),
+        find_unique=AsyncMock(return_value=None),
+        count=AsyncMock(side_effect=lambda where=None: billable_users if where is None else 0),
+        create=AsyncMock(return_value=row),
+        upsert=AsyncMock(return_value=row),
+    )
+    transaction: Final = MagicMock()
+    transaction.__aenter__ = AsyncMock(
+        return_value=SimpleNamespace(litellm_usertable=table, execute_raw=AsyncMock(), query_raw=AsyncMock())
+    )
+    transaction.__aexit__ = AsyncMock(return_value=None)
+    monkeypatch.setattr(
+        "prisma.Prisma",
+        lambda **kwargs: SimpleNamespace(litellm_usertable=table, tx=MagicMock(return_value=transaction)),
+    )
+    database: Final = PrismaClient(
+        database_url="postgresql://unused:unused@localhost/unused",
+        proxy_logging_obj=proxy_server.proxy_logging_obj,
+    )
+    monkeypatch.setattr(proxy_server, "prisma_client", database)
+    monkeypatch.setattr(proxy_server, "premium_user", premium)
+    monkeypatch.setattr("asyncio.sleep", AsyncMock())
+    values: Final[SSOUserDefinedValues] = {
+        "models": [],
+        "user_id": row.user_id,
+        "user_email": row.user_email,
+        "user_role": "internal_user_viewer",
+        "max_budget": None,
+        "budget_duration": None,
+    }
+    if admitted:
+        response: Final = await insert_sso_user(CustomOpenID(id=row.user_id, team_ids=[]), values)
+        assert response.user_id == row.user_id
+        assert table.create.await_count + table.upsert.await_count == 1
+    else:
+        with pytest.raises(ProxyException) as error:
+            await insert_sso_user(CustomOpenID(id=row.user_id, team_ids=[]), values)
+        assert str(error.value.code) == "403"
+        assert "SSO for more than 5 users" in error.value.message
+        table.create.assert_not_awaited()
+        table.upsert.assert_not_awaited()
+
+
 def _sso_http_handler(route: Callable[[httpx.Request], httpx.Response] | None = None) -> AsyncHTTPHandler:
     """litellm's SSO httpx client answered in-process, by ``route`` or by the active respx mock."""
     handler = AsyncHTTPHandler()

tests/unit/proxy/management_helpers/test_bulk_user_creation.py

@@ -191,6 +191,19 @@ async def _run(prisma, users, caller=ADMIN, license=None, generate_key=_no_keys)
     )
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize("user_id", [None, "", "  employee-42  ", " \t"])
+async def test_bulk_keeps_existing_user_id_generation_and_whitespace_behavior(user_id: str | None) -> None:
+    prisma: Final = _FakePrisma()
+    response: Final = await _run(prisma, [{"user_id": user_id}])
+    result: Final = response.data[0]
+    assert result.success
+    assert result.user_id
+    assert result.user_id in prisma.db.litellm_usertable.rows
+    if user_id:
+        assert result.user_id == user_id
+
+
 @pytest.mark.asyncio
 async def test_creates_users_and_team_membership_in_every_store():
     prisma = _FakePrisma(teams=[_team("t1", [Member(user_id="existing", role="admin")]), _team("t2")])

tests/unit/proxy/utils/prisma_and_spend/test_prisma_client_writes.py

@@ -18,8 +18,10 @@ from types import SimpleNamespace
 from typing import Final
 from unittest.mock import AsyncMock, MagicMock, patch
 
+import httpx
 import pytest
 from fastapi import HTTPException
+from prisma.errors import UniqueViolationError
 
 from litellm._service_logger import ServiceTypes
 from litellm.models.user import LiteLLM_UserTable
@@ -192,6 +194,62 @@ async def test_insert_data_team_serializes_members(prisma_client: PrismaClient)
     assert create_payload["team_id"] == "t1"
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    ("metadata", "message"),
+    [
+        ({"target": ["user_id"]}, "User with id requested-user already exists"),
+        ({"target": ["sso_user_id"]}, "User with this sso_user_id already exists"),
+        (None, "User conflicts with an existing unique value"),
+        ({"target": ["unknown_field"]}, "User conflicts with an existing unique value"),
+    ],
+)
+async def test_user_conflict_names_constraint_and_does_not_retry(
+    prisma_client: PrismaClient,
+    metadata: Mapping[str, object] | None,
+    message: str,
+) -> None:
+    create: Final = AsyncMock(side_effect=UniqueViolationError({"user_facing_error": {"meta": metadata}}))
+    prisma_client.db.litellm_usertable.create = create
+    with pytest.raises(HTTPException) as error:
+        await prisma_client.insert_data(
+            data={"user_id": "requested-user", "sso_user_id": "sso-subject"},
+            table_name="user", allow_existing_user=False,
+        )
+    assert error.value.status_code == 409
+    assert error.value.detail == {"error": message}
+    assert create.await_count == 1
+
+
+@pytest.mark.asyncio
+async def test_user_upsert_recovers_from_concurrent_creation(
+    prisma_client: PrismaClient, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    row: Final = LiteLLM_UserTable(user_id="requested-user")
+    conflict: Final = UniqueViolationError({"user_facing_error": {"meta": {"target": ["user_id"]}}})
+    prisma_client.db.litellm_usertable.upsert = AsyncMock(side_effect=[conflict, row])
+    monkeypatch.setattr("asyncio.sleep", AsyncMock())
+
+    result: Final = await prisma_client.insert_data(data={"user_id": row.user_id}, table_name="user")
+
+    assert result is row
+
+
+@pytest.mark.asyncio
+async def test_user_insert_still_retries_transient_connection_failure(
+    prisma_client: PrismaClient, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    row: Final = LiteLLM_UserTable(user_id="requested-user")
+    create: Final = AsyncMock(side_effect=[httpx.ConnectError("connection unavailable"), row])
+    prisma_client.db.litellm_usertable.create = create
+    monkeypatch.setattr("asyncio.sleep", AsyncMock())
+    result: Final = await prisma_client.insert_data(
+        data={"user_id": row.user_id}, table_name="user", allow_existing_user=False
+    )
+    assert result is row
+    assert create.await_count == 2
+
+
 @pytest.mark.asyncio
 async def test_insert_data_user_organization_fk_raises_400(
     prisma_client: PrismaClient,