fix(sso): adopt the concurrently inserted user on a first-login duplicate race

litellm/proxy/management_endpoints/ui_sso.py

@@ -89,7 +89,7 @@ from litellm.proxy._types import (
     TeamMemberAddRequest,
     UserAPIKeyAuth,
 )
-from litellm.proxy.auth.auth_checks import ExperimentalUIJWTToken, get_user_object
+from litellm.proxy.auth.auth_checks import ExperimentalUIJWTToken, forget_missing_user, get_user_object
 from litellm.proxy.auth.auth_utils import (  # noqa: F401  # legacy module exports
     _get_request_ip_address,  # pyright: ignore[reportPrivateUsage,reportUnusedImport]  # backwards-compatible package export
     get_request_ip_address,
@@ -135,6 +135,7 @@ from litellm.proxy.utils import (
     get_custom_url,
     get_server_root_path,
 )
+from litellm.repositories.base_repository import is_unique_violation
 from litellm.repositories.prisma_protocols import TableActions
 from litellm.repositories.table_repositories import SSOConfigRepository
 from litellm.repositories.team_repository import TeamRepository
@@ -2677,6 +2678,48 @@ async def insert_sso_user(
     return response
 
 
+def _is_duplicate_user_error(error: BaseException | None) -> bool:
+    if error is None:
+        return False
+    if isinstance(error, ProxyException) and error.code == str(status.HTTP_409_CONFLICT):
+        return True
+    if isinstance(error, HTTPException) and error.status_code == status.HTTP_409_CONFLICT:
+        return True
+    if is_unique_violation(error):
+        return True
+    return _is_duplicate_user_error(error.__cause__ or error.__context__)
+
+
+async def _concurrently_inserted_sso_user(
+    error: Exception,
+    user_defined_values: SSOUserDefinedValues | None,
+    prisma_client: PrismaClient,
+) -> LiteLLM_UserTable | None:
+    if user_defined_values is None or not _is_duplicate_user_error(error):
+        return None
+    user_id: Final = user_defined_values["user_id"]
+    existing: Final = await UserRepository(prisma_client, use_writer=True).find_by_id(user_id)
+    if existing is None or existing.sso_user_id != user_id:
+        return None
+    forget_missing_user(user_id)
+    return existing
+
+
+async def _insert_sso_user_or_adopt_concurrent_insert(
+    insert: Callable[[], Awaitable[NewUserResponse]],
+    user_defined_values: SSOUserDefinedValues | None,
+    prisma_client: PrismaClient,
+) -> NewUserResponse | LiteLLM_UserTable:
+    try:
+        return await insert()
+    except Exception as error:
+        concurrent: Final = await _concurrently_inserted_sso_user(error, user_defined_values, prisma_client)
+        if concurrent is None:
+            raise
+        verbose_proxy_logger.info("SSO user %s was inserted by a concurrent login, using that row", concurrent.user_id)
+        return concurrent
+
+
 @router.get(
     "/sso/get/ui_settings",
     tags=["experimental"],
@@ -3307,10 +3350,10 @@ class SSOAuthenticationHandler:
                 )
             else:
                 verbose_proxy_logger.info("user not in DB, inserting user into LiteLLM DB")
-                # user not in DB, insert User into LiteLLM DB
-                user_info = await insert_sso_user(
-                    result_openid=result,
+                user_info = await _insert_sso_user_or_adopt_concurrent_insert(
+                    insert=lambda: insert_sso_user(result_openid=result, user_defined_values=user_defined_values),
                     user_defined_values=user_defined_values,
+                    prisma_client=prisma_client,
                 )
             return user_info
         except Exception as e:

tests/integration/authorization/test_cli_sso_login_ui_disabled.py

@@ -32,6 +32,7 @@ from tests.integration._support.client import (
     string_value,
 )
 from tests.integration._support.database import read_rows
+from tests.integration._support.database_relay import HeldStatementRelay, held_statement_relay
 from tests.integration._support.process import OwnedProxy, group_members, owned_proxy_process
 from tests.integration._support.provider import SharedProvider
 from tests.integration._support.wire import Reply, Request, Wire, wire_server
@@ -49,6 +50,9 @@ CLI_SUCCESS_PAGE_TITLE: Final = "<title>CLI Authentication Successful - LiteLLM<
 SESSION_GONE: Final = "CLI login session not found or expired"
 SUBJECT: Final = "cli-sso-audit-subject"
 SUBJECT_EMAIL: Final = "cli-sso-audit-subject@example.com"
+DUPLICATE_USER_ID_CHECK: Final = (
+    b'FROM "public"."LiteLLM_UserTable" WHERE "public"."LiteLLM_UserTable"."user_id" = $1 LIMIT'
+)
 CLIENT_ID: Final = "integration-oidc-client"
 CLIENT_SECRET: Final = "integration-oidc-secret"
 MESSAGE_MODEL: Final = "anthropic/claude-haiku-4-5"
@@ -296,8 +300,11 @@ def _send_message(proxy: Gateway, provider: SharedProvider, key: str) -> str:
     return message_id
 
 
-def _user_rows() -> Sequence[Mapping[str, JsonValue]]:
-    return read_rows('SELECT user_id, user_email FROM "LiteLLM_UserTable" WHERE user_id = %s', (SUBJECT,))
+def _user_rows(subject: str = SUBJECT) -> Sequence[Mapping[str, JsonValue]]:
+    return read_rows(
+        'SELECT user_id, user_email FROM "LiteLLM_UserTable" WHERE user_id = %s OR user_email = %s',
+        (subject, f"{subject}@example.com"),
+    )
 
 
 def _worker_pids(root_pid: int) -> frozenset[int]:
@@ -461,8 +468,32 @@ def test_cli_link_is_reentrant_until_the_poll_consumes_the_session(ui_disabled:
     assert consumed.status_code == 400 and SESSION_GONE in consumed.text, f"{consumed.status_code} {consumed.text}"
 
 
-def test_burst_of_logins_recovers_from_an_idp_token_outage(ui_disabled: OwnedProxy, idp: Idp) -> None:
-    proxy: Final = ui_disabled.gateway
+def test_burst_of_logins_recovers_from_an_idp_token_outage(idp: Idp, provider: SharedProvider, tmp_path: Path) -> None:
+    subject: Final = f"cli-sso-burst-{uuid.uuid4().hex[:12]}"
+    writer_url: Final = os.environ.get("INTEGRATION_PROXY_DATABASE_URL") or os.environ["DATABASE_URL"]
+    with (
+        held_statement_relay(writer_url, DUPLICATE_USER_ID_CHECK) as (relay, relayed_url),
+        gateway_from_environment() as rig,
+        owned_proxy_process(
+            rig,
+            tmp_path,
+            {
+                "DISABLE_ADMIN_UI": "true",
+                "DATABASE_URL": relayed_url,
+                "PRISMA_HEALTH_WATCHDOG_ENABLED": "false",
+                **_sso_environment(idp.wire.url),
+            },
+            config=_gateway_enabled_config(tmp_path),
+            remove_environment=("PROXY_BASE_URL",),
+            workers=2,
+        ) as owned,
+    ):
+        _burst_of_logins_after_an_idp_token_outage(owned.gateway, idp, provider, relay, subject)
+
+
+def _burst_of_logins_after_an_idp_token_outage(
+    proxy: Gateway, idp: Idp, provider: SharedProvider, relay: HeldStatementRelay, subject: str
+) -> None:
     sessions: Final = tuple(_start_lite_login(proxy) for _ in range(BURST))
 
     def failed_exchange(session: CliSession) -> httpx.Response:
@@ -492,16 +523,25 @@ def test_burst_of_logins_recovers_from_an_idp_token_outage(ui_disabled: OwnedPro
 
     def recovered_login(session: CliSession) -> str:
         with _browser() as browser:
-            _sign_in(proxy, idp, browser, session)
-        return _ready_key(proxy, session)
+            _sign_in(proxy, idp, browser, session, subject=subject)
+        return _ready_key(proxy, session, subject=subject)
 
     with ThreadPoolExecutor(max_workers=BURST) as pool:
-        keys: Final = tuple(pool.map(recovered_login, sessions))
+        logins: Final = tuple(pool.submit(recovered_login, session) for session in sessions)
+        try:
+            assert relay.held.wait(60), "no first-time login reached the duplicate user id check"
+            eventually(lambda: _user_rows(subject), bool)
+        finally:
+            relay.release()
+        keys: Final = tuple(login.result() for login in logins)
     assert len(set(keys)) == BURST, keys
     for session in sessions:
         again: Final = _poll(proxy, session)
         assert again.status_code == 400 and SESSION_GONE in again.text, f"{again.status_code} {again.text}"
-    assert [row["user_id"] for row in _user_rows()] == [SUBJECT]
+    users: Final = _user_rows(subject)
+    assert [(row["user_id"], row["user_email"]) for row in users] == [(subject, f"{subject}@example.com")], users
+    for key in keys:
+        _send_message(proxy, provider, key)
 
 
 def test_login_survives_a_worker_kill(idp: Idp, tmp_path: Path) -> None:

tests/unit/proxy/management_endpoints/test_ui_sso.py

@@ -13,7 +13,7 @@ from fastapi import HTTPException, Request
 
 import litellm
 from litellm._uuid import uuid
-from litellm.proxy._types import LiteLLM_UserTable, NewUserResponse
+from litellm.proxy._types import LiteLLM_UserTable, NewUserResponse, ProxyException, SSOUserDefinedValues
 from litellm.proxy.auth.handle_jwt import JWTHandler
 from litellm.proxy.management_endpoints.sso import CustomMicrosoftSSO
 from litellm.proxy.management_endpoints.types import CustomOpenID
@@ -21,6 +21,7 @@ from litellm.proxy.management_endpoints.ui_sso import (
     GoogleSSOHandler,
     MicrosoftSSOHandler,
     SSOAuthenticationHandler,
+    _insert_sso_user_or_adopt_concurrent_insert,
     _setup_team_mappings,
     _sync_user_role_from_jwt_role_map,
     normalize_email,
@@ -1181,6 +1182,125 @@ async def test_insert_sso_user_sets_user_alias_from_display_name():
     assert new_user_request.user_alias == "Doe, Jane"
 
 
+RACE_SUBJECT = "sso-race-subject"
+RACE_VALUES: SSOUserDefinedValues = {
+    "models": [],
+    "user_id": RACE_SUBJECT,
+    "user_email": "sso-race-subject@example.com",
+    "max_budget": None,
+    "user_role": "internal_user",
+    "budget_duration": None,
+}
+
+
+class _UniqueViolation(Exception):
+    code = "P2002"
+
+
+def _caused_by(error: ProxyException, cause: Exception) -> ProxyException:
+    error.__cause__ = cause
+    return error
+
+
+def _writer_with_row(row: dict | None) -> MagicMock:
+    prisma = MagicMock()
+    prisma.writer_db.litellm_usertable.find_unique = AsyncMock(return_value=row)
+    return prisma
+
+
+def _failing_insert(error: Exception):
+    async def insert() -> NewUserResponse:
+        raise error
+
+    return insert
+
+
+DUPLICATE_ID = ProxyException(
+    message="User with id sso-race-subject already exists", type="bad_request", param="user_id", code=409
+)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "conflict",
+    (
+        DUPLICATE_ID,
+        _caused_by(
+            ProxyException(message="insert failed", type="internal_server_error", param=None, code=500),
+            _UniqueViolation("Unique constraint failed on the fields: (`user_id`)"),
+        ),
+    ),
+    ids=("duplicate-id-409", "unique-violation"),
+)
+async def test_concurrent_first_sso_login_adopts_the_row_the_winning_login_inserted(conflict: ProxyException):
+    prisma = _writer_with_row({"user_id": RACE_SUBJECT, "sso_user_id": RACE_SUBJECT, "user_role": "internal_user"})
+
+    user = await _insert_sso_user_or_adopt_concurrent_insert(
+        user_defined_values=RACE_VALUES, prisma_client=prisma, insert=_failing_insert(conflict)
+    )
+
+    assert isinstance(user, LiteLLM_UserTable), user
+    assert (user.user_id, user.sso_user_id, user.user_role) == (RACE_SUBJECT, RACE_SUBJECT, "internal_user")
+    prisma.writer_db.litellm_usertable.find_unique.assert_awaited_once_with(where={"user_id": RACE_SUBJECT})
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "row",
+    (
+        None,
+        {"user_id": RACE_SUBJECT, "sso_user_id": "another-sso-subject"},
+        {"user_id": RACE_SUBJECT, "sso_user_id": None},
+    ),
+    ids=("only-the-email-is-taken", "row-bound-to-another-identity", "row-not-bound-to-sso"),
+)
+async def test_concurrent_first_sso_login_refuses_a_conflicting_row_of_another_identity(row: dict | None):
+    conflict = ProxyException(
+        message="User with email already exists", type="bad_request", param="user_email", code=409
+    )
+    prisma = _writer_with_row(row)
+
+    with pytest.raises(ProxyException) as raised:
+        await _insert_sso_user_or_adopt_concurrent_insert(
+            user_defined_values=RACE_VALUES, prisma_client=prisma, insert=_failing_insert(conflict)
+        )
+
+    assert raised.value is conflict
+
+
+@pytest.mark.asyncio
+async def test_first_sso_login_reraises_an_unrelated_database_failure_without_readback():
+    failure = _caused_by(
+        ProxyException(message="insert failed", type="internal_server_error", param=None, code=500),
+        ConnectionError("server closed the connection unexpectedly"),
+    )
+    prisma = _writer_with_row({"user_id": RACE_SUBJECT, "sso_user_id": RACE_SUBJECT})
+
+    with pytest.raises(ProxyException) as raised:
+        await _insert_sso_user_or_adopt_concurrent_insert(
+            user_defined_values=RACE_VALUES, prisma_client=prisma, insert=_failing_insert(failure)
+        )
+
+    assert raised.value is failure
+    prisma.writer_db.litellm_usertable.find_unique.assert_not_awaited()
+
+
+@pytest.mark.asyncio
+async def test_first_sso_login_returns_the_inserted_user_without_readback():
+    created = NewUserResponse(user_id=RACE_SUBJECT, key="sk-created", teams=None)
+    prisma = _writer_with_row(None)
+
+    async def insert() -> NewUserResponse:
+        return created
+
+    user = await _insert_sso_user_or_adopt_concurrent_insert(
+        user_defined_values=RACE_VALUES, prisma_client=prisma, insert=insert
+    )
+
+    assert user is created
+    prisma.writer_db.litellm_usertable.find_unique.assert_not_awaited()
+
+
 @pytest.mark.asyncio
 async def test_upsert_sso_user_does_not_update_invalid_role():
     """