diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index ed5843f4bce..7cf8f82fb1c 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/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: 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..6aba820ef96 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,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: diff --git a/tests/unit/proxy/management_endpoints/test_ui_sso.py b/tests/unit/proxy/management_endpoints/test_ui_sso.py index 4e128440703..7be77fdaa63 100644 --- a/tests/unit/proxy/management_endpoints/test_ui_sso.py +++ b/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(): """