diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index a4d2977a971..d115e486ad7 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -18,6 +18,7 @@ import re import secrets from collections.abc import Awaitable, Callable, Mapping, Sequence from copy import deepcopy +from dataclasses import dataclass from html import escape from types import MappingProxyType from typing import ( @@ -89,7 +90,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, @@ -110,7 +111,10 @@ from litellm.proxy.common_utils.html_forms.jwt_display_template import ( ) from litellm.proxy.common_utils.html_forms.ui_login import build_ui_login_form from litellm.proxy.common_utils.user_api_key_cache import AUTH_OBJECTS_TARGET, UserApiKeyCache -from litellm.proxy.management_endpoints.internal_user_endpoints import new_user +from litellm.proxy.db.exception_handler import ( + _exception_chain, # pyright: ignore[reportPrivateUsage] # reuse the bounded database exception chain walk +) +from litellm.proxy.management_endpoints.internal_user_endpoints import check_if_default_team_set, new_user from litellm.proxy.management_endpoints.sso import CustomMicrosoftSSO from litellm.proxy.management_endpoints.sso.id_jag_assertion_capture import ( id_jag_assertion_capture_gap, @@ -135,6 +139,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 @@ -2683,6 +2688,74 @@ async def insert_sso_user( return response +def _is_duplicate_user_error(error: BaseException | None) -> bool: + return error is not None and any( + (isinstance(link, ProxyException) and link.code == str(status.HTTP_409_CONFLICT)) + or (isinstance(link, HTTPException) and link.status_code == status.HTTP_409_CONFLICT) + or is_unique_violation(link) + for link in _exception_chain(error) + ) + + +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"] + default_parameters: Final = TypeAdapter[Mapping[str, object]](Mapping[str, object]).validate_python( + getattr(litellm, "default_internal_user_params", None) or {} + ) + required_organizations: Final = TypeAdapter(tuple[str, ...]).validate_python( + () + if user_defined_values["user_role"] == LitellmUserRoles.PROXY_ADMIN + else default_parameters.get("organizations") or () + ) + existing: Final = await UserRepository(prisma_client, use_writer=True).find_by_id( + user_id, include_organization_memberships=bool(required_organizations) + ) + if existing is None or existing.sso_user_id != user_id or existing.user_role != user_defined_values["user_role"]: + return None + default_team_id: Final = TypeAdapter[str | None](str | None).validate_python( + default_parameters.get("team_id") if existing.user_role != LitellmUserRoles.PROXY_ADMIN else None + ) + default_teams: Final = check_if_default_team_set() + required_teams: Final = ( + (default_team_id,) + if default_team_id is not None + else tuple(team if isinstance(team, str) else team.team_id for team in default_teams or ()) + ) + if any(team_id not in existing.teams for team_id in required_teams): + return None + organizations: Final = frozenset(member.organization_id for member in existing.organization_memberships or ()) + if any(organization_id not in organizations for organization_id in required_organizations): + return None + forget_missing_user(user_id) + return existing + + +@dataclass(frozen=True, slots=True) +class SSOInsertFailure: + error: Exception + + +async def _insert_sso_user_or_adopt_concurrent_insert( + insert: Callable[[], Awaitable[NewUserResponse]], + user_defined_values: SSOUserDefinedValues | None, + prisma_client: PrismaClient, +) -> NewUserResponse | LiteLLM_UserTable | SSOInsertFailure: + try: + return await insert() + except Exception as error: # noqa: BLE001 # insert failures remain values for the existing refusal handler + concurrent: Final = await _concurrently_inserted_sso_user(error, user_defined_values, prisma_client) + if concurrent is None: + return SSOInsertFailure(error) + 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"], @@ -3313,11 +3386,17 @@ 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, + outcome: Final = 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, ) + if isinstance(outcome, SSOInsertFailure): + verbose_proxy_logger.exception( + "Error upserting SSO user into LiteLLM DB: %s", outcome.error, exc_info=outcome.error + ) + return user_info + return outcome return user_info except Exception as e: verbose_proxy_logger.exception("Error upserting SSO user into LiteLLM DB: %s", e) diff --git a/litellm/repositories/user_repository.py b/litellm/repositories/user_repository.py index 5ed1574c6bc..537e63ecf18 100644 --- a/litellm/repositories/user_repository.py +++ b/litellm/repositories/user_repository.py @@ -79,8 +79,15 @@ class UserRepository(BaseRepository[LiteLLM_UserTable]): } ) - async def find_by_id(self, id_value: str, id_field: str = "user_id") -> LiteLLM_UserTable | None: - return await super().find_by_id(id_value, id_field) + async def find_by_id( + self, id_value: str, id_field: str = "user_id", *, include_organization_memberships: bool = False + ) -> LiteLLM_UserTable | None: + if not include_organization_memberships: + return await super().find_by_id(id_value, id_field) + record: Final = await self.table.find_unique( + where={id_field: id_value}, include={"organization_memberships": True} + ) + return self._to_model(record) async def find_by_email(self, user_email: str) -> LiteLLM_UserTable | None: """Find a user by email.""" diff --git a/tests/integration/authorization/test_cli_sso_login_ui_disabled.py b/tests/integration/authorization/test_cli_sso_login_ui_disabled.py index a5e9b2752ea..3fd7018eaae 100644 --- a/tests/integration/authorization/test_cli_sso_login_ui_disabled.py +++ b/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 = "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,34 @@ 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"] + reader_url: Final = os.environ.get("INTEGRATION_PROXY_READ_REPLICA_URL") + database_setting: Final = "DATABASE_URL_READ_REPLICA" if reader_url else "DATABASE_URL" + with ( + held_statement_relay(reader_url or 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_setting: 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 +525,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: diff --git a/tests/unit/proxy/management_endpoints/test_ui_sso.py b/tests/unit/proxy/management_endpoints/test_ui_sso.py index 4e128440703..b1899e44887 100644 --- a/tests/unit/proxy/management_endpoints/test_ui_sso.py +++ b/tests/unit/proxy/management_endpoints/test_ui_sso.py @@ -2,30 +2,39 @@ import asyncio import json import logging import os +from collections.abc import Awaitable, Callable, Mapping from contextlib import ExitStack, asynccontextmanager from types import SimpleNamespace +from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import httpx import pytest import respx from fastapi import HTTPException, Request +from pydantic import JsonValue import litellm +import litellm.proxy.auth.auth_checks as auth_checks from litellm._uuid import uuid -from litellm.proxy._types import LiteLLM_UserTable, NewUserResponse +from litellm.caching.dual_cache import LimitedSizeOrderedDict +from litellm.proxy._types import LiteLLM_UserTable, NewUserResponse, ProxyException, SSOUserDefinedValues from litellm.proxy.auth.handle_jwt import JWTHandler +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.proxy.management_endpoints.sso import CustomMicrosoftSSO from litellm.proxy.management_endpoints.types import CustomOpenID from litellm.proxy.management_endpoints.ui_sso import ( GoogleSSOHandler, MicrosoftSSOHandler, SSOAuthenticationHandler, + SSOInsertFailure, + _insert_sso_user_or_adopt_concurrent_insert, _setup_team_mappings, _sync_user_role_from_jwt_role_map, normalize_email, process_sso_jwt_access_token, ) +from litellm.types.proxy.auth.auth_checks import UserNotFoundError from litellm.types.proxy.management_endpoints.ui_sso import ( DefaultTeamSSOParams, MicrosoftGraphAPIUserGroupDirectoryObject, @@ -1181,6 +1190,281 @@ async def test_insert_sso_user_sets_user_alias_from_display_name(): assert new_user_request.user_alias == "Doe, Jane" +RACE_SUBJECT: Final = "sso-race-subject" +RACE_VALUES: Final[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: Final = "P2002" + + +def _caused_by(error: ProxyException, cause: Exception) -> ProxyException: + try: + raise error from cause + except ProxyException as wrapped: + return wrapped + + +def _writer_with_row(row: Mapping[str, JsonValue] | None) -> MagicMock: + prisma: Final = MagicMock() + prisma.writer_db.litellm_usertable.find_unique = AsyncMock(return_value=row) + return prisma + + +def _failing_insert(error: Exception) -> Callable[[], Awaitable[NewUserResponse]]: + async def insert() -> NewUserResponse: + raise error + + return insert + + +DUPLICATE_ID: Final = 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, + HTTPException(status_code=409, detail="User with id already exists"), + _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", "http-conflict", "unique-violation"), +) +async def test_concurrent_first_sso_login_adopts_the_row_the_winning_login_inserted(conflict: Exception) -> None: + prisma: Final = _writer_with_row( + {"user_id": RACE_SUBJECT, "sso_user_id": RACE_SUBJECT, "user_role": "internal_user"} + ) + + user: Final = 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: Mapping[str, str | None] | None, +) -> None: + conflict: Final = ProxyException( + message="User with email already exists", type="bad_request", param="user_email", code=409 + ) + prisma: Final = _writer_with_row(row) + + refused: Final = await _insert_sso_user_or_adopt_concurrent_insert( + user_defined_values=RACE_VALUES, prisma_client=prisma, insert=_failing_insert(conflict) + ) + assert isinstance(refused, SSOInsertFailure) and refused.error is conflict + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "requested_role,persisted_role", (("internal_user", "proxy_admin"), ("proxy_admin", "internal_user")) +) +async def test_concurrent_sso_creation_refuses_a_row_with_a_different_effective_role( + requested_role: str, persisted_role: str +) -> None: + prisma: Final = _writer_with_row( + {"user_id": RACE_SUBJECT, "sso_user_id": RACE_SUBJECT, "user_role": persisted_role} + ) + values: Final[SSOUserDefinedValues] = {**RACE_VALUES, "user_role": requested_role} + refused: Final = await _insert_sso_user_or_adopt_concurrent_insert( + user_defined_values=values, prisma_client=prisma, insert=_failing_insert(DUPLICATE_ID) + ) + assert isinstance(refused, SSOInsertFailure) and refused.error is DUPLICATE_ID + + +@pytest.mark.asyncio +@pytest.mark.parametrize("memberships", ((), ("first-org",), ("first-org", "second-org"))) +async def test_concurrent_sso_creation_requires_every_configured_organization_membership( + memberships: tuple[str, ...], monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "default_internal_user_params", {"organizations": ["first-org", "second-org"]}) + prisma: Final = _writer_with_row( + { + "user_id": RACE_SUBJECT, + "sso_user_id": RACE_SUBJECT, + "user_role": "internal_user", + "organization_memberships": [ + { + "user_id": RACE_SUBJECT, + "organization_id": organization_id, + "created_at": "2000-01-01T00:00:00Z", + "updated_at": "2000-01-01T00:00:00Z", + } + for organization_id in memberships + ], + } + ) + outcome: Final = await _insert_sso_user_or_adopt_concurrent_insert( + user_defined_values=RACE_VALUES, prisma_client=prisma, insert=_failing_insert(DUPLICATE_ID) + ) + if len(memberships) < 2: + assert isinstance(outcome, SSOInsertFailure) and outcome.error is DUPLICATE_ID + return + assert isinstance(outcome, LiteLLM_UserTable) + assert {member.organization_id for member in outcome.organization_memberships or ()} == set(memberships) + prisma.writer_db.litellm_usertable.find_unique.assert_awaited_once_with( + where={"user_id": RACE_SUBJECT}, include={"organization_memberships": True} + ) + + +@pytest.mark.asyncio +async def test_concurrent_sso_creation_refuses_partial_membership_across_multiple_required_teams( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "default_internal_user_params", {"teams": ["joined", "missing"]}) + prisma: Final = _writer_with_row( + {"user_id": RACE_SUBJECT, "sso_user_id": RACE_SUBJECT, "user_role": "internal_user", "teams": ["joined"]} + ) + refused: Final = await _insert_sso_user_or_adopt_concurrent_insert( + user_defined_values=RACE_VALUES, prisma_client=prisma, insert=_failing_insert(DUPLICATE_ID) + ) + assert isinstance(refused, SSOInsertFailure) and refused.error is DUPLICATE_ID + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "defaults,user_role", ((None, "internal_user"), ({"organizations": ["not-for-admins"]}, "proxy_admin")) +) +async def test_public_sso_upsert_recovers_an_actual_duplicate_user_refusal( + defaults: Mapping[str, JsonValue] | None, user_role: str, monkeypatch: pytest.MonkeyPatch +) -> None: + from litellm.proxy import proxy_server + + row: Final = {"user_id": RACE_SUBJECT, "sso_user_id": RACE_SUBJECT, "user_role": user_role} + prisma: Final = _writer_with_row(row) + prisma.db.litellm_usertable.find_first = AsyncMock(return_value=LiteLLM_UserTable.model_validate(row)) + monkeypatch.setattr(proxy_server, "prisma_client", prisma) + monkeypatch.setattr(litellm, "default_internal_user_params", defaults) + values: Final[SSOUserDefinedValues] = {**RACE_VALUES, "user_role": user_role} + recovered: Final = await SSOAuthenticationHandler.upsert_sso_user( + result=CustomOpenID(id=RACE_SUBJECT, email=RACE_VALUES["user_email"], user_role=user_role, team_ids=[]), + user_info=None, + user_email=RACE_VALUES["user_email"], + user_defined_values=values, + prisma_client=prisma, + ) + assert isinstance(recovered, LiteLLM_UserTable) + assert (recovered.user_id, recovered.sso_user_id, recovered.user_role) == ( + RACE_SUBJECT, + RACE_SUBJECT, + user_role, + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "defaults,user_role", + ( + ({"teams": ["default-team"]}, "internal_user"), + ({"teams": [{"team_id": "default-team", "max_budget_in_team": 10}]}, "internal_user"), + ({"team_id": "default-team", "teams": ["other-team"]}, "internal_user"), + ({"team_id": "other-team", "teams": ["default-team"]}, "proxy_admin"), + ), +) +async def test_concurrent_sso_creation_returns_only_a_row_with_the_configured_memberships( + defaults: Mapping[str, JsonValue], user_role: str, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "default_internal_user_params", defaults) + values: Final[SSOUserDefinedValues] = {**RACE_VALUES, "user_role": user_role} + pending: Final = _writer_with_row( + {"user_id": RACE_SUBJECT, "sso_user_id": RACE_SUBJECT, "user_role": user_role, "teams": []} + ) + refused: Final = await _insert_sso_user_or_adopt_concurrent_insert( + user_defined_values=values, prisma_client=pending, insert=_failing_insert(DUPLICATE_ID) + ) + assert isinstance(refused, SSOInsertFailure) and refused.error is DUPLICATE_ID + completed: Final = _writer_with_row( + { + "user_id": RACE_SUBJECT, + "sso_user_id": RACE_SUBJECT, + "user_role": user_role, + "teams": ["default-team"], + } + ) + user: Final = await _insert_sso_user_or_adopt_concurrent_insert( + user_defined_values=values, prisma_client=completed, insert=_failing_insert(DUPLICATE_ID) + ) + assert isinstance(user, LiteLLM_UserTable) and user.teams == ["default-team"] + + +@pytest.mark.asyncio +async def test_recovered_first_login_can_be_authorized_after_an_earlier_missing_user_lookup( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(auth_checks, "last_db_access_time", LimitedSizeOrderedDict(max_size=10)) + prisma: Final = _writer_with_row( + {"user_id": RACE_SUBJECT, "sso_user_id": RACE_SUBJECT, "user_role": "internal_user"} + ) + prisma.db.litellm_usertable.find_unique = AsyncMock(return_value=None) + cache: Final = UserApiKeyCache() + with pytest.raises(UserNotFoundError): + await auth_checks.get_user_object(RACE_SUBJECT, prisma, cache, user_id_upsert=False) + recovered: Final = await _insert_sso_user_or_adopt_concurrent_insert( + user_defined_values=RACE_VALUES, prisma_client=prisma, insert=_failing_insert(DUPLICATE_ID) + ) + assert isinstance(recovered, LiteLLM_UserTable) + prisma.db.litellm_usertable.find_unique.return_value = recovered + authorized: Final = await auth_checks.get_user_object(RACE_SUBJECT, prisma, cache, user_id_upsert=False) + assert authorized is not None and authorized.user_id == recovered.user_id + + +@pytest.mark.asyncio +async def test_first_sso_login_preserves_an_unrelated_database_failure_without_readback() -> None: + failure: Final = _caused_by( + ProxyException(message="insert failed", type="internal_server_error", param=None, code=500), + ConnectionError("server closed the connection unexpectedly"), + ) + prisma: Final = _writer_with_row({"user_id": RACE_SUBJECT, "sso_user_id": RACE_SUBJECT}) + + refused: Final = await _insert_sso_user_or_adopt_concurrent_insert( + user_defined_values=RACE_VALUES, prisma_client=prisma, insert=_failing_insert(failure) + ) + assert isinstance(refused, SSOInsertFailure) and refused.error 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() -> None: + created: Final = NewUserResponse(user_id=RACE_SUBJECT, key="sk-created", teams=None) + prisma: Final = _writer_with_row(None) + + async def insert() -> NewUserResponse: + return created + + user: Final = 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(): """