From fd64fd1294a049ab8cf64df87354bc1cf77a450f Mon Sep 17 00:00:00 2001 From: yuneng Date: Sat, 10 Oct 2026 21:12:49 +0000 Subject: [PATCH 1/5] fix(sso): adopt the concurrently inserted user on a first-login duplicate race --- litellm/proxy/management_endpoints/ui_sso.py | 51 +++++++- .../test_cli_sso_login_ui_disabled.py | 56 ++++++-- .../proxy/management_endpoints/test_ui_sso.py | 122 +++++++++++++++++- 3 files changed, 216 insertions(+), 13 deletions(-) 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(): """ From b966af9e44d3237edac96331d79d193d789240b7 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang <yuneng@berri.ai> Date: Sat, 10 Oct 2026 16:25:11 -0700 Subject: [PATCH 2/5] fix: recover only fully provisioned matching SSO accounts --- litellm/proxy/management_endpoints/ui_sso.py | 66 +- litellm/repositories/user_repository.py | 11 +- .../proxy/management_endpoints/test_ui_sso.py | 1059 +++++++---------- 3 files changed, 498 insertions(+), 638 deletions(-) diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index 7cf8f82fb1c..f8e0d3ab98e 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 ( @@ -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, @@ -2679,15 +2683,12 @@ async def insert_sso_user( 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__) + 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( @@ -2698,24 +2699,51 @@ async def _concurrently_inserted_sso_user( 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: + 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( + 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: +) -> NewUserResponse | LiteLLM_UserTable | SSOInsertFailure: try: return await insert() - except Exception as error: + 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: - raise + return SSOInsertFailure(error) verbose_proxy_logger.info("SSO user %s was inserted by a concurrent login, using that row", concurrent.user_id) return concurrent @@ -3350,11 +3378,17 @@ class SSOAuthenticationHandler: ) else: verbose_proxy_logger.info("user not in DB, inserting user into LiteLLM DB") - user_info = await _insert_sso_user_or_adopt_concurrent_insert( + 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/unit/proxy/management_endpoints/test_ui_sso.py b/tests/unit/proxy/management_endpoints/test_ui_sso.py index 7be77fdaa63..26367fa8733 100644 --- a/tests/unit/proxy/management_endpoints/test_ui_sso.py +++ b/tests/unit/proxy/management_endpoints/test_ui_sso.py @@ -2,31 +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.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, @@ -127,9 +135,7 @@ def test_microsoft_sso_handler_with_empty_response(): # Test with None response # Act - result = MicrosoftSSOHandler.openid_from_response( - response=None, team_ids=[], user_role=None - ) + result = MicrosoftSSOHandler.openid_from_response(response=None, team_ids=[], user_role=None) # Assert assert isinstance(result, CustomOpenID) @@ -165,12 +171,8 @@ def test_microsoft_sso_handler_openid_from_response_with_custom_attributes(): "custom_display_name", ), patch("litellm.constants.MICROSOFT_USER_ID_ATTRIBUTE", "custom_id_field"), - patch( - "litellm.constants.MICROSOFT_USER_FIRST_NAME_ATTRIBUTE", "custom_first_name" - ), - patch( - "litellm.constants.MICROSOFT_USER_LAST_NAME_ATTRIBUTE", "custom_last_name" - ), + patch("litellm.constants.MICROSOFT_USER_FIRST_NAME_ATTRIBUTE", "custom_first_name"), + patch("litellm.constants.MICROSOFT_USER_LAST_NAME_ATTRIBUTE", "custom_last_name"), patch( "litellm.proxy.management_endpoints.ui_sso.MICROSOFT_USER_EMAIL_ATTRIBUTE", "custom_email_field", @@ -210,9 +212,7 @@ def test_microsoft_sso_handler_openid_from_response_with_custom_attributes(): @pytest.fixture def stubbed_graph_api(httpx_transport): with respx.mock: - respx.get(url__regex=r".*graph\.microsoft\.com.*").mock( - return_value=httpx.Response(200, json={"value": []}) - ) + respx.get(url__regex=r".*graph\.microsoft\.com.*").mock(return_value=httpx.Response(200, json={"value": []})) yield @@ -308,9 +308,7 @@ def test_get_google_callback_response(): with patch.dict(os.environ, {"GOOGLE_CLIENT_SECRET": "mock_secret"}): mock_verify = AsyncMock(return_value=mock_response) - with patch( - "fastapi_sso.sso.google.GoogleSSO.verify_and_process", new=mock_verify - ): + with patch("fastapi_sso.sso.google.GoogleSSO.verify_and_process", new=mock_verify): # Act result = asyncio.run( GoogleSSOHandler.get_google_callback_response( @@ -353,16 +351,12 @@ async def test_get_user_groups_from_graph_api(): mock.json.return_value = mock_response return mock - with patch( - "litellm.proxy.management_endpoints.ui_sso.get_async_httpx_client" - ) as mock_client: + with patch("litellm.proxy.management_endpoints.ui_sso.get_async_httpx_client") as mock_client: mock_client.return_value = MagicMock() mock_client.return_value.get = mock_get # Act - result = await MicrosoftSSOHandler.get_user_groups_from_graph_api( - access_token="mock_token" - ) + result = await MicrosoftSSOHandler.get_user_groups_from_graph_api(access_token="mock_token") # Assert assert isinstance(result, list) @@ -384,16 +378,12 @@ async def test_get_user_groups_empty_response(): mock.json.return_value = mock_response return mock - with patch( - "litellm.proxy.management_endpoints.ui_sso.get_async_httpx_client" - ) as mock_client: + with patch("litellm.proxy.management_endpoints.ui_sso.get_async_httpx_client") as mock_client: mock_client.return_value = MagicMock() mock_client.return_value.get = mock_get # Act - result = await MicrosoftSSOHandler.get_user_groups_from_graph_api( - access_token="mock_token" - ) + result = await MicrosoftSSOHandler.get_user_groups_from_graph_api(access_token="mock_token") # Assert assert isinstance(result, list) @@ -406,16 +396,12 @@ async def test_get_user_groups_error_handling(): async def mock_get(*args, **kwargs): raise Exception("API Error") - with patch( - "litellm.proxy.management_endpoints.ui_sso.get_async_httpx_client" - ) as mock_client: + with patch("litellm.proxy.management_endpoints.ui_sso.get_async_httpx_client") as mock_client: mock_client.return_value = MagicMock() mock_client.return_value.get = mock_get # Act - result = await MicrosoftSSOHandler.get_user_groups_from_graph_api( - access_token="mock_token" - ) + result = await MicrosoftSSOHandler.get_user_groups_from_graph_api(access_token="mock_token") # Assert assert isinstance(result, list) @@ -434,9 +420,7 @@ async def test_get_user_groups_uses_default_graph_endpoint(monkeypatch): mock.json.return_value = {"value": []} return mock - with patch( - "litellm.proxy.management_endpoints.ui_sso.get_async_httpx_client" - ) as mock_client: + with patch("litellm.proxy.management_endpoints.ui_sso.get_async_httpx_client") as mock_client: mock_client.return_value = MagicMock() mock_client.return_value.get = mock_get @@ -457,9 +441,7 @@ async def test_get_user_groups_uses_configured_graph_endpoint(monkeypatch): mock.json.return_value = {"value": []} return mock - with patch( - "litellm.proxy.management_endpoints.ui_sso.get_async_httpx_client" - ) as mock_client: + with patch("litellm.proxy.management_endpoints.ui_sso.get_async_httpx_client") as mock_client: mock_client.return_value = MagicMock() mock_client.return_value.get = mock_get @@ -489,9 +471,7 @@ async def test_get_group_ids_from_service_principal_uses_configured_graph_endpoi access_token="mock_token", ) - assert requested_urls == [ - "https://graph.microsoft.us/v1.0/servicePrincipals/sp-123/appRoleAssignedTo" - ] + assert requested_urls == ["https://graph.microsoft.us/v1.0/servicePrincipals/sp-123/appRoleAssignedTo"] @pytest.mark.asyncio @@ -589,9 +569,7 @@ def test_get_group_ids_from_graph_api_response(): "team_params", [ # Test case 1: Using DefaultTeamSSOParams - DefaultTeamSSOParams( - max_budget=10, budget_duration="1d", models=["special-gpt-5"] - ), + DefaultTeamSSOParams(max_budget=10, budget_duration="1d", models=["special-gpt-5"]), # Test case 2: Using Dict {"max_budget": 10, "budget_duration": "1d", "models": ["special-gpt-5"]}, ], @@ -630,9 +608,7 @@ async def test_default_team_params(team_params): # Assert # Verify team was created with correct parameters mock_prisma.db.litellm_teamtable.create.assert_called_once() - create_call_args = mock_prisma.db.litellm_teamtable.create.call_args.kwargs[ - "data" - ] + create_call_args = mock_prisma.db.litellm_teamtable.create.call_args.kwargs["data"] assert create_call_args["team_id"] == team_id assert create_call_args["team_alias"] == "Test Team" assert create_call_args["max_budget"] == 10 @@ -670,10 +646,13 @@ async def test_default_team_params_organization_id_reaches_sso_created_team(team updated_by="admin", ) - with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), patch( - "litellm.proxy.management_endpoints.team_endpoints.get_org_object", - AsyncMock(return_value=mock_org), - ) as mock_get_org: + with ( + patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), + patch( + "litellm.proxy.management_endpoints.team_endpoints.get_org_object", + AsyncMock(return_value=mock_org), + ) as mock_get_org, + ): team_id = str(uuid.uuid4()) await MicrosoftSSOHandler.create_litellm_teams_from_service_principal_team_ids( service_principal_teams=[ @@ -725,9 +704,7 @@ async def test_create_team_without_default_params(): # Assert mock_prisma.db.litellm_teamtable.create.assert_called_once() - create_call_args = mock_prisma.db.litellm_teamtable.create.call_args.kwargs[ - "data" - ] + create_call_args = mock_prisma.db.litellm_teamtable.create.call_args.kwargs["data"] assert create_call_args["team_id"] == team_id assert create_call_args["team_alias"] == "Test Team" # Should not have any of the optional fields @@ -1182,8 +1159,8 @@ 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 = { +RACE_SUBJECT: Final = "sso-race-subject" +RACE_VALUES: Final[SSOUserDefinedValues] = { "models": [], "user_id": RACE_SUBJECT, "user_email": "sso-race-subject@example.com", @@ -1194,28 +1171,30 @@ RACE_VALUES: SSOUserDefinedValues = { class _UniqueViolation(Exception): - code = "P2002" + code: Final = "P2002" def _caused_by(error: ProxyException, cause: Exception) -> ProxyException: - error.__cause__ = cause - return error + try: + raise error from cause + except ProxyException as wrapped: + return wrapped -def _writer_with_row(row: dict | None) -> MagicMock: - prisma = MagicMock() +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): +def _failing_insert(error: Exception) -> Callable[[], Awaitable[NewUserResponse]]: async def insert() -> NewUserResponse: raise error return insert -DUPLICATE_ID = ProxyException( +DUPLICATE_ID: Final = ProxyException( message="User with id sso-race-subject already exists", type="bad_request", param="user_id", code=409 ) @@ -1225,17 +1204,20 @@ DUPLICATE_ID = ProxyException( "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", "unique-violation"), + ids=("duplicate-id-409", "http-conflict", "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"}) +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 = await _insert_sso_user_or_adopt_concurrent_insert( + user: Final = await _insert_sso_user_or_adopt_concurrent_insert( user_defined_values=RACE_VALUES, prisma_client=prisma, insert=_failing_insert(conflict) ) @@ -1254,46 +1236,192 @@ async def test_concurrent_first_sso_login_adopts_the_row_the_winning_login_inser ), 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( +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 = _writer_with_row(row) + prisma: Final = _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 + 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 -async def test_first_sso_login_reraises_an_unrelated_database_failure_without_readback(): - failure = _caused_by( +@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 +async def test_public_sso_upsert_recovers_an_actual_duplicate_user_refusal(monkeypatch: pytest.MonkeyPatch) -> None: + from litellm.proxy import proxy_server + + row: Final = {"user_id": RACE_SUBJECT, "sso_user_id": RACE_SUBJECT, "user_role": "internal_user"} + 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", None) + values: Final[SSOUserDefinedValues] = {**RACE_VALUES} + recovered: Final = await SSOAuthenticationHandler.upsert_sso_user( + result=CustomOpenID(id=RACE_SUBJECT, email=RACE_VALUES["user_email"], user_role="internal_user", 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, + "internal_user", + ) + + +@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 = _writer_with_row({"user_id": RACE_SUBJECT, "sso_user_id": RACE_SUBJECT}) + prisma: Final = _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 + 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(): - created = NewUserResponse(user_id=RACE_SUBJECT, key="sk-created", teams=None) - prisma = _writer_with_row(None) +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 = await _insert_sso_user_or_adopt_concurrent_insert( + user: Final = await _insert_sso_user_or_adopt_concurrent_insert( user_defined_values=RACE_VALUES, prisma_client=prisma, insert=insert ) @@ -1462,9 +1590,7 @@ async def test_get_user_info_from_db_user_exists(): "user_email": user_email, "user_defined_values": user_defined_values, } - with patch( - "litellm.proxy.management_endpoints.ui_sso.get_user_object" - ) as mock_get_user_object: + with patch("litellm.proxy.management_endpoints.ui_sso.get_user_object") as mock_get_user_object: await get_user_info_from_db(**args) mock_get_user_object.assert_called_once() assert mock_get_user_object.call_args.kwargs["user_id"] == "krrishd" @@ -1504,9 +1630,7 @@ async def test_get_user_info_from_db_user_exists_alternate_user_id(): "user_defined_values": user_defined_values, "alternate_user_id": "krrishd-email1234", } - with patch( - "litellm.proxy.management_endpoints.ui_sso.get_user_object" - ) as mock_get_user_object: + with patch("litellm.proxy.management_endpoints.ui_sso.get_user_object") as mock_get_user_object: await get_user_info_from_db(**args) mock_get_user_object.assert_called_once() assert mock_get_user_object.call_args.kwargs["user_id"] == "krrishd-email1234" @@ -1808,17 +1932,13 @@ async def test_get_generic_sso_response_with_additional_headers(): # Mock the SSO provider and its methods mock_sso_instance = MagicMock() mock_sso_instance.verify_and_process = AsyncMock(return_value=mock_sso_response) - mock_sso_instance.access_token = ( - None # Avoid triggering JWT decode in process_sso_jwt_access_token - ) + mock_sso_instance.access_token = None # Avoid triggering JWT decode in process_sso_jwt_access_token mock_sso_class = MagicMock(return_value=mock_sso_instance) with patch.dict(os.environ, test_env_vars): with patch("fastapi_sso.sso.base.DiscoveryDocument"): - with patch( - "fastapi_sso.sso.generic.create_provider", return_value=mock_sso_class - ): + with patch("fastapi_sso.sso.generic.create_provider", return_value=mock_sso_class): # Act result, received_response, _, _ = await get_generic_sso_response( request=mock_request, @@ -1872,17 +1992,13 @@ async def test_get_generic_sso_response_with_empty_headers(): # Mock the SSO provider and its methods mock_sso_instance = MagicMock() mock_sso_instance.verify_and_process = AsyncMock(return_value=mock_sso_response) - mock_sso_instance.access_token = ( - None # Avoid triggering JWT decode in process_sso_jwt_access_token - ) + mock_sso_instance.access_token = None # Avoid triggering JWT decode in process_sso_jwt_access_token mock_sso_class = MagicMock(return_value=mock_sso_instance) with patch.dict(os.environ, test_env_vars): with patch("fastapi_sso.sso.base.DiscoveryDocument"): - with patch( - "fastapi_sso.sso.generic.create_provider", return_value=mock_sso_class - ): + with patch("fastapi_sso.sso.generic.create_provider", return_value=mock_sso_class): # Act result, received_response, _, _ = await get_generic_sso_response( request=mock_request, @@ -2263,9 +2379,7 @@ class TestAuthCallbackRouting: for state in non_cli_states: # This mimics the routing logic in auth_callback - should_route_to_cli = state and state.startswith( - f"{LITELLM_CLI_SESSION_TOKEN_PREFIX}:" - ) + should_route_to_cli = state and state.startswith(f"{LITELLM_CLI_SESSION_TOKEN_PREFIX}:") assert not should_route_to_cli, f"State '{state}' should not route to CLI" @@ -2300,9 +2414,7 @@ class TestGoogleLoginCLIIntegration: for source, key in test_cases: cli_state = SSOAuthenticationHandler._get_cli_state(source=source, key=key) - assert ( - cli_state is None - ), f"CLI state should not be generated for source='{source}', key='{key}'" + assert cli_state is None, f"CLI state should not be generated for source='{source}', key='{key}'" class TestSSOHandlerIntegration: @@ -2313,24 +2425,13 @@ class TestSSOHandlerIntegration: from litellm.proxy.management_endpoints.ui_sso import SSOAuthenticationHandler # Test that SSO handler is used when client IDs are provided - assert ( - SSOAuthenticationHandler.should_use_sso_handler(google_client_id="test") - is True - ) - assert ( - SSOAuthenticationHandler.should_use_sso_handler(microsoft_client_id="test") - is True - ) - assert ( - SSOAuthenticationHandler.should_use_sso_handler(generic_client_id="test") - is True - ) + assert SSOAuthenticationHandler.should_use_sso_handler(google_client_id="test") is True + assert SSOAuthenticationHandler.should_use_sso_handler(microsoft_client_id="test") is True + assert SSOAuthenticationHandler.should_use_sso_handler(generic_client_id="test") is True # Test that SSO handler is not used when no client IDs are provided assert SSOAuthenticationHandler.should_use_sso_handler() is False - assert ( - SSOAuthenticationHandler.should_use_sso_handler(None, None, None) is False - ) + assert SSOAuthenticationHandler.should_use_sso_handler(None, None, None) is False @patch.dict(os.environ, {}, clear=False) def test_get_redirect_url_for_sso(self): @@ -2401,9 +2502,7 @@ class TestSSOStateHandling: """Test generating CLI state with valid parameters""" from litellm.proxy.management_endpoints.ui_sso import SSOAuthenticationHandler - state = SSOAuthenticationHandler._get_cli_state( - source="litellm-cli", key="cli-test1234567890" - ) + state = SSOAuthenticationHandler._get_cli_state(source="litellm-cli", key="cli-test1234567890") assert state is not None assert state.startswith("litellm-session-token:") @@ -2413,9 +2512,7 @@ class TestSSOStateHandling: """Test generating CLI state with invalid source""" from litellm.proxy.management_endpoints.ui_sso import SSOAuthenticationHandler - state = SSOAuthenticationHandler._get_cli_state( - source="invalid_source", key="cli-test1234567890" - ) + state = SSOAuthenticationHandler._get_cli_state(source="invalid_source", key="cli-test1234567890") assert state is None @@ -2431,9 +2528,7 @@ class TestSSOStateHandling: """Test generating CLI state without source""" from litellm.proxy.management_endpoints.ui_sso import SSOAuthenticationHandler - state = SSOAuthenticationHandler._get_cli_state( - source=None, key="cli-test1234567890" - ) + state = SSOAuthenticationHandler._get_cli_state(source=None, key="cli-test1234567890") assert state is None @@ -2602,9 +2697,7 @@ class TestCustomUISSO: # Test that the ValueError is raised with the correct message import pytest - with pytest.raises( - ValueError, match="Enterprise features are not available" - ): + with pytest.raises(ValueError, match="Enterprise features are not available"): asyncio.run(mock_google_login()) @pytest.mark.asyncio @@ -2638,9 +2731,7 @@ class TestCustomUISSO: picture=None, provider="custom", ) - mock_custom_handler.handle_custom_ui_sso_sign_in = AsyncMock( - return_value=expected_openid - ) + mock_custom_handler.handle_custom_ui_sso_sign_in = AsyncMock(return_value=expected_openid) # Mock the redirect response method mock_redirect_response = MagicMock() @@ -2661,15 +2752,11 @@ class TestCustomUISSO: return_value=mock_redirect_response, ) as mock_get_redirect: # Act - result = await EnterpriseCustomSSOHandler.handle_custom_ui_sso_sign_in( - request=mock_request - ) + result = await EnterpriseCustomSSOHandler.handle_custom_ui_sso_sign_in(request=mock_request) # Assert # Verify the custom handler was called with the request - mock_custom_handler.handle_custom_ui_sso_sign_in.assert_called_once_with( - request=mock_request - ) + mock_custom_handler.handle_custom_ui_sso_sign_in.assert_called_once_with(request=mock_request) # Verify the redirect response was generated with correct OpenID mock_get_redirect.assert_called_once_with( @@ -2714,9 +2801,7 @@ class TestCustomUISSO: {"trusted_proxy_ranges": ["10.0.0.0/24"]}, ): with pytest.raises(ValueError, match="not trusted"): - await EnterpriseCustomSSOHandler.handle_custom_ui_sso_sign_in( - request=mock_request - ) + await EnterpriseCustomSSOHandler.handle_custom_ui_sso_sign_in(request=mock_request) mock_custom_handler.handle_custom_ui_sso_sign_in.assert_not_called() @@ -2748,9 +2833,7 @@ class TestCustomUISSO: request_headers_dict = dict(request.headers) return OpenID( id=request_headers_dict.get("x-litellm-user-id", "default_user"), - email=request_headers_dict.get( - "x-litellm-user-email", "default@test.com" - ), + email=request_headers_dict.get("x-litellm-user-email", "default@test.com"), first_name="Custom", last_name="Handler", display_name="Custom Handler Test", @@ -2790,9 +2873,7 @@ class TestCustomUISSO: return_value=mock_redirect_response, ) as mock_get_redirect: # Act - result = await EnterpriseCustomSSOHandler.handle_custom_ui_sso_sign_in( - request=mock_request - ) + result = await EnterpriseCustomSSOHandler.handle_custom_ui_sso_sign_in(request=mock_request) # Assert that our custom handler was executed assert test_handler_instance.method_called is True @@ -2964,9 +3045,7 @@ class TestCLIKeyRegenerationFlow: redis_cache.set_cache.side_effect = lambda key, value, ttl: redis_store.__setitem__( key, str(value).encode("utf-8") ) - redis_cache.get_cache.side_effect = lambda key: RedisCache._get_cache_logic( - MagicMock(), redis_store.get(key) - ) + redis_cache.get_cache.side_effect = lambda key: RedisCache._get_cache_logic(MagicMock(), redis_store.get(key)) cache = MagicMock() cache.redis_cache = redis_cache @@ -3006,12 +3085,8 @@ class TestCLIKeyRegenerationFlow: assert mock_cache.increment_cache.call_args.kwargs["ttl"] == 60 mock_cache.set_cache.assert_called_once() flow_data = mock_cache.set_cache.call_args.kwargs["value"] - assert flow_data["poll_secret_hash"] == _hash_cli_sso_secret( - result["poll_secret"] - ) - assert flow_data["user_code_hash"] == _hash_cli_sso_secret( - _normalize_cli_sso_user_code(result["user_code"]) - ) + assert flow_data["poll_secret_hash"] == _hash_cli_sso_secret(result["poll_secret"]) + assert flow_data["user_code_hash"] == _hash_cli_sso_secret(_normalize_cli_sso_user_code(result["user_code"])) assert flow_data["poll_secret_hash"] != result["poll_secret"] assert flow_data["user_code_hash"] != result["user_code"] @@ -3119,9 +3194,7 @@ class TestCLIKeyRegenerationFlow: mock_cache = MagicMock(redis_cache=None) mock_cache.get_cache.return_value = {"poll_secret_hash": "h"} env_without_sso_providers = { - name: value - for name, value in os.environ.items() - if name not in _SSO_PROVIDER_ENV_VARS + name: value for name, value in os.environ.items() if name not in _SSO_PROVIDER_ENV_VARS } async def drive(enabled: bool): @@ -3172,9 +3245,7 @@ class TestCLIKeyRegenerationFlow: ) from litellm.proxy.management_endpoints.ui_sso import SSOAuthenticationHandler - manual_state = SSOAuthenticationHandler._get_cli_state( - source=LITELLM_CLI_SOURCE_IDENTIFIER, key="cli-abc123" - ) + manual_state = SSOAuthenticationHandler._get_cli_state(source=LITELLM_CLI_SOURCE_IDENTIFIER, key="cli-abc123") prefill_state = SSOAuthenticationHandler._get_cli_state( source=LITELLM_CLI_SOURCE_IDENTIFIER, key="cli-abc123", @@ -3182,14 +3253,9 @@ class TestCLIKeyRegenerationFlow: ) assert manual_state == f"{LITELLM_CLI_SESSION_TOKEN_PREFIX}:cli-abc123" + assert prefill_state == f"{LITELLM_CLI_SESSION_TOKEN_PREFIX}:cli-abc123:WXYZ-2345" assert ( - prefill_state == f"{LITELLM_CLI_SESSION_TOKEN_PREFIX}:cli-abc123:WXYZ-2345" - ) - assert ( - SSOAuthenticationHandler._get_cli_state( - source="not-cli", key="cli-abc123", user_code="WXYZ-2345" - ) - is None + SSOAuthenticationHandler._get_cli_state(source="not-cli", key="cli-abc123", user_code="WXYZ-2345") is None ) def test_get_cli_state_drops_malformed_user_code(self): @@ -3273,9 +3339,7 @@ class TestCLIKeyRegenerationFlow: browser_complete_token="browser-token", ) - input_line = next( - line for line in html.splitlines() if 'name="user_code"' in line - ) + input_line = next(line for line in html.splitlines() if 'name="user_code"' in line) assert "value=" not in input_line assert "shown in your terminal" in html @@ -3338,15 +3402,11 @@ class TestCLIKeyRegenerationFlow: ) mock_request = MagicMock(spec=Request) - mock_request.body = AsyncMock( - return_value=b"user_code=ABCD-EFGH&browser_complete_token=browser-token" - ) + mock_request.body = AsyncMock(return_value=b"user_code=ABCD-EFGH&browser_complete_token=browser-token") mock_cache = MagicMock(redis_cache=None) mock_cache.get_cache.return_value = { "poll_secret_hash": _hash_cli_sso_secret("poll-secret"), - "user_code_hash": _hash_cli_sso_secret( - _normalize_cli_sso_user_code("ABCD-EFGH") - ), + "user_code_hash": _hash_cli_sso_secret(_normalize_cli_sso_user_code("ABCD-EFGH")), "browser_complete_token_hash": _hash_cli_sso_secret("browser-token"), "sso_complete": True, "user_code_verified": False, @@ -3361,9 +3421,7 @@ class TestCLIKeyRegenerationFlow: return_value="<html>Success</html>", ), ): - result = await cli_sso_complete( - request=mock_request, login_id="cli-session-4567890" - ) + result = await cli_sso_complete(request=mock_request, login_id="cli-session-4567890") assert result.status_code == 200 flow_data = mock_cache.set_cache.call_args.kwargs["value"] @@ -3383,9 +3441,7 @@ class TestCLIKeyRegenerationFlow: mock_cache = MagicMock(redis_cache=None) mock_cache.get_cache.return_value = { "poll_secret_hash": _hash_cli_sso_secret("poll-secret"), - "user_code_hash": _hash_cli_sso_secret( - _normalize_cli_sso_user_code("ABCD-EFGH") - ), + "user_code_hash": _hash_cli_sso_secret(_normalize_cli_sso_user_code("ABCD-EFGH")), "browser_complete_token_hash": _hash_cli_sso_secret("browser-token"), "sso_complete": True, "user_code_verified": False, @@ -3397,9 +3453,7 @@ class TestCLIKeyRegenerationFlow: patch("litellm.proxy.proxy_server.cli_sso_session_cache", mock_cache), ): with pytest.raises(HTTPException) as exc_info: - await cli_sso_complete( - request=mock_request, login_id="cli-session-4567890" - ) + await cli_sso_complete(request=mock_request, login_id="cli-session-4567890") assert exc_info.value.status_code == 400 mock_cache.set_cache.assert_not_called() @@ -3414,15 +3468,11 @@ class TestCLIKeyRegenerationFlow: ) mock_request = MagicMock(spec=Request) - mock_request.body = AsyncMock( - return_value=b"user_code=ABCD-EFGH&browser_complete_token=browser-token" - ) + mock_request.body = AsyncMock(return_value=b"user_code=ABCD-EFGH&browser_complete_token=browser-token") mock_cache = MagicMock(redis_cache=None) mock_cache.get_cache.return_value = { "poll_secret_hash": _hash_cli_sso_secret("poll-secret"), - "user_code_hash": _hash_cli_sso_secret( - _normalize_cli_sso_user_code("ABCD-EFGH") - ), + "user_code_hash": _hash_cli_sso_secret(_normalize_cli_sso_user_code("ABCD-EFGH")), "sso_complete": False, "user_code_verified": False, "session_data": None, @@ -3433,9 +3483,7 @@ class TestCLIKeyRegenerationFlow: patch("litellm.proxy.proxy_server.cli_sso_session_cache", mock_cache), ): with pytest.raises(HTTPException) as exc_info: - await cli_sso_complete( - request=mock_request, login_id="cli-session-4567890" - ) + await cli_sso_complete(request=mock_request, login_id="cli-session-4567890") assert exc_info.value.status_code == 400 assert exc_info.value.detail == "CLI login is not ready" @@ -3541,10 +3589,7 @@ class TestCLIKeyRegenerationFlow: assert result.status_code == 200 # Verify response contains success message (response is HTML) assert result.body is not None - assert ( - 'action="https://test.example.com/sso/cli/complete/cli-session-4567890"' - in result.body.decode() - ) + assert 'action="https://test.example.com/sso/cli/complete/cli-session-4567890"' in result.body.decode() @pytest.mark.asyncio async def test_cli_poll_key_returns_teams_for_selection(self): @@ -3672,9 +3717,7 @@ class TestCLIKeyRegenerationFlow: mock_result = {"user_id": "test-user", "email": "test@example.com"} with ( - patch( - "litellm.proxy.management_endpoints.ui_sso.cli_sso_callback" - ) as mock_cli_callback, + patch("litellm.proxy.management_endpoints.ui_sso.cli_sso_callback") as mock_cli_callback, patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), patch("litellm.proxy.proxy_server.master_key", "test-master-key"), patch("litellm.proxy.proxy_server.general_settings", {}), @@ -3708,15 +3751,11 @@ class TestCLIKeyRegenerationFlow: mock_request = MagicMock(spec=Request) mock_request.query_params = {"code": "some-auth-code"} - cli_state = ( - f"{LITELLM_CLI_SESSION_TOKEN_PREFIX}:cli-new-session-key-456:WXYZ-2345" - ) + cli_state = f"{LITELLM_CLI_SESSION_TOKEN_PREFIX}:cli-new-session-key-456:WXYZ-2345" mock_result = {"user_id": "test-user", "email": "test@example.com"} with ( - patch( - "litellm.proxy.management_endpoints.ui_sso.cli_sso_callback" - ) as mock_cli_callback, + patch("litellm.proxy.management_endpoints.ui_sso.cli_sso_callback") as mock_cli_callback, patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), patch("litellm.proxy.proxy_server.master_key", "test-master-key"), patch("litellm.proxy.proxy_server.general_settings", {}), @@ -3749,9 +3788,7 @@ class TestCLIKeyRegenerationFlow: mock_request = MagicMock() mock_request.base_url = "https://test.litellm.ai/" - with patch( - "litellm.proxy.utils.get_custom_url", return_value="https://test.litellm.ai" - ): + with patch("litellm.proxy.utils.get_custom_url", return_value="https://test.litellm.ai"): # Test with existing_key - should NOT be in URL redirect_url = SSOAuthenticationHandler.get_redirect_url_for_sso( request=mock_request, @@ -3771,9 +3808,7 @@ class TestCLIKeyRegenerationFlow: mock_request = MagicMock() mock_request.base_url = "https://test.litellm.ai/" - with patch( - "litellm.proxy.utils.get_custom_url", return_value="https://test.litellm.ai" - ): + with patch("litellm.proxy.utils.get_custom_url", return_value="https://test.litellm.ai"): # Test without existing_key redirect_url = SSOAuthenticationHandler.get_redirect_url_for_sso( request=mock_request, sso_callback_route="sso/callback" @@ -3896,9 +3931,7 @@ class TestCLIKeyRegenerationFlow: prisma_client = MagicMock() prisma_client.db.litellm_teamtable.find_many = find_many - details = await fetch_cli_sso_team_details( - prisma_client=prisma_client, teams=["team-a"] - ) + details = await fetch_cli_sso_team_details(prisma_client=prisma_client, teams=["team-a"]) assert find_many.await_args.kwargs["include"] == {"litellm_model_table": True} assert [detail.model_dump() for detail in details] == [ @@ -3923,24 +3956,12 @@ class TestCLIKeyRegenerationFlow: ) failing_client = MagicMock() - failing_client.db.litellm_teamtable.find_many = AsyncMock( - side_effect=Exception("connection reset") - ) - assert ( - await fetch_cli_sso_team_details( - prisma_client=failing_client, teams=["team-a"] - ) - is None - ) + failing_client.db.litellm_teamtable.find_many = AsyncMock(side_effect=Exception("connection reset")) + assert await fetch_cli_sso_team_details(prisma_client=failing_client, teams=["team-a"]) is None empty_client = MagicMock() empty_client.db.litellm_teamtable.find_many = AsyncMock(return_value=[]) - assert ( - await fetch_cli_sso_team_details( - prisma_client=empty_client, teams=["team-a"] - ) - == () - ) + assert await fetch_cli_sso_team_details(prisma_client=empty_client, teams=["team-a"]) == () @pytest.mark.asyncio async def test_cli_poll_key_mints_jwt_with_selected_team_grants(self): @@ -4004,18 +4025,14 @@ class TestCLIKeyRegenerationFlow: "team_details", [ pytest.param(None, id="detail_fetch_failed"), - pytest.param( - [{"team_id": "team-other", "team_models": []}], id="selected_team_absent" - ), + pytest.param([{"team_id": "team-other", "team_models": []}], id="selected_team_absent"), pytest.param( [{"team_id": "team-a", "team_alias": "Team A"}], id="legacy_detail_without_grants", ), ], ) - async def test_cli_poll_key_refuses_to_mint_when_team_grants_are_unknown( - self, team_details - ): + async def test_cli_poll_key_refuses_to_mint_when_team_grants_are_unknown(self, team_details): """An unknown team grant must never be minted as an empty one. get_complete_model_list falls through to the whole proxy model list when both @@ -4151,9 +4168,7 @@ class TestCLIKeyRegenerationFlow: ), patch( "litellm.proxy.auth.auth_checks.get_team_object", - new=AsyncMock( - side_effect=AssertionError("team lookup must be skipped") - ), + new=AsyncMock(side_effect=AssertionError("team lookup must be skipped")), ), ): result = await cli_poll_key( @@ -4235,9 +4250,7 @@ class TestGetAppRolesFromIdToken: # Assert assert result == ["Admin", "User", "Developer"] - mock_jwt_decode.assert_called_once_with( - mock_token, options={"verify_signature": False} - ) + mock_jwt_decode.assert_called_once_with(mock_token, options={"verify_signature": False}) def test_app_roles_picked_when_both_exist(self): """Test that 'app_roles' takes precedence when both 'app_roles' and 'roles' exist""" @@ -4359,9 +4372,7 @@ class TestProcessSSOJWTAccessToken: "groups": ["team1", "team2", "team3"], } - def test_process_sso_jwt_access_token_with_existing_team_ids( - self, mock_jwt_handler, sample_jwt_token - ): + def test_process_sso_jwt_access_token_with_existing_team_ids(self, mock_jwt_handler, sample_jwt_token): """Test that existing team IDs are not overwritten""" from litellm.proxy.management_endpoints.ui_sso import ( process_sso_jwt_access_token, @@ -4417,20 +4428,14 @@ class TestProcessSSOJWTAccessToken: ) # Assert - mock_jwt_decode.assert_called_once_with( - sample_jwt_token, options={"verify_signature": False} - ) - mock_jwt_handler.get_team_ids_from_jwt.assert_called_once_with( - sample_jwt_payload - ) + mock_jwt_decode.assert_called_once_with(sample_jwt_token, options={"verify_signature": False}) + mock_jwt_handler.get_team_ids_from_jwt.assert_called_once_with(sample_jwt_payload) # Verify team_ids was added to the dict as a key assert "team_ids" in result assert result["team_ids"] == ["team1", "team2", "team3"] - def test_process_sso_jwt_access_token_with_dict_existing_team_ids( - self, mock_jwt_handler, sample_jwt_token - ): + def test_process_sso_jwt_access_token_with_dict_existing_team_ids(self, mock_jwt_handler, sample_jwt_token): """Test that existing team IDs in dictionary are not overwritten""" from litellm.proxy.management_endpoints.ui_sso import ( process_sso_jwt_access_token, @@ -4473,9 +4478,7 @@ class TestProcessSSOJWTAccessToken: # Test with None access token with patch("jwt.decode") as mock_jwt_decode: - process_sso_jwt_access_token( - access_token_str=None, sso_jwt_handler=mock_jwt_handler, result=result - ) + process_sso_jwt_access_token(access_token_str=None, sso_jwt_handler=mock_jwt_handler, result=result) # Assert nothing was processed mock_jwt_decode.assert_not_called() @@ -4484,18 +4487,14 @@ class TestProcessSSOJWTAccessToken: # Test with empty string access token with patch("jwt.decode") as mock_jwt_decode: - process_sso_jwt_access_token( - access_token_str="", sso_jwt_handler=mock_jwt_handler, result=result - ) + process_sso_jwt_access_token(access_token_str="", sso_jwt_handler=mock_jwt_handler, result=result) # Assert nothing was processed mock_jwt_decode.assert_not_called() mock_jwt_handler.get_team_ids_from_jwt.assert_not_called() assert result.team_ids == [] - def test_process_sso_jwt_access_token_no_result( - self, mock_jwt_handler, sample_jwt_token - ): + def test_process_sso_jwt_access_token_no_result(self, mock_jwt_handler, sample_jwt_token): """Test that nothing happens when result is None""" from litellm.proxy.management_endpoints.ui_sso import ( process_sso_jwt_access_token, @@ -4513,9 +4512,7 @@ class TestProcessSSOJWTAccessToken: mock_jwt_decode.assert_not_called() mock_jwt_handler.get_team_ids_from_jwt.assert_not_called() - def test_process_sso_jwt_access_token_non_decode_exception_propagates( - self, mock_jwt_handler, sample_jwt_token - ): + def test_process_sso_jwt_access_token_non_decode_exception_propagates(self, mock_jwt_handler, sample_jwt_token): """Test that non-DecodeError JWT exceptions still propagate up.""" import jwt as pyjwt @@ -4525,9 +4522,7 @@ class TestProcessSSOJWTAccessToken: result = CustomOpenID(id="test_user", email="test@example.com", team_ids=[]) - with patch( - "jwt.decode", side_effect=pyjwt.exceptions.InvalidKeyError("Invalid key") - ) as mock_jwt_decode: + with patch("jwt.decode", side_effect=pyjwt.exceptions.InvalidKeyError("Invalid key")) as mock_jwt_decode: with pytest.raises(pyjwt.exceptions.InvalidKeyError, match="Invalid key"): process_sso_jwt_access_token( access_token_str=sample_jwt_token, @@ -4561,9 +4556,7 @@ class TestProcessSSOJWTAccessToken: # Assert mock_jwt_decode.assert_called_once() - mock_jwt_handler.get_team_ids_from_jwt.assert_called_once_with( - sample_jwt_payload - ) + mock_jwt_handler.get_team_ids_from_jwt.assert_called_once_with(sample_jwt_payload) # Even empty team IDs should be set assert result.team_ids == [] @@ -4600,9 +4593,7 @@ class TestProcessSSOJWTAccessToken: assert result.team_ids == ["existing_team"] assert result.user_role is None - def test_process_sso_jwt_access_token_real_jwt_with_role_and_teams( - self, mock_jwt_handler - ): + def test_process_sso_jwt_access_token_real_jwt_with_role_and_teams(self, mock_jwt_handler): """Test that a real JWT containing role and team fields is correctly processed.""" import jwt as pyjwt @@ -4775,15 +4766,9 @@ class TestGenericResponseConvertorNestedAttributes: # Current behavior: returns None for nested paths # Expected behavior with current implementation (no nested path support): assert result.id == "nested-user-456" - assert ( - result.email == "john.doe@example.com" - ) # Can't access "attributes.email" with .get() - assert ( - result.first_name == "John" - ) # Can't access "attributes.given_name" with .get() - assert ( - result.last_name == "Doe" - ) # Can't access "attributes.family_name" with .get() + assert result.email == "john.doe@example.com" # Can't access "attributes.email" with .get() + assert result.first_name == "John" # Can't access "attributes.given_name" with .get() + assert result.last_name == "Doe" # Can't access "attributes.family_name" with .get() assert result.display_name == "user-sub-123" # Top-level attribute works @@ -4959,14 +4944,8 @@ class TestGetGenericSSORedirectParams: assert redirect_params["code_challenge_method"] == "S256" # Verify code_challenge is correctly derived from code_verifier - expected_challenge_bytes = hashlib.sha256( - code_verifier.encode("utf-8") - ).digest() - expected_challenge = ( - base64.urlsafe_b64encode(expected_challenge_bytes) - .decode("utf-8") - .rstrip("=") - ) + expected_challenge_bytes = hashlib.sha256(code_verifier.encode("utf-8")).digest() + expected_challenge = base64.urlsafe_b64encode(expected_challenge_bytes).decode("utf-8").rstrip("=") assert redirect_params["code_challenge"] == expected_challenge def test_state_with_pkce_disabled(self): @@ -5099,14 +5078,8 @@ class TestPKCEFunctionality: assert isinstance(code_verifier, str) # Verify code_challenge is correctly generated from code_verifier - expected_challenge_bytes = hashlib.sha256( - code_verifier.encode("utf-8") - ).digest() - expected_challenge = ( - base64.urlsafe_b64encode(expected_challenge_bytes) - .decode("utf-8") - .rstrip("=") - ) + expected_challenge_bytes = hashlib.sha256(code_verifier.encode("utf-8")).digest() + expected_challenge = base64.urlsafe_b64encode(expected_challenge_bytes).decode("utf-8").rstrip("=") assert code_challenge == expected_challenge # Verify both are base64url encoded (no padding) @@ -5128,9 +5101,7 @@ class TestPKCEFunctionality: # Mock cache with async methods — use dict format (primary path) mock_cache = MagicMock(redis_cache=None) test_code_verifier = "test_code_verifier_abc123xyz" - mock_cache.async_get_cache = AsyncMock( - return_value={"code_verifier": test_code_verifier} - ) + mock_cache.async_get_cache = AsyncMock(return_value={"code_verifier": test_code_verifier}) mock_cache.async_delete_cache = AsyncMock() with ( @@ -5139,10 +5110,8 @@ class TestPKCEFunctionality: patch.dict(os.environ, {"GENERIC_CLIENT_USE_PKCE": "true"}), ): # Act - token_params = ( - await SSOAuthenticationHandler.prepare_token_exchange_parameters( - request=mock_request, generic_include_client_id=False - ) + token_params = await SSOAuthenticationHandler.prepare_token_exchange_parameters( + request=mock_request, generic_include_client_id=False ) # Assert @@ -5153,9 +5122,7 @@ class TestPKCEFunctionality: # Verify cache was read but NOT deleted yet (deletion is deferred to after # successful token exchange to preserve the verifier for retries) - mock_cache.async_get_cache.assert_called_once_with( - key=f"pkce_verifier:{test_state}" - ) + mock_cache.async_get_cache.assert_called_once_with(key=f"pkce_verifier:{test_state}") mock_cache.async_delete_cache.assert_not_called() @pytest.mark.asyncio @@ -5168,9 +5135,7 @@ class TestPKCEFunctionality: # Mock SSO provider mock_sso = MagicMock() mock_redirect_response = MagicMock() - original_location = ( - "https://auth.example.com/authorize?state=test456&client_id=abc" - ) + original_location = "https://auth.example.com/authorize?state=test456&client_id=abc" mock_redirect_response.headers = {"location": original_location} mock_sso.get_login_redirect = AsyncMock(return_value=mock_redirect_response) mock_sso.__enter__ = MagicMock(return_value=mock_sso) @@ -5256,9 +5221,7 @@ class TestPKCEFunctionality: with patch.dict(os.environ, {"GENERIC_CLIENT_USE_PKCE": "true"}): with patch("litellm.proxy.proxy_server.redis_usage_cache", mock_redis): - with patch( - "litellm.proxy.proxy_server.user_api_key_cache", mock_in_memory - ): + with patch("litellm.proxy.proxy_server.user_api_key_cache", mock_in_memory): # Pod A: start login, store code_verifier in "Redis" await SSOAuthenticationHandler.get_generic_sso_redirect_response( generic_sso=mock_sso, @@ -5272,9 +5235,7 @@ class TestPKCEFunctionality: stored_value = mock_redis._store[stored_key] # Stored as JSON-serialized dict for Redis compatibility stored_dict = json.loads(stored_value) - assert ( - isinstance(stored_dict, dict) and "code_verifier" in stored_dict - ) + assert isinstance(stored_dict, dict) and "code_verifier" in stored_dict assert len(stored_dict["code_verifier"]) == 43 # Pod B: callback with same state, retrieve from "Redis" @@ -5329,9 +5290,7 @@ class TestPKCEFunctionality: with patch.dict(os.environ, {"GENERIC_CLIENT_USE_PKCE": "true"}): with patch("litellm.proxy.proxy_server.redis_usage_cache", None): - with patch( - "litellm.proxy.proxy_server.user_api_key_cache", mock_in_memory - ): + with patch("litellm.proxy.proxy_server.user_api_key_cache", mock_in_memory): # Pod A: start login, store code_verifier in in-memory cache await SSOAuthenticationHandler.get_generic_sso_redirect_response( generic_sso=mock_sso, @@ -5340,14 +5299,9 @@ class TestPKCEFunctionality: ) mock_in_memory.async_set_cache.assert_called_once() stored_key = mock_in_memory.async_set_cache.call_args.kwargs["key"] - stored_value = mock_in_memory.async_set_cache.call_args.kwargs[ - "value" - ] + stored_value = mock_in_memory.async_set_cache.call_args.kwargs["value"] assert stored_key == "pkce_verifier:fallback_state_xyz" - assert ( - isinstance(stored_value, dict) - and len(stored_value["code_verifier"]) == 43 - ) + assert isinstance(stored_value, dict) and len(stored_value["code_verifier"]) == 43 # Same pod: callback retrieves from in-memory cache mock_request = MagicMock(spec=Request) @@ -5356,14 +5310,10 @@ class TestPKCEFunctionality: request=mock_request, generic_include_client_id=False ) assert "code_verifier" in token_params - assert ( - token_params["code_verifier"] == stored_value["code_verifier"] - ) + assert token_params["code_verifier"] == stored_value["code_verifier"] # Cache key returned for deferred deletion after successful exchange assert token_params["_pkce_cache_key"] == stored_key - mock_in_memory.async_get_cache.assert_called_once_with( - key=stored_key - ) + mock_in_memory.async_get_cache.assert_called_once_with(key=stored_key) # Deletion is deferred — not called by prepare_token_exchange_parameters mock_in_memory.async_delete_cache.assert_not_called() @@ -5385,10 +5335,8 @@ class TestPKCEFunctionality: ): mock_request = MagicMock(spec=Request) mock_request.query_params = {} - token_params = ( - await SSOAuthenticationHandler.prepare_token_exchange_parameters( - request=mock_request, generic_include_client_id=False - ) + token_params = await SSOAuthenticationHandler.prepare_token_exchange_parameters( + request=mock_request, generic_include_client_id=False ) assert "code_verifier" not in token_params mock_redis.async_get_cache.assert_not_called() @@ -5424,12 +5372,10 @@ class TestPKCEFunctionality: # Verify redirect_uri is forwarded (required by strict OAuth providers) assert post_data.get("redirect_uri") == "https://proxy.example.com/callback" # Verify credentials are NOT double-sent in the POST body when using Basic Auth - assert ( - "client_secret" not in post_data - ), "client_secret must not appear in POST body when using Basic Auth" - assert ( - "client_id" not in post_data - ), "client_id must not appear in POST body when using Basic Auth (include_client_id=False)" + assert "client_secret" not in post_data, "client_secret must not appear in POST body when using Basic Auth" + assert "client_id" not in post_data, ( + "client_id must not appear in POST body when using Basic Auth (include_client_id=False)" + ) return mock_response # get_async_httpx_client returns a client directly (no context manager). @@ -5439,9 +5385,7 @@ class TestPKCEFunctionality: mock_userinfo_client = MagicMock() mock_userinfo_client.get = AsyncMock(return_value=mock_userinfo_response) - with patch( - "litellm.proxy.management_endpoints.ui_sso.get_async_httpx_client" - ) as mock_get_client: + with patch("litellm.proxy.management_endpoints.ui_sso.get_async_httpx_client") as mock_get_client: mock_get_client.side_effect = [mock_token_client, mock_userinfo_client] result = await SSOAuthenticationHandler._pkce_token_exchange( @@ -5460,9 +5404,7 @@ class TestPKCEFunctionality: assert result["email"] == "user@example.com" # id_token was explicit null in token_response — the merge loop must remove it # rather than leaving "id_token": None in the result. - assert ( - "id_token" not in result - ), "null id_token from token endpoint must be absent in merged result" + assert "id_token" not in result, "null id_token from token endpoint must be absent in merged result" # Verify userinfo GET used the correct Bearer token header get_call = mock_userinfo_client.get.call_args assert get_call is not None @@ -5482,18 +5424,12 @@ class TestPKCEFunctionality: async def fake_post(*args, **kwargs): headers = kwargs.get("headers", {}) auth_header = headers.get("Authorization", "") - assert not auth_header.startswith( - "Basic " - ), "Should NOT use Basic Auth when include_client_id=True" + assert not auth_header.startswith("Basic "), "Should NOT use Basic Auth when include_client_id=True" data = kwargs.get("data", {}) assert "client_id" in data assert "client_secret" in data - assert ( - data.get("code_verifier") == "verifier_xyz" - ), "code_verifier must be in POST body" - assert ( - data.get("redirect_uri") == "https://proxy.example.com/callback" - ), "redirect_uri must be forwarded" + assert data.get("code_verifier") == "verifier_xyz", "code_verifier must be in POST body" + assert data.get("redirect_uri") == "https://proxy.example.com/callback", "redirect_uri must be forwarded" mock = MagicMock() mock.status_code = 200 mock.json.return_value = token_resp @@ -5509,9 +5445,7 @@ class TestPKCEFunctionality: mock_userinfo_client = MagicMock() mock_userinfo_client.get = AsyncMock(return_value=mock_userinfo) - with patch( - "litellm.proxy.management_endpoints.ui_sso.get_async_httpx_client" - ) as mock_get_client: + with patch("litellm.proxy.management_endpoints.ui_sso.get_async_httpx_client") as mock_get_client: mock_get_client.side_effect = [mock_token_client, mock_userinfo_client] result = await SSOAuthenticationHandler._pkce_token_exchange( @@ -5543,9 +5477,7 @@ class TestPKCEFunctionality: "error_description": "Code already used", } - with patch( - "litellm.proxy.management_endpoints.ui_sso.get_async_httpx_client" - ) as mock_get_client: + with patch("litellm.proxy.management_endpoints.ui_sso.get_async_httpx_client") as mock_get_client: mock_client = MagicMock() mock_resp = MagicMock() mock_resp.status_code = 200 @@ -5577,16 +5509,10 @@ class TestPKCEFunctionality: payload = {"sub": "user_from_jwt", "email": "jwt@example.com"} # Build a minimal JWT (header.payload.signature — signature not verified) - encoded_payload = ( - base64.urlsafe_b64encode(_json.dumps(payload).encode()) - .rstrip(b"=") - .decode() - ) + encoded_payload = base64.urlsafe_b64encode(_json.dumps(payload).encode()).rstrip(b"=").decode() fake_id_token = f"eyJhbGciOiJSUzI1NiJ9.{encoded_payload}.fakesig" - with patch( - "litellm.proxy.management_endpoints.ui_sso.get_async_httpx_client" - ) as mock_get_client: + with patch("litellm.proxy.management_endpoints.ui_sso.get_async_httpx_client") as mock_get_client: mock_client = MagicMock() mock_fail = MagicMock() mock_fail.status_code = 503 @@ -5612,11 +5538,7 @@ class TestPKCEFunctionality: from litellm.proxy.management_endpoints.ui_sso import SSOAuthenticationHandler payload = {"sub": "id_token_user", "email": "id@example.com"} - encoded_payload = ( - base64.urlsafe_b64encode(_json.dumps(payload).encode()) - .rstrip(b"=") - .decode() - ) + encoded_payload = base64.urlsafe_b64encode(_json.dumps(payload).encode()).rstrip(b"=").decode() fake_id_token = f"eyJhbGciOiJSUzI1NiJ9.{encoded_payload}.fakesig" # No httpx call should happen when userinfo_endpoint is None @@ -5636,9 +5558,7 @@ class TestPKCEFunctionality: from litellm.proxy._types import ProxyException from litellm.proxy.management_endpoints.ui_sso import SSOAuthenticationHandler - with patch( - "litellm.proxy.management_endpoints.ui_sso.get_async_httpx_client" - ) as mock_get_client: + with patch("litellm.proxy.management_endpoints.ui_sso.get_async_httpx_client") as mock_get_client: mock_client = MagicMock() mock_fail = MagicMock() mock_fail.status_code = 503 @@ -5667,9 +5587,7 @@ class TestPKCEFunctionality: mock_resp.status_code = 200 mock_resp.json.return_value = None # HTTP 200 with null JSON body - with patch( - "litellm.proxy.management_endpoints.ui_sso.get_async_httpx_client" - ) as mock_get_client: + with patch("litellm.proxy.management_endpoints.ui_sso.get_async_httpx_client") as mock_get_client: mock_client = MagicMock() mock_client.get = AsyncMock(return_value=mock_resp) mock_get_client.return_value = mock_client @@ -5720,10 +5638,7 @@ class TestPKCEFunctionality: request=mock_request, generic_include_client_id=False ) - assert ( - "verifier not found" in exc_info.value.message.lower() - or "cache" in exc_info.value.message.lower() - ) + assert "verifier not found" in exc_info.value.message.lower() or "cache" in exc_info.value.message.lower() assert str(exc_info.value.code) == "401" @pytest.mark.asyncio @@ -5742,14 +5657,10 @@ class TestPKCEFunctionality: async def fake_post(*args, **kwargs): headers = kwargs.get("headers", {}) auth_header = headers.get("Authorization", "") - assert not auth_header.startswith( - "Basic " - ), "Public client must not use Basic Auth" + assert not auth_header.startswith("Basic "), "Public client must not use Basic Auth" data = kwargs.get("data", {}) assert data.get("client_id") == "public_client_id" - assert ( - "client_secret" not in data - ), "No secret should be sent for public client" + assert "client_secret" not in data, "No secret should be sent for public client" assert data.get("code_verifier") == "public_verifier" mock = MagicMock() mock.status_code = 200 @@ -5766,9 +5677,7 @@ class TestPKCEFunctionality: mock_userinfo_client = MagicMock() mock_userinfo_client.get = AsyncMock(return_value=mock_userinfo) - with patch( - "litellm.proxy.management_endpoints.ui_sso.get_async_httpx_client" - ) as mock_get_client: + with patch("litellm.proxy.management_endpoints.ui_sso.get_async_httpx_client") as mock_get_client: mock_get_client.side_effect = [mock_token_client, mock_userinfo_client] result = await SSOAuthenticationHandler._pkce_token_exchange( @@ -5796,22 +5705,16 @@ class TestPKCEFunctionality: from litellm.proxy.management_endpoints.ui_sso import SSOAuthenticationHandler failing_cache = MagicMock() - failing_cache.async_delete_cache = AsyncMock( - side_effect=Exception("Redis down") - ) + failing_cache.async_delete_cache = AsyncMock(side_effect=Exception("Redis down")) # Should NOT raise even though the underlying cache delete fails with ( patch("litellm.proxy.proxy_server.redis_usage_cache", None), patch("litellm.proxy.proxy_server.user_api_key_cache", failing_cache), ): - await SSOAuthenticationHandler._delete_pkce_verifier( - "pkce_verifier:test_state" - ) + await SSOAuthenticationHandler._delete_pkce_verifier("pkce_verifier:test_state") - failing_cache.async_delete_cache.assert_called_once_with( - key="pkce_verifier:test_state" - ) + failing_cache.async_delete_cache.assert_called_once_with(key="pkce_verifier:test_state") @pytest.mark.asyncio async def test_pkce_cache_miss_unexpected_format_raises_proxy_exception(self): @@ -5899,8 +5802,7 @@ class TestPKCEFunctionality: mock_cache.async_get_cache.assert_called_once() # Verify the warning was actually logged assert any( - "verifier not found" in r.message.lower() - or "code_verifier" in r.message.lower() + "verifier not found" in r.message.lower() or "code_verifier" in r.message.lower() for r in caplog.records if r.levelno >= logging.WARNING ), f"Expected a cache-miss warning. Records: {[r.message for r in caplog.records]}" @@ -5918,9 +5820,7 @@ class TestPKCEFunctionality: mock_response.status_code = 401 mock_response.text = "Unauthorized" - with patch( - "litellm.proxy.management_endpoints.ui_sso.get_async_httpx_client" - ) as mock_get_client: + with patch("litellm.proxy.management_endpoints.ui_sso.get_async_httpx_client") as mock_get_client: mock_client = MagicMock() mock_client.post = AsyncMock(return_value=mock_response) mock_get_client.return_value = mock_client @@ -5942,9 +5842,7 @@ class TestPKCEFunctionality: assert str(exc_info.value.code) == "401" @pytest.mark.asyncio - async def test_pkce_cache_miss_unexpected_format_non_strict_logs_warning( - self, caplog - ): + async def test_pkce_cache_miss_unexpected_format_non_strict_logs_warning(self, caplog): """When cached data has an unexpected format (e.g. integer from corrupt Redis) in non-strict mode, prepare_token_exchange_parameters logs a warning and returns params without code_verifier rather than raising.""" @@ -5988,9 +5886,7 @@ class TestPKCEFunctionality: mock_cache.async_get_cache.assert_called_once() # Verify a warning was logged about the unexpected format or cache miss assert any( - "verifier" in r.message.lower() - or "format" in r.message.lower() - or "cache" in r.message.lower() + "verifier" in r.message.lower() or "format" in r.message.lower() or "cache" in r.message.lower() for r in caplog.records if r.levelno >= logging.WARNING ), f"Expected a format/cache warning. Records: {[r.message for r in caplog.records]}" @@ -6034,9 +5930,7 @@ class TestPKCEFunctionality: from litellm.proxy._types import ProxyException from litellm.proxy.management_endpoints.ui_sso import SSOAuthenticationHandler - with patch( - "litellm.proxy.management_endpoints.ui_sso.get_async_httpx_client" - ) as mock_get_client: + with patch("litellm.proxy.management_endpoints.ui_sso.get_async_httpx_client") as mock_get_client: mock_client = MagicMock() mock_resp = MagicMock() mock_resp.status_code = 200 @@ -6070,9 +5964,7 @@ class TestPKCEFunctionality: body_without_token = {"token_type": "Bearer", "scope": "openid"} - with patch( - "litellm.proxy.management_endpoints.ui_sso.get_async_httpx_client" - ) as mock_get_client: + with patch("litellm.proxy.management_endpoints.ui_sso.get_async_httpx_client") as mock_get_client: mock_client = MagicMock() mock_resp = MagicMock() mock_resp.status_code = 200 @@ -6093,10 +5985,7 @@ class TestPKCEFunctionality: additional_headers={}, ) - assert ( - "no access_token" in exc_info.value.message - or "access_token" in exc_info.value.message - ) + assert "no access_token" in exc_info.value.message or "access_token" in exc_info.value.message assert str(exc_info.value.code) == "401" @@ -6124,18 +6013,16 @@ class TestAddMissingTeamMember: sso_teams = ["team-from-entra-1", "team-from-entra-2"] - with patch( - "litellm.proxy.management_endpoints.ui_sso.create_team_member_add_task" - ) as mock_add_task: + with patch("litellm.proxy.management_endpoints.ui_sso.create_team_member_add_task") as mock_add_task: mock_add_task.return_value = AsyncMock() await add_missing_team_member(user_info=new_user, sso_teams=sso_teams) # Bug: This assertion currently FAILS - no teams are added # because function returns early when teams is None - assert ( - mock_add_task.call_count == 2 - ), f"Expected 2 calls to add user to teams, but got {mock_add_task.call_count}" + assert mock_add_task.call_count == 2, ( + f"Expected 2 calls to add user to teams, but got {mock_add_task.call_count}" + ) called_team_ids = [call.args[0] for call in mock_add_task.call_args_list] assert set(called_team_ids) == { "team-from-entra-1", @@ -6161,9 +6048,7 @@ class TestAddMissingTeamMember: sso_teams = ["team-from-entra-1", "team-from-entra-2"] - with patch( - "litellm.proxy.management_endpoints.ui_sso.create_team_member_add_task" - ) as mock_add_task: + with patch("litellm.proxy.management_endpoints.ui_sso.create_team_member_add_task") as mock_add_task: mock_add_task.return_value = AsyncMock() await add_missing_team_member(user_info=existing_user, sso_teams=sso_teams) @@ -6197,9 +6082,7 @@ class TestAddMissingTeamMember: teams=None, # Bug: NewUserResponse defaults to None ) - with patch( - "litellm.proxy.management_endpoints.ui_sso.add_missing_team_member" - ) as mock_add_member: + with patch("litellm.proxy.management_endpoints.ui_sso.add_missing_team_member") as mock_add_member: await SSOAuthenticationHandler.add_user_to_teams_from_sso_response( result=sso_result, user_info=new_user_info, @@ -6241,9 +6124,7 @@ class TestAddMissingTeamMember: # Bug: With current code, team_member_calls will be empty # After fix: Should have 2 entries - assert ( - len(team_member_calls) == 2 - ), f"Expected 2 teams to be added, but got {len(team_member_calls)}" + assert len(team_member_calls) == 2, f"Expected 2 teams to be added, but got {len(team_member_calls)}" assert {c["team_id"] for c in team_member_calls} == { "entra-team-alpha", "entra-team-beta", @@ -6299,9 +6180,9 @@ class TestAddMissingTeamMember: ): await add_missing_team_member(user_info=user_info, sso_teams=sso_teams) - assert set(added_teams) == set( - expected_teams_added - ), f"Expected teams {expected_teams_added}, but got {added_teams}" + assert set(added_teams) == set(expected_teams_added), ( + f"Expected teams {expected_teams_added}, but got {added_teams}" + ) @pytest.mark.asyncio @@ -6367,17 +6248,15 @@ async def test_role_mappings_override_default_internal_user_params(): new_user_request = call_args.kwargs["data"] # The role from SSO should be preserved, not overridden by default_internal_user_params - assert ( - new_user_request.user_role == "proxy_admin" - ), "SSO-extracted role should override default_internal_user_params role" + assert new_user_request.user_role == "proxy_admin", ( + "SSO-extracted role should override default_internal_user_params role" + ) # Other default params should still be applied - assert ( - new_user_request.max_budget == 100 - ), "max_budget from default_internal_user_params should be applied" - assert ( - new_user_request.budget_duration == "30d" - ), "budget_duration from default_internal_user_params should be applied" + assert new_user_request.max_budget == 100, "max_budget from default_internal_user_params should be applied" + assert new_user_request.budget_duration == "30d", ( + "budget_duration from default_internal_user_params should be applied" + ) finally: # Restore original default_internal_user_params (always assign, never delattr — @@ -6448,19 +6327,17 @@ async def test_sso_role_preserved_without_role_mappings(): new_user_request = call_args.kwargs["data"] # SSO role should be preserved even without role_mappings configured - assert ( - new_user_request.user_role == "proxy_admin" - ), "SSO role from app_roles should not be overwritten by default_internal_user_params" + assert new_user_request.user_role == "proxy_admin", ( + "SSO role from app_roles should not be overwritten by default_internal_user_params" + ) # Other defaults should still apply - assert ( - new_user_request.max_budget == 50 - ), "max_budget from default_internal_user_params should be applied" + assert new_user_request.max_budget == 50, "max_budget from default_internal_user_params should be applied" # Verify user_defined_values was also updated (it's mutated in-place) - assert ( - user_defined_values["user_role"] == "proxy_admin" - ), "user_defined_values should retain the SSO role after insert_sso_user" + assert user_defined_values["user_role"] == "proxy_admin", ( + "user_defined_values should retain the SSO role after insert_sso_user" + ) finally: # Restore original default_internal_user_params (always assign, never delattr — @@ -6569,10 +6446,7 @@ class TestSSOReadinessEndpoint: assert data["sso_configured"] is True assert data["provider"] == "google" assert "GOOGLE_CLIENT_SECRET" in data["missing_environment_variables"] - assert ( - "Google SSO is configured but missing required environment variables" - in data["message"] - ) + assert "Google SSO is configured but missing required environment variables" in data["message"] finally: app.dependency_overrides.clear() @@ -6633,9 +6507,7 @@ class TestSSOReadinessEndpoint: assert data["sso_configured"] is True assert data["provider"] == expected_provider assert data["status"] == "unhealthy" - assert set(data["missing_environment_variables"]) == set( - expected_missing_vars - ) + assert set(data["missing_environment_variables"]) == set(expected_missing_vars) finally: app.dependency_overrides.clear() @@ -6703,9 +6575,7 @@ class TestSSOReadinessEndpoint: assert data["sso_configured"] is True assert data["provider"] == expected_provider assert data["status"] == "unhealthy" - assert set(data["missing_environment_variables"]) == set( - expected_missing_vars - ) + assert set(data["missing_environment_variables"]) == set(expected_missing_vars) finally: app.dependency_overrides.clear() @@ -6737,13 +6607,9 @@ class TestCustomMicrosoftSSO: discovery = await sso.get_discovery_document() assert ( - discovery["authorization_endpoint"] - == "https://login.microsoftonline.com/test-tenant/oauth2/v2.0/authorize" - ) - assert ( - discovery["token_endpoint"] - == "https://login.microsoftonline.com/test-tenant/oauth2/v2.0/token" + discovery["authorization_endpoint"] == "https://login.microsoftonline.com/test-tenant/oauth2/v2.0/authorize" ) + assert discovery["token_endpoint"] == "https://login.microsoftonline.com/test-tenant/oauth2/v2.0/token" assert discovery["userinfo_endpoint"] == "https://graph.microsoft.com/v1.0/me" @pytest.mark.asyncio @@ -6807,13 +6673,8 @@ class TestCustomMicrosoftSSO: # Custom auth endpoint assert discovery["authorization_endpoint"] == custom_auth_endpoint # Default token and userinfo endpoints - assert ( - discovery["token_endpoint"] - == "https://login.microsoftonline.com/test-tenant/oauth2/v2.0/token" - ) - assert ( - discovery["userinfo_endpoint"] == "https://graph.microsoft.com/v1.0/me" - ) + assert discovery["token_endpoint"] == "https://login.microsoftonline.com/test-tenant/oauth2/v2.0/token" + assert discovery["userinfo_endpoint"] == "https://graph.microsoft.com/v1.0/me" def test_custom_microsoft_sso_uses_common_tenant_when_none(self): """ @@ -6851,9 +6712,7 @@ async def test_setup_team_mappings(): mock_prisma = MagicMock() mock_sso_config = MagicMock() mock_sso_config.sso_settings = {"team_mappings": {"team_ids_jwt_field": "groups"}} - mock_prisma.db.litellm_ssoconfig.find_unique = AsyncMock( - return_value=mock_sso_config - ) + mock_prisma.db.litellm_ssoconfig.find_unique = AsyncMock(return_value=mock_sso_config) with patch( "litellm.proxy.utils.get_prisma_client_or_throw", @@ -6866,9 +6725,7 @@ async def test_setup_team_mappings(): assert result is not None assert isinstance(result, TeamMappings) assert result.team_ids_jwt_field == "groups" - mock_prisma.db.litellm_ssoconfig.find_unique.assert_called_once_with( - where={"id": "sso_config"} - ) + mock_prisma.db.litellm_ssoconfig.find_unique.assert_called_once_with(where={"id": "sso_config"}) # ============================================================================ @@ -7148,9 +7005,7 @@ def test_process_sso_jwt_access_token_extracts_role_from_nested_field(): user_role=None, ) - with patch.dict( - os.environ, {"GENERIC_USER_ROLE_ATTRIBUTE": "resource_access.my-client.roles"} - ): + with patch.dict(os.environ, {"GENERIC_USER_ROLE_ATTRIBUTE": "resource_access.my-client.roles"}): process_sso_jwt_access_token( access_token_str=access_token_str, sso_jwt_handler=None, @@ -7211,9 +7066,7 @@ def test_generic_response_convertor_with_extra_attributes(monkeypatch): from litellm.proxy.management_endpoints.ui_sso import generic_response_convertor monkeypatch.setenv("GENERIC_CLIENT_ID", "test_client") - monkeypatch.setenv( - "GENERIC_USER_EXTRA_ATTRIBUTES", "custom_field1,custom_field2,custom_field3" - ) + monkeypatch.setenv("GENERIC_USER_EXTRA_ATTRIBUTES", "custom_field1,custom_field2,custom_field3") mock_response = { "sub": "user-id-123", @@ -7279,9 +7132,7 @@ def test_generic_response_convertor_extra_attributes_with_nested_paths(monkeypat from litellm.proxy.management_endpoints.ui_sso import generic_response_convertor monkeypatch.setenv("GENERIC_CLIENT_ID", "test_client") - monkeypatch.setenv( - "GENERIC_USER_EXTRA_ATTRIBUTES", "org_info.department,org_info.manager" - ) + monkeypatch.setenv("GENERIC_USER_EXTRA_ATTRIBUTES", "org_info.department,org_info.manager") mock_response = { "sub": "user-id-123", @@ -7443,10 +7294,7 @@ class TestCliSsoAttributionMetadata: get_user_info_mock.assert_awaited_once() assert get_user_info_mock.call_args.kwargs["user_defined_values"] is not None - assert ( - get_user_info_mock.call_args.kwargs["user_defined_values"]["user_id"] - == "cli-test-user" - ) + assert get_user_info_mock.call_args.kwargs["user_defined_values"]["user_id"] == "cli-test-user" @pytest.mark.asyncio async def test_cli_sso_callback_rejects_restricted_sso_group(self): @@ -7581,13 +7429,9 @@ class TestCliSsoAttributionMetadata: ) flow_data = mock_cache.set_cache.call_args.kwargs["value"] - assert flow_data["session_data"]["attribution_metadata"] == { - "acme_employment_type": "contractor" - } + assert flow_data["session_data"]["attribution_metadata"] == {"acme_employment_type": "contractor"} mock_prisma.db.litellm_usertable.update_many.assert_awaited_once() - update_data = mock_prisma.db.litellm_usertable.update_many.call_args.kwargs[ - "data" - ] + update_data = mock_prisma.db.litellm_usertable.update_many.call_args.kwargs["data"] assert update_data["metadata"]["acme_employment_type"] == "contractor" assert update_data["metadata"]["auth_provider"] == "generic" @@ -7639,9 +7483,7 @@ class TestValidateReturnTo: def test_returns_false_when_no_control_plane_url_configured(self, monkeypatch): """return_to should be silently ignored if control_plane_url is not in general_settings.""" monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {}) - result = SSOAuthenticationHandler._validate_return_to( - "https://cp.example.com/ui" - ) + result = SSOAuthenticationHandler._validate_return_to("https://cp.example.com/ui") assert result is False def test_allows_matching_origin(self, monkeypatch): @@ -7651,9 +7493,7 @@ class TestValidateReturnTo: {"control_plane_url": "https://cp.example.com"}, ) # Should not raise - SSOAuthenticationHandler._validate_return_to( - "https://cp.example.com/ui?page=models" - ) + SSOAuthenticationHandler._validate_return_to("https://cp.example.com/ui?page=models") def test_allows_matching_origin_with_trailing_slash(self, monkeypatch): """Trailing slash on control_plane_url should not affect origin comparison.""" @@ -7670,9 +7510,7 @@ class TestValidateReturnTo: {"control_plane_url": "https://cp.example.com"}, ) with pytest.raises(HTTPException) as exc_info: - SSOAuthenticationHandler._validate_return_to( - "https://cp.example.com.evil.com/steal" - ) + SSOAuthenticationHandler._validate_return_to("https://cp.example.com.evil.com/steal") assert exc_info.value.status_code == 400 def test_rejects_different_origin(self, monkeypatch): @@ -7711,9 +7549,7 @@ class TestValidateReturnTo: {"control_plane_url": "https://cp.example.com"}, ) with pytest.raises(HTTPException) as exc_info: - SSOAuthenticationHandler._validate_return_to( - "https://cp.example.com:8443/ui" - ) + SSOAuthenticationHandler._validate_return_to("https://cp.example.com:8443/ui") assert exc_info.value.status_code == 400 def test_allows_explicit_default_port(self, monkeypatch): @@ -7839,9 +7675,7 @@ class TestSyncUserRoleFromJwtRoleMap: user_id=user_id, user_role=LitellmUserRoles.INTERNAL_USER_VIEW_ONLY.value, ) - await cache.async_set_cache( - key=user_id, value=existing_user.model_dump(), ttl=60 - ) + await cache.async_set_cache(key=user_id, value=existing_user.model_dump(), ttl=60) sso_values = self._make_sso_values( user_role=LitellmUserRoles.INTERNAL_USER_VIEW_ONLY.value, @@ -7922,9 +7756,7 @@ class TestPKCEStateCookieBinding: SSOAuthenticationHandler, ) - mock_redirect = RedirectResponse( - url="https://idp.example.com/authorize?state=test-state-xyz" - ) + mock_redirect = RedirectResponse(url="https://idp.example.com/authorize?state=test-state-xyz") mock_generic_sso = MagicMock() mock_generic_sso.__enter__ = MagicMock(return_value=mock_generic_sso) mock_generic_sso.__exit__ = MagicMock(return_value=None) @@ -7945,12 +7777,8 @@ class TestPKCEStateCookieBinding: assert response is not None cookie_headers = response.headers.getlist("set-cookie") - cookie_str = next( - (c for c in cookie_headers if "litellm_oauth_state=" in c), None - ) - assert ( - cookie_str is not None - ), f"litellm_oauth_state cookie not set; got: {cookie_headers}" + cookie_str = next((c for c in cookie_headers if "litellm_oauth_state=" in c), None) + assert cookie_str is not None, f"litellm_oauth_state cookie not set; got: {cookie_headers}" assert "test-state-xyz" in cookie_str assert "HttpOnly" in cookie_str assert "SameSite=lax" in cookie_str @@ -7972,9 +7800,7 @@ class TestPKCEStateCookieBinding: SSOAuthenticationHandler, ) - mock_redirect = RedirectResponse( - url="https://idp.example.com/authorize?state=test-state-xyz" - ) + mock_redirect = RedirectResponse(url="https://idp.example.com/authorize?state=test-state-xyz") mock_generic_sso = MagicMock() mock_generic_sso.__enter__ = MagicMock(return_value=mock_generic_sso) mock_generic_sso.__exit__ = MagicMock(return_value=None) @@ -7995,9 +7821,9 @@ class TestPKCEStateCookieBinding: assert response is not None cookie_headers = response.headers.getlist("set-cookie") - assert not any( - "litellm_oauth_state=" in c for c in cookie_headers - ), f"litellm_oauth_state cookie set on non-PKCE flow; got: {cookie_headers}" + assert not any("litellm_oauth_state=" in c for c in cookie_headers), ( + f"litellm_oauth_state cookie set on non-PKCE flow; got: {cookie_headers}" + ) @pytest.mark.asyncio async def test_redirect_response_drops_secure_flag_for_http_dev(self): @@ -8010,9 +7836,7 @@ class TestPKCEStateCookieBinding: SSOAuthenticationHandler, ) - mock_redirect = RedirectResponse( - url="http://idp.local/authorize?state=local-dev-state" - ) + mock_redirect = RedirectResponse(url="http://idp.local/authorize?state=local-dev-state") mock_generic_sso = MagicMock() mock_generic_sso.__enter__ = MagicMock(return_value=mock_generic_sso) mock_generic_sso.__exit__ = MagicMock(return_value=None) @@ -8036,16 +7860,12 @@ class TestPKCEStateCookieBinding: ) cookie_headers = response.headers.getlist("set-cookie") - cookie_str = next( - (c for c in cookie_headers if "litellm_oauth_state=" in c), None - ) + cookie_str = next((c for c in cookie_headers if "litellm_oauth_state=" in c), None) assert cookie_str is not None assert "Secure" not in cookie_str @pytest.mark.asyncio - async def test_redirect_response_sets_secure_flag_behind_trusted_tls_terminating_proxy( - self, monkeypatch - ): + async def test_redirect_response_sets_secure_flag_behind_trusted_tls_terminating_proxy(self, monkeypatch): """Regression: litellm sees a plain-HTTP hop when TLS terminates at a reverse proxy. The Secure flag must still be set when the direct peer is a configured trusted proxy and it reports X-Forwarded-Proto: https -- but NOT from an @@ -8057,9 +7877,7 @@ class TestPKCEStateCookieBinding: SSOAuthenticationHandler, ) - mock_redirect = RedirectResponse( - url="http://idp.internal/authorize?state=behind-proxy-state" - ) + mock_redirect = RedirectResponse(url="http://idp.internal/authorize?state=behind-proxy-state") mock_generic_sso = MagicMock() mock_generic_sso.__enter__ = MagicMock(return_value=mock_generic_sso) mock_generic_sso.__exit__ = MagicMock(return_value=None) @@ -8091,16 +7909,12 @@ class TestPKCEStateCookieBinding: ) cookie_headers = response.headers.getlist("set-cookie") - cookie_str = next( - (c for c in cookie_headers if "litellm_oauth_state=" in c), None - ) + cookie_str = next((c for c in cookie_headers if "litellm_oauth_state=" in c), None) assert cookie_str is not None assert "Secure" in cookie_str @pytest.mark.asyncio - async def test_redirect_response_ignores_spoofed_forwarded_proto_without_trust_config( - self, monkeypatch - ): + async def test_redirect_response_ignores_spoofed_forwarded_proto_without_trust_config(self, monkeypatch): """The same X-Forwarded-Proto: https header must NOT flip Secure on when no trusted-proxy config is present -- honoring it unconditionally would let any client spoof the header and would not itself be the vulnerability the ticket @@ -8111,9 +7925,7 @@ class TestPKCEStateCookieBinding: SSOAuthenticationHandler, ) - mock_redirect = RedirectResponse( - url="http://idp.internal/authorize?state=spoofed-state" - ) + mock_redirect = RedirectResponse(url="http://idp.internal/authorize?state=spoofed-state") mock_generic_sso = MagicMock() mock_generic_sso.__enter__ = MagicMock(return_value=mock_generic_sso) mock_generic_sso.__exit__ = MagicMock(return_value=None) @@ -8142,9 +7954,7 @@ class TestPKCEStateCookieBinding: ) cookie_headers = response.headers.getlist("set-cookie") - cookie_str = next( - (c for c in cookie_headers if "litellm_oauth_state=" in c), None - ) + cookie_str = next((c for c in cookie_headers if "litellm_oauth_state=" in c), None) assert cookie_str is not None assert "Secure" not in cookie_str @@ -8551,9 +8361,15 @@ async def _render_debug_page(provider_env, id_jag_registered, force_inert=False) AsyncMock(return_value=id_jag_registered), ), patch("litellm.proxy.proxy_server.general_settings", {}), # test-quality-ok: debug endpoint reads proxy globals - patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), # test-quality-ok: debug endpoint reads proxy DB - patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()), # test-quality-ok: debug endpoint reads proxy globals - patch("litellm.proxy.proxy_server.jwt_handler", MagicMock(spec=JWTHandler)), # test-quality-ok: debug endpoint reads proxy globals + patch( + "litellm.proxy.proxy_server.prisma_client", MagicMock() + ), # test-quality-ok: debug endpoint reads proxy DB + patch( + "litellm.proxy.proxy_server.user_api_key_cache", MagicMock() + ), # test-quality-ok: debug endpoint reads proxy globals + patch( + "litellm.proxy.proxy_server.jwt_handler", MagicMock(spec=JWTHandler) + ), # test-quality-ok: debug endpoint reads proxy globals ] if force_inert: stack.append( @@ -8591,9 +8407,7 @@ async def test_debug_page_logs_the_capture_gap_but_never_renders_it(caplog): async def test_debug_page_is_byte_identical_when_the_provider_captures(): """A deployment with no gap must get the page it got before this change, to the byte. The comparison is against the endpoint with the diagnostic forced inert, not against a guess.""" - with_feature = await _render_debug_page( - {"GENERIC_CLIENT_ID": _GENERIC_DEBUG_CLIENT_ID}, id_jag_registered=True - ) + with_feature = await _render_debug_page({"GENERIC_CLIENT_ID": _GENERIC_DEBUG_CLIENT_ID}, id_jag_registered=True) pre_change = await _render_debug_page( {"GENERIC_CLIENT_ID": _GENERIC_DEBUG_CLIENT_ID}, id_jag_registered=True, @@ -8608,9 +8422,7 @@ async def test_debug_page_is_byte_identical_when_the_provider_captures(): async def test_debug_page_is_byte_identical_when_no_id_jag_server_is_registered(): """Most deployments run Google SSO and no id_jag server at all; their debug page must not grow an ID-JAG section about a feature they do not use.""" - with_feature = await _render_debug_page( - {"GOOGLE_CLIENT_ID": _GOOGLE_DEBUG_CLIENT_ID}, id_jag_registered=False - ) + with_feature = await _render_debug_page({"GOOGLE_CLIENT_ID": _GOOGLE_DEBUG_CLIENT_ID}, id_jag_registered=False) pre_change = await _render_debug_page( {"GOOGLE_CLIENT_ID": _GOOGLE_DEBUG_CLIENT_ID}, id_jag_registered=False, @@ -8774,16 +8586,19 @@ async def test_saml_callback_enforces_free_sso_user_limit_after_validation(): request_double = SimpleNamespace(cookies={}, headers={}, stream=_stream) - with patch.dict(os.environ, {"DISABLE_ADMIN_UI": "false"}), patch( - "litellm.proxy.proxy_server.premium_user", False - ), patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), patch( - "litellm.proxy.proxy_server.master_key", MASTER_KEY - ), patch( - "litellm.proxy.management_endpoints.sso.saml_sso.SAMLAuthHandler.handle_acs", - new=_fake_handle_acs, - ), patch( - "litellm.repositories.user_repository.UserRepository.count_billable_users", - new=AsyncMock(side_effect=_fake_count_billable_users), + with ( + patch.dict(os.environ, {"DISABLE_ADMIN_UI": "false"}), + patch("litellm.proxy.proxy_server.premium_user", False), + patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), + patch("litellm.proxy.proxy_server.master_key", MASTER_KEY), + patch( + "litellm.proxy.management_endpoints.sso.saml_sso.SAMLAuthHandler.handle_acs", + new=_fake_handle_acs, + ), + patch( + "litellm.repositories.user_repository.UserRepository.count_billable_users", + new=AsyncMock(side_effect=_fake_count_billable_users), + ), ): with pytest.raises(ProxyException) as exc: await saml_callback(request_double) @@ -8999,9 +8814,7 @@ async def test_verify_and_process_arm_captures_sso_assertion(): mock_jwt_handler.get_team_ids_from_jwt.return_value = [] mock_sso_instance = MagicMock() - mock_sso_instance.verify_and_process = AsyncMock( - return_value={"sub": "u1", "email": "u@example.com"} - ) + mock_sso_instance.verify_and_process = AsyncMock(return_value={"sub": "u1", "email": "u@example.com"}) mock_sso_instance.access_token = None mock_sso_instance.id_token = id_token mock_sso_instance.refresh_token = "rt_from_idp" @@ -9017,9 +8830,7 @@ async def test_verify_and_process_arm_captures_sso_assertion(): }, ): with patch("fastapi_sso.sso.base.DiscoveryDocument"): - with patch( - "fastapi_sso.sso.generic.create_provider", return_value=mock_sso_class - ): + with patch("fastapi_sso.sso.generic.create_provider", return_value=mock_sso_class): _, _, _, sso_assertion = await get_generic_sso_response( request=mock_request, jwt_handler=mock_jwt_handler, @@ -9062,9 +8873,7 @@ async def test_redirect_from_openid_persists_assertion_under_canonical_user_id() patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()), patch( "litellm.proxy.proxy_server.generate_key_helper_fn", - AsyncMock( - return_value={"token": "sk-ui-key", "user_id": "canonical-user-id"} - ), + AsyncMock(return_value={"token": "sk-ui-key", "user_id": "canonical-user-id"}), ), patch( "litellm.proxy.management_endpoints.ui_sso.get_user_info_from_db", @@ -9099,9 +8908,7 @@ async def test_redirect_from_openid_persists_assertion_under_canonical_user_id() sso_assertion=assertion, ) - retain_mock.assert_awaited_once_with( - user_id="canonical-user-id", assertion=assertion - ) + retain_mock.assert_awaited_once_with(user_id="canonical-user-id", assertion=assertion) assert response is not None @@ -9298,9 +9105,13 @@ async def test_browser_funnel_reports_an_uncaptured_assertion(monkeypatch, caplo patch("litellm.proxy.proxy_server.general_settings", {}), # test-quality-ok: endpoint reads proxy globals patch("litellm.proxy.proxy_server.premium_user", False), # test-quality-ok: endpoint reads proxy globals patch("litellm.proxy.proxy_server.user_custom_sso", None), # test-quality-ok: endpoint reads proxy globals - patch("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock()), # test-quality-ok: endpoint reads proxy globals + patch( + "litellm.proxy.proxy_server.proxy_logging_obj", MagicMock() + ), # test-quality-ok: endpoint reads proxy globals patch("litellm.proxy.proxy_server.redis_usage_cache", None), # test-quality-ok: endpoint reads proxy globals - patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()), # test-quality-ok: endpoint reads proxy globals + patch( + "litellm.proxy.proxy_server.user_api_key_cache", MagicMock() + ), # test-quality-ok: endpoint reads proxy globals patch( # test-quality-ok: endpoint test stubs key generation at its module boundary "litellm.proxy.proxy_server.generate_key_helper_fn", AsyncMock(return_value={"token": "sk-ui-key", "user_id": "canonical-user-id"}), @@ -9461,9 +9272,7 @@ async def test_cli_completion_drops_teams_whose_rows_no_longer_exist(): _complete_cli_sso_callback_session, ) - live_detail = CliSsoTeamDetail( - team_id="team-live", team_alias="Live", team_models=("gpt-4.1",) - ) + live_detail = CliSsoTeamDetail(team_id="team-live", team_alias="Live", team_models=("gpt-4.1",)) flow = {} with ( patch( @@ -9775,8 +9584,16 @@ async def test_cli_sign_in_enrolls_only_verified_subjects_before_completing( if trusted: table.upsert.assert_awaited_once_with( where={"issuer_tenant_id_oid": {"issuer": "issuer", "tenant_id": "tenant", "oid": "subject"}}, - data={"create": {"issuer": "issuer", "tenant_id": "tenant", "oid": "subject", - "user_id": "cli-user-id", "verified_via": "sso_interactive"}, "update": {}}, + data={ + "create": { + "issuer": "issuer", + "tenant_id": "tenant", + "oid": "subject", + "user_id": "cli-user-id", + "verified_via": "sso_interactive", + }, + "update": {}, + }, ) else: table.upsert.assert_not_awaited() @@ -9795,7 +9612,9 @@ async def _sso_key_generate_on_ui_disabled_node(*, source, key, google_sso_confi discovery document.""" from litellm.proxy.management_endpoints.ui_sso import _get_cli_sso_flow_cache_key, google_login - env_without_sso_providers = {name: value for name, value in os.environ.items() if name not in _SSO_PROVIDER_ENV_VARS} + env_without_sso_providers = { + name: value for name, value in os.environ.items() if name not in _SSO_PROVIDER_ENV_VARS + } env = { **env_without_sso_providers, "DISABLE_ADMIN_UI": "true", From 1a87606395ad89e07d52037d5fbe88e7a7db669d Mon Sep 17 00:00:00 2001 From: Yuneng Jiang <yuneng@berri.ai> Date: Sat, 10 Oct 2026 16:35:51 -0700 Subject: [PATCH 3/5] test: keep SSO regression changes focused --- .../proxy/management_endpoints/test_ui_sso.py | 830 ++++++++++++------ 1 file changed, 585 insertions(+), 245 deletions(-) diff --git a/tests/unit/proxy/management_endpoints/test_ui_sso.py b/tests/unit/proxy/management_endpoints/test_ui_sso.py index 26367fa8733..1b69b511543 100644 --- a/tests/unit/proxy/management_endpoints/test_ui_sso.py +++ b/tests/unit/proxy/management_endpoints/test_ui_sso.py @@ -135,7 +135,9 @@ def test_microsoft_sso_handler_with_empty_response(): # Test with None response # Act - result = MicrosoftSSOHandler.openid_from_response(response=None, team_ids=[], user_role=None) + result = MicrosoftSSOHandler.openid_from_response( + response=None, team_ids=[], user_role=None + ) # Assert assert isinstance(result, CustomOpenID) @@ -171,8 +173,12 @@ def test_microsoft_sso_handler_openid_from_response_with_custom_attributes(): "custom_display_name", ), patch("litellm.constants.MICROSOFT_USER_ID_ATTRIBUTE", "custom_id_field"), - patch("litellm.constants.MICROSOFT_USER_FIRST_NAME_ATTRIBUTE", "custom_first_name"), - patch("litellm.constants.MICROSOFT_USER_LAST_NAME_ATTRIBUTE", "custom_last_name"), + patch( + "litellm.constants.MICROSOFT_USER_FIRST_NAME_ATTRIBUTE", "custom_first_name" + ), + patch( + "litellm.constants.MICROSOFT_USER_LAST_NAME_ATTRIBUTE", "custom_last_name" + ), patch( "litellm.proxy.management_endpoints.ui_sso.MICROSOFT_USER_EMAIL_ATTRIBUTE", "custom_email_field", @@ -212,7 +218,9 @@ def test_microsoft_sso_handler_openid_from_response_with_custom_attributes(): @pytest.fixture def stubbed_graph_api(httpx_transport): with respx.mock: - respx.get(url__regex=r".*graph\.microsoft\.com.*").mock(return_value=httpx.Response(200, json={"value": []})) + respx.get(url__regex=r".*graph\.microsoft\.com.*").mock( + return_value=httpx.Response(200, json={"value": []}) + ) yield @@ -308,7 +316,9 @@ def test_get_google_callback_response(): with patch.dict(os.environ, {"GOOGLE_CLIENT_SECRET": "mock_secret"}): mock_verify = AsyncMock(return_value=mock_response) - with patch("fastapi_sso.sso.google.GoogleSSO.verify_and_process", new=mock_verify): + with patch( + "fastapi_sso.sso.google.GoogleSSO.verify_and_process", new=mock_verify + ): # Act result = asyncio.run( GoogleSSOHandler.get_google_callback_response( @@ -351,12 +361,16 @@ async def test_get_user_groups_from_graph_api(): mock.json.return_value = mock_response return mock - with patch("litellm.proxy.management_endpoints.ui_sso.get_async_httpx_client") as mock_client: + with patch( + "litellm.proxy.management_endpoints.ui_sso.get_async_httpx_client" + ) as mock_client: mock_client.return_value = MagicMock() mock_client.return_value.get = mock_get # Act - result = await MicrosoftSSOHandler.get_user_groups_from_graph_api(access_token="mock_token") + result = await MicrosoftSSOHandler.get_user_groups_from_graph_api( + access_token="mock_token" + ) # Assert assert isinstance(result, list) @@ -378,12 +392,16 @@ async def test_get_user_groups_empty_response(): mock.json.return_value = mock_response return mock - with patch("litellm.proxy.management_endpoints.ui_sso.get_async_httpx_client") as mock_client: + with patch( + "litellm.proxy.management_endpoints.ui_sso.get_async_httpx_client" + ) as mock_client: mock_client.return_value = MagicMock() mock_client.return_value.get = mock_get # Act - result = await MicrosoftSSOHandler.get_user_groups_from_graph_api(access_token="mock_token") + result = await MicrosoftSSOHandler.get_user_groups_from_graph_api( + access_token="mock_token" + ) # Assert assert isinstance(result, list) @@ -396,12 +414,16 @@ async def test_get_user_groups_error_handling(): async def mock_get(*args, **kwargs): raise Exception("API Error") - with patch("litellm.proxy.management_endpoints.ui_sso.get_async_httpx_client") as mock_client: + with patch( + "litellm.proxy.management_endpoints.ui_sso.get_async_httpx_client" + ) as mock_client: mock_client.return_value = MagicMock() mock_client.return_value.get = mock_get # Act - result = await MicrosoftSSOHandler.get_user_groups_from_graph_api(access_token="mock_token") + result = await MicrosoftSSOHandler.get_user_groups_from_graph_api( + access_token="mock_token" + ) # Assert assert isinstance(result, list) @@ -420,7 +442,9 @@ async def test_get_user_groups_uses_default_graph_endpoint(monkeypatch): mock.json.return_value = {"value": []} return mock - with patch("litellm.proxy.management_endpoints.ui_sso.get_async_httpx_client") as mock_client: + with patch( + "litellm.proxy.management_endpoints.ui_sso.get_async_httpx_client" + ) as mock_client: mock_client.return_value = MagicMock() mock_client.return_value.get = mock_get @@ -441,7 +465,9 @@ async def test_get_user_groups_uses_configured_graph_endpoint(monkeypatch): mock.json.return_value = {"value": []} return mock - with patch("litellm.proxy.management_endpoints.ui_sso.get_async_httpx_client") as mock_client: + with patch( + "litellm.proxy.management_endpoints.ui_sso.get_async_httpx_client" + ) as mock_client: mock_client.return_value = MagicMock() mock_client.return_value.get = mock_get @@ -471,7 +497,9 @@ async def test_get_group_ids_from_service_principal_uses_configured_graph_endpoi access_token="mock_token", ) - assert requested_urls == ["https://graph.microsoft.us/v1.0/servicePrincipals/sp-123/appRoleAssignedTo"] + assert requested_urls == [ + "https://graph.microsoft.us/v1.0/servicePrincipals/sp-123/appRoleAssignedTo" + ] @pytest.mark.asyncio @@ -569,7 +597,9 @@ def test_get_group_ids_from_graph_api_response(): "team_params", [ # Test case 1: Using DefaultTeamSSOParams - DefaultTeamSSOParams(max_budget=10, budget_duration="1d", models=["special-gpt-5"]), + DefaultTeamSSOParams( + max_budget=10, budget_duration="1d", models=["special-gpt-5"] + ), # Test case 2: Using Dict {"max_budget": 10, "budget_duration": "1d", "models": ["special-gpt-5"]}, ], @@ -608,7 +638,9 @@ async def test_default_team_params(team_params): # Assert # Verify team was created with correct parameters mock_prisma.db.litellm_teamtable.create.assert_called_once() - create_call_args = mock_prisma.db.litellm_teamtable.create.call_args.kwargs["data"] + create_call_args = mock_prisma.db.litellm_teamtable.create.call_args.kwargs[ + "data" + ] assert create_call_args["team_id"] == team_id assert create_call_args["team_alias"] == "Test Team" assert create_call_args["max_budget"] == 10 @@ -646,13 +678,10 @@ async def test_default_team_params_organization_id_reaches_sso_created_team(team updated_by="admin", ) - with ( - patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), - patch( - "litellm.proxy.management_endpoints.team_endpoints.get_org_object", - AsyncMock(return_value=mock_org), - ) as mock_get_org, - ): + with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), patch( + "litellm.proxy.management_endpoints.team_endpoints.get_org_object", + AsyncMock(return_value=mock_org), + ) as mock_get_org: team_id = str(uuid.uuid4()) await MicrosoftSSOHandler.create_litellm_teams_from_service_principal_team_ids( service_principal_teams=[ @@ -704,7 +733,9 @@ async def test_create_team_without_default_params(): # Assert mock_prisma.db.litellm_teamtable.create.assert_called_once() - create_call_args = mock_prisma.db.litellm_teamtable.create.call_args.kwargs["data"] + create_call_args = mock_prisma.db.litellm_teamtable.create.call_args.kwargs[ + "data" + ] assert create_call_args["team_id"] == team_id assert create_call_args["team_alias"] == "Test Team" # Should not have any of the optional fields @@ -1590,7 +1621,9 @@ async def test_get_user_info_from_db_user_exists(): "user_email": user_email, "user_defined_values": user_defined_values, } - with patch("litellm.proxy.management_endpoints.ui_sso.get_user_object") as mock_get_user_object: + with patch( + "litellm.proxy.management_endpoints.ui_sso.get_user_object" + ) as mock_get_user_object: await get_user_info_from_db(**args) mock_get_user_object.assert_called_once() assert mock_get_user_object.call_args.kwargs["user_id"] == "krrishd" @@ -1630,7 +1663,9 @@ async def test_get_user_info_from_db_user_exists_alternate_user_id(): "user_defined_values": user_defined_values, "alternate_user_id": "krrishd-email1234", } - with patch("litellm.proxy.management_endpoints.ui_sso.get_user_object") as mock_get_user_object: + with patch( + "litellm.proxy.management_endpoints.ui_sso.get_user_object" + ) as mock_get_user_object: await get_user_info_from_db(**args) mock_get_user_object.assert_called_once() assert mock_get_user_object.call_args.kwargs["user_id"] == "krrishd-email1234" @@ -1932,13 +1967,17 @@ async def test_get_generic_sso_response_with_additional_headers(): # Mock the SSO provider and its methods mock_sso_instance = MagicMock() mock_sso_instance.verify_and_process = AsyncMock(return_value=mock_sso_response) - mock_sso_instance.access_token = None # Avoid triggering JWT decode in process_sso_jwt_access_token + mock_sso_instance.access_token = ( + None # Avoid triggering JWT decode in process_sso_jwt_access_token + ) mock_sso_class = MagicMock(return_value=mock_sso_instance) with patch.dict(os.environ, test_env_vars): with patch("fastapi_sso.sso.base.DiscoveryDocument"): - with patch("fastapi_sso.sso.generic.create_provider", return_value=mock_sso_class): + with patch( + "fastapi_sso.sso.generic.create_provider", return_value=mock_sso_class + ): # Act result, received_response, _, _ = await get_generic_sso_response( request=mock_request, @@ -1992,13 +2031,17 @@ async def test_get_generic_sso_response_with_empty_headers(): # Mock the SSO provider and its methods mock_sso_instance = MagicMock() mock_sso_instance.verify_and_process = AsyncMock(return_value=mock_sso_response) - mock_sso_instance.access_token = None # Avoid triggering JWT decode in process_sso_jwt_access_token + mock_sso_instance.access_token = ( + None # Avoid triggering JWT decode in process_sso_jwt_access_token + ) mock_sso_class = MagicMock(return_value=mock_sso_instance) with patch.dict(os.environ, test_env_vars): with patch("fastapi_sso.sso.base.DiscoveryDocument"): - with patch("fastapi_sso.sso.generic.create_provider", return_value=mock_sso_class): + with patch( + "fastapi_sso.sso.generic.create_provider", return_value=mock_sso_class + ): # Act result, received_response, _, _ = await get_generic_sso_response( request=mock_request, @@ -2379,7 +2422,9 @@ class TestAuthCallbackRouting: for state in non_cli_states: # This mimics the routing logic in auth_callback - should_route_to_cli = state and state.startswith(f"{LITELLM_CLI_SESSION_TOKEN_PREFIX}:") + should_route_to_cli = state and state.startswith( + f"{LITELLM_CLI_SESSION_TOKEN_PREFIX}:" + ) assert not should_route_to_cli, f"State '{state}' should not route to CLI" @@ -2414,7 +2459,9 @@ class TestGoogleLoginCLIIntegration: for source, key in test_cases: cli_state = SSOAuthenticationHandler._get_cli_state(source=source, key=key) - assert cli_state is None, f"CLI state should not be generated for source='{source}', key='{key}'" + assert ( + cli_state is None + ), f"CLI state should not be generated for source='{source}', key='{key}'" class TestSSOHandlerIntegration: @@ -2425,13 +2472,24 @@ class TestSSOHandlerIntegration: from litellm.proxy.management_endpoints.ui_sso import SSOAuthenticationHandler # Test that SSO handler is used when client IDs are provided - assert SSOAuthenticationHandler.should_use_sso_handler(google_client_id="test") is True - assert SSOAuthenticationHandler.should_use_sso_handler(microsoft_client_id="test") is True - assert SSOAuthenticationHandler.should_use_sso_handler(generic_client_id="test") is True + assert ( + SSOAuthenticationHandler.should_use_sso_handler(google_client_id="test") + is True + ) + assert ( + SSOAuthenticationHandler.should_use_sso_handler(microsoft_client_id="test") + is True + ) + assert ( + SSOAuthenticationHandler.should_use_sso_handler(generic_client_id="test") + is True + ) # Test that SSO handler is not used when no client IDs are provided assert SSOAuthenticationHandler.should_use_sso_handler() is False - assert SSOAuthenticationHandler.should_use_sso_handler(None, None, None) is False + assert ( + SSOAuthenticationHandler.should_use_sso_handler(None, None, None) is False + ) @patch.dict(os.environ, {}, clear=False) def test_get_redirect_url_for_sso(self): @@ -2502,7 +2560,9 @@ class TestSSOStateHandling: """Test generating CLI state with valid parameters""" from litellm.proxy.management_endpoints.ui_sso import SSOAuthenticationHandler - state = SSOAuthenticationHandler._get_cli_state(source="litellm-cli", key="cli-test1234567890") + state = SSOAuthenticationHandler._get_cli_state( + source="litellm-cli", key="cli-test1234567890" + ) assert state is not None assert state.startswith("litellm-session-token:") @@ -2512,7 +2572,9 @@ class TestSSOStateHandling: """Test generating CLI state with invalid source""" from litellm.proxy.management_endpoints.ui_sso import SSOAuthenticationHandler - state = SSOAuthenticationHandler._get_cli_state(source="invalid_source", key="cli-test1234567890") + state = SSOAuthenticationHandler._get_cli_state( + source="invalid_source", key="cli-test1234567890" + ) assert state is None @@ -2528,7 +2590,9 @@ class TestSSOStateHandling: """Test generating CLI state without source""" from litellm.proxy.management_endpoints.ui_sso import SSOAuthenticationHandler - state = SSOAuthenticationHandler._get_cli_state(source=None, key="cli-test1234567890") + state = SSOAuthenticationHandler._get_cli_state( + source=None, key="cli-test1234567890" + ) assert state is None @@ -2697,7 +2761,9 @@ class TestCustomUISSO: # Test that the ValueError is raised with the correct message import pytest - with pytest.raises(ValueError, match="Enterprise features are not available"): + with pytest.raises( + ValueError, match="Enterprise features are not available" + ): asyncio.run(mock_google_login()) @pytest.mark.asyncio @@ -2731,7 +2797,9 @@ class TestCustomUISSO: picture=None, provider="custom", ) - mock_custom_handler.handle_custom_ui_sso_sign_in = AsyncMock(return_value=expected_openid) + mock_custom_handler.handle_custom_ui_sso_sign_in = AsyncMock( + return_value=expected_openid + ) # Mock the redirect response method mock_redirect_response = MagicMock() @@ -2752,11 +2820,15 @@ class TestCustomUISSO: return_value=mock_redirect_response, ) as mock_get_redirect: # Act - result = await EnterpriseCustomSSOHandler.handle_custom_ui_sso_sign_in(request=mock_request) + result = await EnterpriseCustomSSOHandler.handle_custom_ui_sso_sign_in( + request=mock_request + ) # Assert # Verify the custom handler was called with the request - mock_custom_handler.handle_custom_ui_sso_sign_in.assert_called_once_with(request=mock_request) + mock_custom_handler.handle_custom_ui_sso_sign_in.assert_called_once_with( + request=mock_request + ) # Verify the redirect response was generated with correct OpenID mock_get_redirect.assert_called_once_with( @@ -2801,7 +2873,9 @@ class TestCustomUISSO: {"trusted_proxy_ranges": ["10.0.0.0/24"]}, ): with pytest.raises(ValueError, match="not trusted"): - await EnterpriseCustomSSOHandler.handle_custom_ui_sso_sign_in(request=mock_request) + await EnterpriseCustomSSOHandler.handle_custom_ui_sso_sign_in( + request=mock_request + ) mock_custom_handler.handle_custom_ui_sso_sign_in.assert_not_called() @@ -2833,7 +2907,9 @@ class TestCustomUISSO: request_headers_dict = dict(request.headers) return OpenID( id=request_headers_dict.get("x-litellm-user-id", "default_user"), - email=request_headers_dict.get("x-litellm-user-email", "default@test.com"), + email=request_headers_dict.get( + "x-litellm-user-email", "default@test.com" + ), first_name="Custom", last_name="Handler", display_name="Custom Handler Test", @@ -2873,7 +2949,9 @@ class TestCustomUISSO: return_value=mock_redirect_response, ) as mock_get_redirect: # Act - result = await EnterpriseCustomSSOHandler.handle_custom_ui_sso_sign_in(request=mock_request) + result = await EnterpriseCustomSSOHandler.handle_custom_ui_sso_sign_in( + request=mock_request + ) # Assert that our custom handler was executed assert test_handler_instance.method_called is True @@ -3045,7 +3123,9 @@ class TestCLIKeyRegenerationFlow: redis_cache.set_cache.side_effect = lambda key, value, ttl: redis_store.__setitem__( key, str(value).encode("utf-8") ) - redis_cache.get_cache.side_effect = lambda key: RedisCache._get_cache_logic(MagicMock(), redis_store.get(key)) + redis_cache.get_cache.side_effect = lambda key: RedisCache._get_cache_logic( + MagicMock(), redis_store.get(key) + ) cache = MagicMock() cache.redis_cache = redis_cache @@ -3085,8 +3165,12 @@ class TestCLIKeyRegenerationFlow: assert mock_cache.increment_cache.call_args.kwargs["ttl"] == 60 mock_cache.set_cache.assert_called_once() flow_data = mock_cache.set_cache.call_args.kwargs["value"] - assert flow_data["poll_secret_hash"] == _hash_cli_sso_secret(result["poll_secret"]) - assert flow_data["user_code_hash"] == _hash_cli_sso_secret(_normalize_cli_sso_user_code(result["user_code"])) + assert flow_data["poll_secret_hash"] == _hash_cli_sso_secret( + result["poll_secret"] + ) + assert flow_data["user_code_hash"] == _hash_cli_sso_secret( + _normalize_cli_sso_user_code(result["user_code"]) + ) assert flow_data["poll_secret_hash"] != result["poll_secret"] assert flow_data["user_code_hash"] != result["user_code"] @@ -3194,7 +3278,9 @@ class TestCLIKeyRegenerationFlow: mock_cache = MagicMock(redis_cache=None) mock_cache.get_cache.return_value = {"poll_secret_hash": "h"} env_without_sso_providers = { - name: value for name, value in os.environ.items() if name not in _SSO_PROVIDER_ENV_VARS + name: value + for name, value in os.environ.items() + if name not in _SSO_PROVIDER_ENV_VARS } async def drive(enabled: bool): @@ -3245,7 +3331,9 @@ class TestCLIKeyRegenerationFlow: ) from litellm.proxy.management_endpoints.ui_sso import SSOAuthenticationHandler - manual_state = SSOAuthenticationHandler._get_cli_state(source=LITELLM_CLI_SOURCE_IDENTIFIER, key="cli-abc123") + manual_state = SSOAuthenticationHandler._get_cli_state( + source=LITELLM_CLI_SOURCE_IDENTIFIER, key="cli-abc123" + ) prefill_state = SSOAuthenticationHandler._get_cli_state( source=LITELLM_CLI_SOURCE_IDENTIFIER, key="cli-abc123", @@ -3253,9 +3341,14 @@ class TestCLIKeyRegenerationFlow: ) assert manual_state == f"{LITELLM_CLI_SESSION_TOKEN_PREFIX}:cli-abc123" - assert prefill_state == f"{LITELLM_CLI_SESSION_TOKEN_PREFIX}:cli-abc123:WXYZ-2345" assert ( - SSOAuthenticationHandler._get_cli_state(source="not-cli", key="cli-abc123", user_code="WXYZ-2345") is None + prefill_state == f"{LITELLM_CLI_SESSION_TOKEN_PREFIX}:cli-abc123:WXYZ-2345" + ) + assert ( + SSOAuthenticationHandler._get_cli_state( + source="not-cli", key="cli-abc123", user_code="WXYZ-2345" + ) + is None ) def test_get_cli_state_drops_malformed_user_code(self): @@ -3339,7 +3432,9 @@ class TestCLIKeyRegenerationFlow: browser_complete_token="browser-token", ) - input_line = next(line for line in html.splitlines() if 'name="user_code"' in line) + input_line = next( + line for line in html.splitlines() if 'name="user_code"' in line + ) assert "value=" not in input_line assert "shown in your terminal" in html @@ -3402,11 +3497,15 @@ class TestCLIKeyRegenerationFlow: ) mock_request = MagicMock(spec=Request) - mock_request.body = AsyncMock(return_value=b"user_code=ABCD-EFGH&browser_complete_token=browser-token") + mock_request.body = AsyncMock( + return_value=b"user_code=ABCD-EFGH&browser_complete_token=browser-token" + ) mock_cache = MagicMock(redis_cache=None) mock_cache.get_cache.return_value = { "poll_secret_hash": _hash_cli_sso_secret("poll-secret"), - "user_code_hash": _hash_cli_sso_secret(_normalize_cli_sso_user_code("ABCD-EFGH")), + "user_code_hash": _hash_cli_sso_secret( + _normalize_cli_sso_user_code("ABCD-EFGH") + ), "browser_complete_token_hash": _hash_cli_sso_secret("browser-token"), "sso_complete": True, "user_code_verified": False, @@ -3421,7 +3520,9 @@ class TestCLIKeyRegenerationFlow: return_value="<html>Success</html>", ), ): - result = await cli_sso_complete(request=mock_request, login_id="cli-session-4567890") + result = await cli_sso_complete( + request=mock_request, login_id="cli-session-4567890" + ) assert result.status_code == 200 flow_data = mock_cache.set_cache.call_args.kwargs["value"] @@ -3441,7 +3542,9 @@ class TestCLIKeyRegenerationFlow: mock_cache = MagicMock(redis_cache=None) mock_cache.get_cache.return_value = { "poll_secret_hash": _hash_cli_sso_secret("poll-secret"), - "user_code_hash": _hash_cli_sso_secret(_normalize_cli_sso_user_code("ABCD-EFGH")), + "user_code_hash": _hash_cli_sso_secret( + _normalize_cli_sso_user_code("ABCD-EFGH") + ), "browser_complete_token_hash": _hash_cli_sso_secret("browser-token"), "sso_complete": True, "user_code_verified": False, @@ -3453,7 +3556,9 @@ class TestCLIKeyRegenerationFlow: patch("litellm.proxy.proxy_server.cli_sso_session_cache", mock_cache), ): with pytest.raises(HTTPException) as exc_info: - await cli_sso_complete(request=mock_request, login_id="cli-session-4567890") + await cli_sso_complete( + request=mock_request, login_id="cli-session-4567890" + ) assert exc_info.value.status_code == 400 mock_cache.set_cache.assert_not_called() @@ -3468,11 +3573,15 @@ class TestCLIKeyRegenerationFlow: ) mock_request = MagicMock(spec=Request) - mock_request.body = AsyncMock(return_value=b"user_code=ABCD-EFGH&browser_complete_token=browser-token") + mock_request.body = AsyncMock( + return_value=b"user_code=ABCD-EFGH&browser_complete_token=browser-token" + ) mock_cache = MagicMock(redis_cache=None) mock_cache.get_cache.return_value = { "poll_secret_hash": _hash_cli_sso_secret("poll-secret"), - "user_code_hash": _hash_cli_sso_secret(_normalize_cli_sso_user_code("ABCD-EFGH")), + "user_code_hash": _hash_cli_sso_secret( + _normalize_cli_sso_user_code("ABCD-EFGH") + ), "sso_complete": False, "user_code_verified": False, "session_data": None, @@ -3483,7 +3592,9 @@ class TestCLIKeyRegenerationFlow: patch("litellm.proxy.proxy_server.cli_sso_session_cache", mock_cache), ): with pytest.raises(HTTPException) as exc_info: - await cli_sso_complete(request=mock_request, login_id="cli-session-4567890") + await cli_sso_complete( + request=mock_request, login_id="cli-session-4567890" + ) assert exc_info.value.status_code == 400 assert exc_info.value.detail == "CLI login is not ready" @@ -3589,7 +3700,10 @@ class TestCLIKeyRegenerationFlow: assert result.status_code == 200 # Verify response contains success message (response is HTML) assert result.body is not None - assert 'action="https://test.example.com/sso/cli/complete/cli-session-4567890"' in result.body.decode() + assert ( + 'action="https://test.example.com/sso/cli/complete/cli-session-4567890"' + in result.body.decode() + ) @pytest.mark.asyncio async def test_cli_poll_key_returns_teams_for_selection(self): @@ -3717,7 +3831,9 @@ class TestCLIKeyRegenerationFlow: mock_result = {"user_id": "test-user", "email": "test@example.com"} with ( - patch("litellm.proxy.management_endpoints.ui_sso.cli_sso_callback") as mock_cli_callback, + patch( + "litellm.proxy.management_endpoints.ui_sso.cli_sso_callback" + ) as mock_cli_callback, patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), patch("litellm.proxy.proxy_server.master_key", "test-master-key"), patch("litellm.proxy.proxy_server.general_settings", {}), @@ -3751,11 +3867,15 @@ class TestCLIKeyRegenerationFlow: mock_request = MagicMock(spec=Request) mock_request.query_params = {"code": "some-auth-code"} - cli_state = f"{LITELLM_CLI_SESSION_TOKEN_PREFIX}:cli-new-session-key-456:WXYZ-2345" + cli_state = ( + f"{LITELLM_CLI_SESSION_TOKEN_PREFIX}:cli-new-session-key-456:WXYZ-2345" + ) mock_result = {"user_id": "test-user", "email": "test@example.com"} with ( - patch("litellm.proxy.management_endpoints.ui_sso.cli_sso_callback") as mock_cli_callback, + patch( + "litellm.proxy.management_endpoints.ui_sso.cli_sso_callback" + ) as mock_cli_callback, patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), patch("litellm.proxy.proxy_server.master_key", "test-master-key"), patch("litellm.proxy.proxy_server.general_settings", {}), @@ -3788,7 +3908,9 @@ class TestCLIKeyRegenerationFlow: mock_request = MagicMock() mock_request.base_url = "https://test.litellm.ai/" - with patch("litellm.proxy.utils.get_custom_url", return_value="https://test.litellm.ai"): + with patch( + "litellm.proxy.utils.get_custom_url", return_value="https://test.litellm.ai" + ): # Test with existing_key - should NOT be in URL redirect_url = SSOAuthenticationHandler.get_redirect_url_for_sso( request=mock_request, @@ -3808,7 +3930,9 @@ class TestCLIKeyRegenerationFlow: mock_request = MagicMock() mock_request.base_url = "https://test.litellm.ai/" - with patch("litellm.proxy.utils.get_custom_url", return_value="https://test.litellm.ai"): + with patch( + "litellm.proxy.utils.get_custom_url", return_value="https://test.litellm.ai" + ): # Test without existing_key redirect_url = SSOAuthenticationHandler.get_redirect_url_for_sso( request=mock_request, sso_callback_route="sso/callback" @@ -3931,7 +4055,9 @@ class TestCLIKeyRegenerationFlow: prisma_client = MagicMock() prisma_client.db.litellm_teamtable.find_many = find_many - details = await fetch_cli_sso_team_details(prisma_client=prisma_client, teams=["team-a"]) + details = await fetch_cli_sso_team_details( + prisma_client=prisma_client, teams=["team-a"] + ) assert find_many.await_args.kwargs["include"] == {"litellm_model_table": True} assert [detail.model_dump() for detail in details] == [ @@ -3956,12 +4082,24 @@ class TestCLIKeyRegenerationFlow: ) failing_client = MagicMock() - failing_client.db.litellm_teamtable.find_many = AsyncMock(side_effect=Exception("connection reset")) - assert await fetch_cli_sso_team_details(prisma_client=failing_client, teams=["team-a"]) is None + failing_client.db.litellm_teamtable.find_many = AsyncMock( + side_effect=Exception("connection reset") + ) + assert ( + await fetch_cli_sso_team_details( + prisma_client=failing_client, teams=["team-a"] + ) + is None + ) empty_client = MagicMock() empty_client.db.litellm_teamtable.find_many = AsyncMock(return_value=[]) - assert await fetch_cli_sso_team_details(prisma_client=empty_client, teams=["team-a"]) == () + assert ( + await fetch_cli_sso_team_details( + prisma_client=empty_client, teams=["team-a"] + ) + == () + ) @pytest.mark.asyncio async def test_cli_poll_key_mints_jwt_with_selected_team_grants(self): @@ -4025,14 +4163,18 @@ class TestCLIKeyRegenerationFlow: "team_details", [ pytest.param(None, id="detail_fetch_failed"), - pytest.param([{"team_id": "team-other", "team_models": []}], id="selected_team_absent"), + pytest.param( + [{"team_id": "team-other", "team_models": []}], id="selected_team_absent" + ), pytest.param( [{"team_id": "team-a", "team_alias": "Team A"}], id="legacy_detail_without_grants", ), ], ) - async def test_cli_poll_key_refuses_to_mint_when_team_grants_are_unknown(self, team_details): + async def test_cli_poll_key_refuses_to_mint_when_team_grants_are_unknown( + self, team_details + ): """An unknown team grant must never be minted as an empty one. get_complete_model_list falls through to the whole proxy model list when both @@ -4168,7 +4310,9 @@ class TestCLIKeyRegenerationFlow: ), patch( "litellm.proxy.auth.auth_checks.get_team_object", - new=AsyncMock(side_effect=AssertionError("team lookup must be skipped")), + new=AsyncMock( + side_effect=AssertionError("team lookup must be skipped") + ), ), ): result = await cli_poll_key( @@ -4250,7 +4394,9 @@ class TestGetAppRolesFromIdToken: # Assert assert result == ["Admin", "User", "Developer"] - mock_jwt_decode.assert_called_once_with(mock_token, options={"verify_signature": False}) + mock_jwt_decode.assert_called_once_with( + mock_token, options={"verify_signature": False} + ) def test_app_roles_picked_when_both_exist(self): """Test that 'app_roles' takes precedence when both 'app_roles' and 'roles' exist""" @@ -4372,7 +4518,9 @@ class TestProcessSSOJWTAccessToken: "groups": ["team1", "team2", "team3"], } - def test_process_sso_jwt_access_token_with_existing_team_ids(self, mock_jwt_handler, sample_jwt_token): + def test_process_sso_jwt_access_token_with_existing_team_ids( + self, mock_jwt_handler, sample_jwt_token + ): """Test that existing team IDs are not overwritten""" from litellm.proxy.management_endpoints.ui_sso import ( process_sso_jwt_access_token, @@ -4428,14 +4576,20 @@ class TestProcessSSOJWTAccessToken: ) # Assert - mock_jwt_decode.assert_called_once_with(sample_jwt_token, options={"verify_signature": False}) - mock_jwt_handler.get_team_ids_from_jwt.assert_called_once_with(sample_jwt_payload) + mock_jwt_decode.assert_called_once_with( + sample_jwt_token, options={"verify_signature": False} + ) + mock_jwt_handler.get_team_ids_from_jwt.assert_called_once_with( + sample_jwt_payload + ) # Verify team_ids was added to the dict as a key assert "team_ids" in result assert result["team_ids"] == ["team1", "team2", "team3"] - def test_process_sso_jwt_access_token_with_dict_existing_team_ids(self, mock_jwt_handler, sample_jwt_token): + def test_process_sso_jwt_access_token_with_dict_existing_team_ids( + self, mock_jwt_handler, sample_jwt_token + ): """Test that existing team IDs in dictionary are not overwritten""" from litellm.proxy.management_endpoints.ui_sso import ( process_sso_jwt_access_token, @@ -4478,7 +4632,9 @@ class TestProcessSSOJWTAccessToken: # Test with None access token with patch("jwt.decode") as mock_jwt_decode: - process_sso_jwt_access_token(access_token_str=None, sso_jwt_handler=mock_jwt_handler, result=result) + process_sso_jwt_access_token( + access_token_str=None, sso_jwt_handler=mock_jwt_handler, result=result + ) # Assert nothing was processed mock_jwt_decode.assert_not_called() @@ -4487,14 +4643,18 @@ class TestProcessSSOJWTAccessToken: # Test with empty string access token with patch("jwt.decode") as mock_jwt_decode: - process_sso_jwt_access_token(access_token_str="", sso_jwt_handler=mock_jwt_handler, result=result) + process_sso_jwt_access_token( + access_token_str="", sso_jwt_handler=mock_jwt_handler, result=result + ) # Assert nothing was processed mock_jwt_decode.assert_not_called() mock_jwt_handler.get_team_ids_from_jwt.assert_not_called() assert result.team_ids == [] - def test_process_sso_jwt_access_token_no_result(self, mock_jwt_handler, sample_jwt_token): + def test_process_sso_jwt_access_token_no_result( + self, mock_jwt_handler, sample_jwt_token + ): """Test that nothing happens when result is None""" from litellm.proxy.management_endpoints.ui_sso import ( process_sso_jwt_access_token, @@ -4512,7 +4672,9 @@ class TestProcessSSOJWTAccessToken: mock_jwt_decode.assert_not_called() mock_jwt_handler.get_team_ids_from_jwt.assert_not_called() - def test_process_sso_jwt_access_token_non_decode_exception_propagates(self, mock_jwt_handler, sample_jwt_token): + def test_process_sso_jwt_access_token_non_decode_exception_propagates( + self, mock_jwt_handler, sample_jwt_token + ): """Test that non-DecodeError JWT exceptions still propagate up.""" import jwt as pyjwt @@ -4522,7 +4684,9 @@ class TestProcessSSOJWTAccessToken: result = CustomOpenID(id="test_user", email="test@example.com", team_ids=[]) - with patch("jwt.decode", side_effect=pyjwt.exceptions.InvalidKeyError("Invalid key")) as mock_jwt_decode: + with patch( + "jwt.decode", side_effect=pyjwt.exceptions.InvalidKeyError("Invalid key") + ) as mock_jwt_decode: with pytest.raises(pyjwt.exceptions.InvalidKeyError, match="Invalid key"): process_sso_jwt_access_token( access_token_str=sample_jwt_token, @@ -4556,7 +4720,9 @@ class TestProcessSSOJWTAccessToken: # Assert mock_jwt_decode.assert_called_once() - mock_jwt_handler.get_team_ids_from_jwt.assert_called_once_with(sample_jwt_payload) + mock_jwt_handler.get_team_ids_from_jwt.assert_called_once_with( + sample_jwt_payload + ) # Even empty team IDs should be set assert result.team_ids == [] @@ -4593,7 +4759,9 @@ class TestProcessSSOJWTAccessToken: assert result.team_ids == ["existing_team"] assert result.user_role is None - def test_process_sso_jwt_access_token_real_jwt_with_role_and_teams(self, mock_jwt_handler): + def test_process_sso_jwt_access_token_real_jwt_with_role_and_teams( + self, mock_jwt_handler + ): """Test that a real JWT containing role and team fields is correctly processed.""" import jwt as pyjwt @@ -4766,9 +4934,15 @@ class TestGenericResponseConvertorNestedAttributes: # Current behavior: returns None for nested paths # Expected behavior with current implementation (no nested path support): assert result.id == "nested-user-456" - assert result.email == "john.doe@example.com" # Can't access "attributes.email" with .get() - assert result.first_name == "John" # Can't access "attributes.given_name" with .get() - assert result.last_name == "Doe" # Can't access "attributes.family_name" with .get() + assert ( + result.email == "john.doe@example.com" + ) # Can't access "attributes.email" with .get() + assert ( + result.first_name == "John" + ) # Can't access "attributes.given_name" with .get() + assert ( + result.last_name == "Doe" + ) # Can't access "attributes.family_name" with .get() assert result.display_name == "user-sub-123" # Top-level attribute works @@ -4944,8 +5118,14 @@ class TestGetGenericSSORedirectParams: assert redirect_params["code_challenge_method"] == "S256" # Verify code_challenge is correctly derived from code_verifier - expected_challenge_bytes = hashlib.sha256(code_verifier.encode("utf-8")).digest() - expected_challenge = base64.urlsafe_b64encode(expected_challenge_bytes).decode("utf-8").rstrip("=") + expected_challenge_bytes = hashlib.sha256( + code_verifier.encode("utf-8") + ).digest() + expected_challenge = ( + base64.urlsafe_b64encode(expected_challenge_bytes) + .decode("utf-8") + .rstrip("=") + ) assert redirect_params["code_challenge"] == expected_challenge def test_state_with_pkce_disabled(self): @@ -5078,8 +5258,14 @@ class TestPKCEFunctionality: assert isinstance(code_verifier, str) # Verify code_challenge is correctly generated from code_verifier - expected_challenge_bytes = hashlib.sha256(code_verifier.encode("utf-8")).digest() - expected_challenge = base64.urlsafe_b64encode(expected_challenge_bytes).decode("utf-8").rstrip("=") + expected_challenge_bytes = hashlib.sha256( + code_verifier.encode("utf-8") + ).digest() + expected_challenge = ( + base64.urlsafe_b64encode(expected_challenge_bytes) + .decode("utf-8") + .rstrip("=") + ) assert code_challenge == expected_challenge # Verify both are base64url encoded (no padding) @@ -5101,7 +5287,9 @@ class TestPKCEFunctionality: # Mock cache with async methods — use dict format (primary path) mock_cache = MagicMock(redis_cache=None) test_code_verifier = "test_code_verifier_abc123xyz" - mock_cache.async_get_cache = AsyncMock(return_value={"code_verifier": test_code_verifier}) + mock_cache.async_get_cache = AsyncMock( + return_value={"code_verifier": test_code_verifier} + ) mock_cache.async_delete_cache = AsyncMock() with ( @@ -5110,8 +5298,10 @@ class TestPKCEFunctionality: patch.dict(os.environ, {"GENERIC_CLIENT_USE_PKCE": "true"}), ): # Act - token_params = await SSOAuthenticationHandler.prepare_token_exchange_parameters( - request=mock_request, generic_include_client_id=False + token_params = ( + await SSOAuthenticationHandler.prepare_token_exchange_parameters( + request=mock_request, generic_include_client_id=False + ) ) # Assert @@ -5122,7 +5312,9 @@ class TestPKCEFunctionality: # Verify cache was read but NOT deleted yet (deletion is deferred to after # successful token exchange to preserve the verifier for retries) - mock_cache.async_get_cache.assert_called_once_with(key=f"pkce_verifier:{test_state}") + mock_cache.async_get_cache.assert_called_once_with( + key=f"pkce_verifier:{test_state}" + ) mock_cache.async_delete_cache.assert_not_called() @pytest.mark.asyncio @@ -5135,7 +5327,9 @@ class TestPKCEFunctionality: # Mock SSO provider mock_sso = MagicMock() mock_redirect_response = MagicMock() - original_location = "https://auth.example.com/authorize?state=test456&client_id=abc" + original_location = ( + "https://auth.example.com/authorize?state=test456&client_id=abc" + ) mock_redirect_response.headers = {"location": original_location} mock_sso.get_login_redirect = AsyncMock(return_value=mock_redirect_response) mock_sso.__enter__ = MagicMock(return_value=mock_sso) @@ -5221,7 +5415,9 @@ class TestPKCEFunctionality: with patch.dict(os.environ, {"GENERIC_CLIENT_USE_PKCE": "true"}): with patch("litellm.proxy.proxy_server.redis_usage_cache", mock_redis): - with patch("litellm.proxy.proxy_server.user_api_key_cache", mock_in_memory): + with patch( + "litellm.proxy.proxy_server.user_api_key_cache", mock_in_memory + ): # Pod A: start login, store code_verifier in "Redis" await SSOAuthenticationHandler.get_generic_sso_redirect_response( generic_sso=mock_sso, @@ -5235,7 +5431,9 @@ class TestPKCEFunctionality: stored_value = mock_redis._store[stored_key] # Stored as JSON-serialized dict for Redis compatibility stored_dict = json.loads(stored_value) - assert isinstance(stored_dict, dict) and "code_verifier" in stored_dict + assert ( + isinstance(stored_dict, dict) and "code_verifier" in stored_dict + ) assert len(stored_dict["code_verifier"]) == 43 # Pod B: callback with same state, retrieve from "Redis" @@ -5290,7 +5488,9 @@ class TestPKCEFunctionality: with patch.dict(os.environ, {"GENERIC_CLIENT_USE_PKCE": "true"}): with patch("litellm.proxy.proxy_server.redis_usage_cache", None): - with patch("litellm.proxy.proxy_server.user_api_key_cache", mock_in_memory): + with patch( + "litellm.proxy.proxy_server.user_api_key_cache", mock_in_memory + ): # Pod A: start login, store code_verifier in in-memory cache await SSOAuthenticationHandler.get_generic_sso_redirect_response( generic_sso=mock_sso, @@ -5299,9 +5499,14 @@ class TestPKCEFunctionality: ) mock_in_memory.async_set_cache.assert_called_once() stored_key = mock_in_memory.async_set_cache.call_args.kwargs["key"] - stored_value = mock_in_memory.async_set_cache.call_args.kwargs["value"] + stored_value = mock_in_memory.async_set_cache.call_args.kwargs[ + "value" + ] assert stored_key == "pkce_verifier:fallback_state_xyz" - assert isinstance(stored_value, dict) and len(stored_value["code_verifier"]) == 43 + assert ( + isinstance(stored_value, dict) + and len(stored_value["code_verifier"]) == 43 + ) # Same pod: callback retrieves from in-memory cache mock_request = MagicMock(spec=Request) @@ -5310,10 +5515,14 @@ class TestPKCEFunctionality: request=mock_request, generic_include_client_id=False ) assert "code_verifier" in token_params - assert token_params["code_verifier"] == stored_value["code_verifier"] + assert ( + token_params["code_verifier"] == stored_value["code_verifier"] + ) # Cache key returned for deferred deletion after successful exchange assert token_params["_pkce_cache_key"] == stored_key - mock_in_memory.async_get_cache.assert_called_once_with(key=stored_key) + mock_in_memory.async_get_cache.assert_called_once_with( + key=stored_key + ) # Deletion is deferred — not called by prepare_token_exchange_parameters mock_in_memory.async_delete_cache.assert_not_called() @@ -5335,8 +5544,10 @@ class TestPKCEFunctionality: ): mock_request = MagicMock(spec=Request) mock_request.query_params = {} - token_params = await SSOAuthenticationHandler.prepare_token_exchange_parameters( - request=mock_request, generic_include_client_id=False + token_params = ( + await SSOAuthenticationHandler.prepare_token_exchange_parameters( + request=mock_request, generic_include_client_id=False + ) ) assert "code_verifier" not in token_params mock_redis.async_get_cache.assert_not_called() @@ -5372,10 +5583,12 @@ class TestPKCEFunctionality: # Verify redirect_uri is forwarded (required by strict OAuth providers) assert post_data.get("redirect_uri") == "https://proxy.example.com/callback" # Verify credentials are NOT double-sent in the POST body when using Basic Auth - assert "client_secret" not in post_data, "client_secret must not appear in POST body when using Basic Auth" - assert "client_id" not in post_data, ( - "client_id must not appear in POST body when using Basic Auth (include_client_id=False)" - ) + assert ( + "client_secret" not in post_data + ), "client_secret must not appear in POST body when using Basic Auth" + assert ( + "client_id" not in post_data + ), "client_id must not appear in POST body when using Basic Auth (include_client_id=False)" return mock_response # get_async_httpx_client returns a client directly (no context manager). @@ -5385,7 +5598,9 @@ class TestPKCEFunctionality: mock_userinfo_client = MagicMock() mock_userinfo_client.get = AsyncMock(return_value=mock_userinfo_response) - with patch("litellm.proxy.management_endpoints.ui_sso.get_async_httpx_client") as mock_get_client: + with patch( + "litellm.proxy.management_endpoints.ui_sso.get_async_httpx_client" + ) as mock_get_client: mock_get_client.side_effect = [mock_token_client, mock_userinfo_client] result = await SSOAuthenticationHandler._pkce_token_exchange( @@ -5404,7 +5619,9 @@ class TestPKCEFunctionality: assert result["email"] == "user@example.com" # id_token was explicit null in token_response — the merge loop must remove it # rather than leaving "id_token": None in the result. - assert "id_token" not in result, "null id_token from token endpoint must be absent in merged result" + assert ( + "id_token" not in result + ), "null id_token from token endpoint must be absent in merged result" # Verify userinfo GET used the correct Bearer token header get_call = mock_userinfo_client.get.call_args assert get_call is not None @@ -5424,12 +5641,18 @@ class TestPKCEFunctionality: async def fake_post(*args, **kwargs): headers = kwargs.get("headers", {}) auth_header = headers.get("Authorization", "") - assert not auth_header.startswith("Basic "), "Should NOT use Basic Auth when include_client_id=True" + assert not auth_header.startswith( + "Basic " + ), "Should NOT use Basic Auth when include_client_id=True" data = kwargs.get("data", {}) assert "client_id" in data assert "client_secret" in data - assert data.get("code_verifier") == "verifier_xyz", "code_verifier must be in POST body" - assert data.get("redirect_uri") == "https://proxy.example.com/callback", "redirect_uri must be forwarded" + assert ( + data.get("code_verifier") == "verifier_xyz" + ), "code_verifier must be in POST body" + assert ( + data.get("redirect_uri") == "https://proxy.example.com/callback" + ), "redirect_uri must be forwarded" mock = MagicMock() mock.status_code = 200 mock.json.return_value = token_resp @@ -5445,7 +5668,9 @@ class TestPKCEFunctionality: mock_userinfo_client = MagicMock() mock_userinfo_client.get = AsyncMock(return_value=mock_userinfo) - with patch("litellm.proxy.management_endpoints.ui_sso.get_async_httpx_client") as mock_get_client: + with patch( + "litellm.proxy.management_endpoints.ui_sso.get_async_httpx_client" + ) as mock_get_client: mock_get_client.side_effect = [mock_token_client, mock_userinfo_client] result = await SSOAuthenticationHandler._pkce_token_exchange( @@ -5477,7 +5702,9 @@ class TestPKCEFunctionality: "error_description": "Code already used", } - with patch("litellm.proxy.management_endpoints.ui_sso.get_async_httpx_client") as mock_get_client: + with patch( + "litellm.proxy.management_endpoints.ui_sso.get_async_httpx_client" + ) as mock_get_client: mock_client = MagicMock() mock_resp = MagicMock() mock_resp.status_code = 200 @@ -5509,10 +5736,16 @@ class TestPKCEFunctionality: payload = {"sub": "user_from_jwt", "email": "jwt@example.com"} # Build a minimal JWT (header.payload.signature — signature not verified) - encoded_payload = base64.urlsafe_b64encode(_json.dumps(payload).encode()).rstrip(b"=").decode() + encoded_payload = ( + base64.urlsafe_b64encode(_json.dumps(payload).encode()) + .rstrip(b"=") + .decode() + ) fake_id_token = f"eyJhbGciOiJSUzI1NiJ9.{encoded_payload}.fakesig" - with patch("litellm.proxy.management_endpoints.ui_sso.get_async_httpx_client") as mock_get_client: + with patch( + "litellm.proxy.management_endpoints.ui_sso.get_async_httpx_client" + ) as mock_get_client: mock_client = MagicMock() mock_fail = MagicMock() mock_fail.status_code = 503 @@ -5538,7 +5771,11 @@ class TestPKCEFunctionality: from litellm.proxy.management_endpoints.ui_sso import SSOAuthenticationHandler payload = {"sub": "id_token_user", "email": "id@example.com"} - encoded_payload = base64.urlsafe_b64encode(_json.dumps(payload).encode()).rstrip(b"=").decode() + encoded_payload = ( + base64.urlsafe_b64encode(_json.dumps(payload).encode()) + .rstrip(b"=") + .decode() + ) fake_id_token = f"eyJhbGciOiJSUzI1NiJ9.{encoded_payload}.fakesig" # No httpx call should happen when userinfo_endpoint is None @@ -5558,7 +5795,9 @@ class TestPKCEFunctionality: from litellm.proxy._types import ProxyException from litellm.proxy.management_endpoints.ui_sso import SSOAuthenticationHandler - with patch("litellm.proxy.management_endpoints.ui_sso.get_async_httpx_client") as mock_get_client: + with patch( + "litellm.proxy.management_endpoints.ui_sso.get_async_httpx_client" + ) as mock_get_client: mock_client = MagicMock() mock_fail = MagicMock() mock_fail.status_code = 503 @@ -5587,7 +5826,9 @@ class TestPKCEFunctionality: mock_resp.status_code = 200 mock_resp.json.return_value = None # HTTP 200 with null JSON body - with patch("litellm.proxy.management_endpoints.ui_sso.get_async_httpx_client") as mock_get_client: + with patch( + "litellm.proxy.management_endpoints.ui_sso.get_async_httpx_client" + ) as mock_get_client: mock_client = MagicMock() mock_client.get = AsyncMock(return_value=mock_resp) mock_get_client.return_value = mock_client @@ -5638,7 +5879,10 @@ class TestPKCEFunctionality: request=mock_request, generic_include_client_id=False ) - assert "verifier not found" in exc_info.value.message.lower() or "cache" in exc_info.value.message.lower() + assert ( + "verifier not found" in exc_info.value.message.lower() + or "cache" in exc_info.value.message.lower() + ) assert str(exc_info.value.code) == "401" @pytest.mark.asyncio @@ -5657,10 +5901,14 @@ class TestPKCEFunctionality: async def fake_post(*args, **kwargs): headers = kwargs.get("headers", {}) auth_header = headers.get("Authorization", "") - assert not auth_header.startswith("Basic "), "Public client must not use Basic Auth" + assert not auth_header.startswith( + "Basic " + ), "Public client must not use Basic Auth" data = kwargs.get("data", {}) assert data.get("client_id") == "public_client_id" - assert "client_secret" not in data, "No secret should be sent for public client" + assert ( + "client_secret" not in data + ), "No secret should be sent for public client" assert data.get("code_verifier") == "public_verifier" mock = MagicMock() mock.status_code = 200 @@ -5677,7 +5925,9 @@ class TestPKCEFunctionality: mock_userinfo_client = MagicMock() mock_userinfo_client.get = AsyncMock(return_value=mock_userinfo) - with patch("litellm.proxy.management_endpoints.ui_sso.get_async_httpx_client") as mock_get_client: + with patch( + "litellm.proxy.management_endpoints.ui_sso.get_async_httpx_client" + ) as mock_get_client: mock_get_client.side_effect = [mock_token_client, mock_userinfo_client] result = await SSOAuthenticationHandler._pkce_token_exchange( @@ -5705,16 +5955,22 @@ class TestPKCEFunctionality: from litellm.proxy.management_endpoints.ui_sso import SSOAuthenticationHandler failing_cache = MagicMock() - failing_cache.async_delete_cache = AsyncMock(side_effect=Exception("Redis down")) + failing_cache.async_delete_cache = AsyncMock( + side_effect=Exception("Redis down") + ) # Should NOT raise even though the underlying cache delete fails with ( patch("litellm.proxy.proxy_server.redis_usage_cache", None), patch("litellm.proxy.proxy_server.user_api_key_cache", failing_cache), ): - await SSOAuthenticationHandler._delete_pkce_verifier("pkce_verifier:test_state") + await SSOAuthenticationHandler._delete_pkce_verifier( + "pkce_verifier:test_state" + ) - failing_cache.async_delete_cache.assert_called_once_with(key="pkce_verifier:test_state") + failing_cache.async_delete_cache.assert_called_once_with( + key="pkce_verifier:test_state" + ) @pytest.mark.asyncio async def test_pkce_cache_miss_unexpected_format_raises_proxy_exception(self): @@ -5802,7 +6058,8 @@ class TestPKCEFunctionality: mock_cache.async_get_cache.assert_called_once() # Verify the warning was actually logged assert any( - "verifier not found" in r.message.lower() or "code_verifier" in r.message.lower() + "verifier not found" in r.message.lower() + or "code_verifier" in r.message.lower() for r in caplog.records if r.levelno >= logging.WARNING ), f"Expected a cache-miss warning. Records: {[r.message for r in caplog.records]}" @@ -5820,7 +6077,9 @@ class TestPKCEFunctionality: mock_response.status_code = 401 mock_response.text = "Unauthorized" - with patch("litellm.proxy.management_endpoints.ui_sso.get_async_httpx_client") as mock_get_client: + with patch( + "litellm.proxy.management_endpoints.ui_sso.get_async_httpx_client" + ) as mock_get_client: mock_client = MagicMock() mock_client.post = AsyncMock(return_value=mock_response) mock_get_client.return_value = mock_client @@ -5842,7 +6101,9 @@ class TestPKCEFunctionality: assert str(exc_info.value.code) == "401" @pytest.mark.asyncio - async def test_pkce_cache_miss_unexpected_format_non_strict_logs_warning(self, caplog): + async def test_pkce_cache_miss_unexpected_format_non_strict_logs_warning( + self, caplog + ): """When cached data has an unexpected format (e.g. integer from corrupt Redis) in non-strict mode, prepare_token_exchange_parameters logs a warning and returns params without code_verifier rather than raising.""" @@ -5886,7 +6147,9 @@ class TestPKCEFunctionality: mock_cache.async_get_cache.assert_called_once() # Verify a warning was logged about the unexpected format or cache miss assert any( - "verifier" in r.message.lower() or "format" in r.message.lower() or "cache" in r.message.lower() + "verifier" in r.message.lower() + or "format" in r.message.lower() + or "cache" in r.message.lower() for r in caplog.records if r.levelno >= logging.WARNING ), f"Expected a format/cache warning. Records: {[r.message for r in caplog.records]}" @@ -5930,7 +6193,9 @@ class TestPKCEFunctionality: from litellm.proxy._types import ProxyException from litellm.proxy.management_endpoints.ui_sso import SSOAuthenticationHandler - with patch("litellm.proxy.management_endpoints.ui_sso.get_async_httpx_client") as mock_get_client: + with patch( + "litellm.proxy.management_endpoints.ui_sso.get_async_httpx_client" + ) as mock_get_client: mock_client = MagicMock() mock_resp = MagicMock() mock_resp.status_code = 200 @@ -5964,7 +6229,9 @@ class TestPKCEFunctionality: body_without_token = {"token_type": "Bearer", "scope": "openid"} - with patch("litellm.proxy.management_endpoints.ui_sso.get_async_httpx_client") as mock_get_client: + with patch( + "litellm.proxy.management_endpoints.ui_sso.get_async_httpx_client" + ) as mock_get_client: mock_client = MagicMock() mock_resp = MagicMock() mock_resp.status_code = 200 @@ -5985,7 +6252,10 @@ class TestPKCEFunctionality: additional_headers={}, ) - assert "no access_token" in exc_info.value.message or "access_token" in exc_info.value.message + assert ( + "no access_token" in exc_info.value.message + or "access_token" in exc_info.value.message + ) assert str(exc_info.value.code) == "401" @@ -6013,16 +6283,18 @@ class TestAddMissingTeamMember: sso_teams = ["team-from-entra-1", "team-from-entra-2"] - with patch("litellm.proxy.management_endpoints.ui_sso.create_team_member_add_task") as mock_add_task: + with patch( + "litellm.proxy.management_endpoints.ui_sso.create_team_member_add_task" + ) as mock_add_task: mock_add_task.return_value = AsyncMock() await add_missing_team_member(user_info=new_user, sso_teams=sso_teams) # Bug: This assertion currently FAILS - no teams are added # because function returns early when teams is None - assert mock_add_task.call_count == 2, ( - f"Expected 2 calls to add user to teams, but got {mock_add_task.call_count}" - ) + assert ( + mock_add_task.call_count == 2 + ), f"Expected 2 calls to add user to teams, but got {mock_add_task.call_count}" called_team_ids = [call.args[0] for call in mock_add_task.call_args_list] assert set(called_team_ids) == { "team-from-entra-1", @@ -6048,7 +6320,9 @@ class TestAddMissingTeamMember: sso_teams = ["team-from-entra-1", "team-from-entra-2"] - with patch("litellm.proxy.management_endpoints.ui_sso.create_team_member_add_task") as mock_add_task: + with patch( + "litellm.proxy.management_endpoints.ui_sso.create_team_member_add_task" + ) as mock_add_task: mock_add_task.return_value = AsyncMock() await add_missing_team_member(user_info=existing_user, sso_teams=sso_teams) @@ -6082,7 +6356,9 @@ class TestAddMissingTeamMember: teams=None, # Bug: NewUserResponse defaults to None ) - with patch("litellm.proxy.management_endpoints.ui_sso.add_missing_team_member") as mock_add_member: + with patch( + "litellm.proxy.management_endpoints.ui_sso.add_missing_team_member" + ) as mock_add_member: await SSOAuthenticationHandler.add_user_to_teams_from_sso_response( result=sso_result, user_info=new_user_info, @@ -6124,7 +6400,9 @@ class TestAddMissingTeamMember: # Bug: With current code, team_member_calls will be empty # After fix: Should have 2 entries - assert len(team_member_calls) == 2, f"Expected 2 teams to be added, but got {len(team_member_calls)}" + assert ( + len(team_member_calls) == 2 + ), f"Expected 2 teams to be added, but got {len(team_member_calls)}" assert {c["team_id"] for c in team_member_calls} == { "entra-team-alpha", "entra-team-beta", @@ -6180,9 +6458,9 @@ class TestAddMissingTeamMember: ): await add_missing_team_member(user_info=user_info, sso_teams=sso_teams) - assert set(added_teams) == set(expected_teams_added), ( - f"Expected teams {expected_teams_added}, but got {added_teams}" - ) + assert set(added_teams) == set( + expected_teams_added + ), f"Expected teams {expected_teams_added}, but got {added_teams}" @pytest.mark.asyncio @@ -6248,15 +6526,17 @@ async def test_role_mappings_override_default_internal_user_params(): new_user_request = call_args.kwargs["data"] # The role from SSO should be preserved, not overridden by default_internal_user_params - assert new_user_request.user_role == "proxy_admin", ( - "SSO-extracted role should override default_internal_user_params role" - ) + assert ( + new_user_request.user_role == "proxy_admin" + ), "SSO-extracted role should override default_internal_user_params role" # Other default params should still be applied - assert new_user_request.max_budget == 100, "max_budget from default_internal_user_params should be applied" - assert new_user_request.budget_duration == "30d", ( - "budget_duration from default_internal_user_params should be applied" - ) + assert ( + new_user_request.max_budget == 100 + ), "max_budget from default_internal_user_params should be applied" + assert ( + new_user_request.budget_duration == "30d" + ), "budget_duration from default_internal_user_params should be applied" finally: # Restore original default_internal_user_params (always assign, never delattr — @@ -6327,17 +6607,19 @@ async def test_sso_role_preserved_without_role_mappings(): new_user_request = call_args.kwargs["data"] # SSO role should be preserved even without role_mappings configured - assert new_user_request.user_role == "proxy_admin", ( - "SSO role from app_roles should not be overwritten by default_internal_user_params" - ) + assert ( + new_user_request.user_role == "proxy_admin" + ), "SSO role from app_roles should not be overwritten by default_internal_user_params" # Other defaults should still apply - assert new_user_request.max_budget == 50, "max_budget from default_internal_user_params should be applied" + assert ( + new_user_request.max_budget == 50 + ), "max_budget from default_internal_user_params should be applied" # Verify user_defined_values was also updated (it's mutated in-place) - assert user_defined_values["user_role"] == "proxy_admin", ( - "user_defined_values should retain the SSO role after insert_sso_user" - ) + assert ( + user_defined_values["user_role"] == "proxy_admin" + ), "user_defined_values should retain the SSO role after insert_sso_user" finally: # Restore original default_internal_user_params (always assign, never delattr — @@ -6446,7 +6728,10 @@ class TestSSOReadinessEndpoint: assert data["sso_configured"] is True assert data["provider"] == "google" assert "GOOGLE_CLIENT_SECRET" in data["missing_environment_variables"] - assert "Google SSO is configured but missing required environment variables" in data["message"] + assert ( + "Google SSO is configured but missing required environment variables" + in data["message"] + ) finally: app.dependency_overrides.clear() @@ -6507,7 +6792,9 @@ class TestSSOReadinessEndpoint: assert data["sso_configured"] is True assert data["provider"] == expected_provider assert data["status"] == "unhealthy" - assert set(data["missing_environment_variables"]) == set(expected_missing_vars) + assert set(data["missing_environment_variables"]) == set( + expected_missing_vars + ) finally: app.dependency_overrides.clear() @@ -6575,7 +6862,9 @@ class TestSSOReadinessEndpoint: assert data["sso_configured"] is True assert data["provider"] == expected_provider assert data["status"] == "unhealthy" - assert set(data["missing_environment_variables"]) == set(expected_missing_vars) + assert set(data["missing_environment_variables"]) == set( + expected_missing_vars + ) finally: app.dependency_overrides.clear() @@ -6607,9 +6896,13 @@ class TestCustomMicrosoftSSO: discovery = await sso.get_discovery_document() assert ( - discovery["authorization_endpoint"] == "https://login.microsoftonline.com/test-tenant/oauth2/v2.0/authorize" + discovery["authorization_endpoint"] + == "https://login.microsoftonline.com/test-tenant/oauth2/v2.0/authorize" + ) + assert ( + discovery["token_endpoint"] + == "https://login.microsoftonline.com/test-tenant/oauth2/v2.0/token" ) - assert discovery["token_endpoint"] == "https://login.microsoftonline.com/test-tenant/oauth2/v2.0/token" assert discovery["userinfo_endpoint"] == "https://graph.microsoft.com/v1.0/me" @pytest.mark.asyncio @@ -6673,8 +6966,13 @@ class TestCustomMicrosoftSSO: # Custom auth endpoint assert discovery["authorization_endpoint"] == custom_auth_endpoint # Default token and userinfo endpoints - assert discovery["token_endpoint"] == "https://login.microsoftonline.com/test-tenant/oauth2/v2.0/token" - assert discovery["userinfo_endpoint"] == "https://graph.microsoft.com/v1.0/me" + assert ( + discovery["token_endpoint"] + == "https://login.microsoftonline.com/test-tenant/oauth2/v2.0/token" + ) + assert ( + discovery["userinfo_endpoint"] == "https://graph.microsoft.com/v1.0/me" + ) def test_custom_microsoft_sso_uses_common_tenant_when_none(self): """ @@ -6712,7 +7010,9 @@ async def test_setup_team_mappings(): mock_prisma = MagicMock() mock_sso_config = MagicMock() mock_sso_config.sso_settings = {"team_mappings": {"team_ids_jwt_field": "groups"}} - mock_prisma.db.litellm_ssoconfig.find_unique = AsyncMock(return_value=mock_sso_config) + mock_prisma.db.litellm_ssoconfig.find_unique = AsyncMock( + return_value=mock_sso_config + ) with patch( "litellm.proxy.utils.get_prisma_client_or_throw", @@ -6725,7 +7025,9 @@ async def test_setup_team_mappings(): assert result is not None assert isinstance(result, TeamMappings) assert result.team_ids_jwt_field == "groups" - mock_prisma.db.litellm_ssoconfig.find_unique.assert_called_once_with(where={"id": "sso_config"}) + mock_prisma.db.litellm_ssoconfig.find_unique.assert_called_once_with( + where={"id": "sso_config"} + ) # ============================================================================ @@ -7005,7 +7307,9 @@ def test_process_sso_jwt_access_token_extracts_role_from_nested_field(): user_role=None, ) - with patch.dict(os.environ, {"GENERIC_USER_ROLE_ATTRIBUTE": "resource_access.my-client.roles"}): + with patch.dict( + os.environ, {"GENERIC_USER_ROLE_ATTRIBUTE": "resource_access.my-client.roles"} + ): process_sso_jwt_access_token( access_token_str=access_token_str, sso_jwt_handler=None, @@ -7066,7 +7370,9 @@ def test_generic_response_convertor_with_extra_attributes(monkeypatch): from litellm.proxy.management_endpoints.ui_sso import generic_response_convertor monkeypatch.setenv("GENERIC_CLIENT_ID", "test_client") - monkeypatch.setenv("GENERIC_USER_EXTRA_ATTRIBUTES", "custom_field1,custom_field2,custom_field3") + monkeypatch.setenv( + "GENERIC_USER_EXTRA_ATTRIBUTES", "custom_field1,custom_field2,custom_field3" + ) mock_response = { "sub": "user-id-123", @@ -7132,7 +7438,9 @@ def test_generic_response_convertor_extra_attributes_with_nested_paths(monkeypat from litellm.proxy.management_endpoints.ui_sso import generic_response_convertor monkeypatch.setenv("GENERIC_CLIENT_ID", "test_client") - monkeypatch.setenv("GENERIC_USER_EXTRA_ATTRIBUTES", "org_info.department,org_info.manager") + monkeypatch.setenv( + "GENERIC_USER_EXTRA_ATTRIBUTES", "org_info.department,org_info.manager" + ) mock_response = { "sub": "user-id-123", @@ -7294,7 +7602,10 @@ class TestCliSsoAttributionMetadata: get_user_info_mock.assert_awaited_once() assert get_user_info_mock.call_args.kwargs["user_defined_values"] is not None - assert get_user_info_mock.call_args.kwargs["user_defined_values"]["user_id"] == "cli-test-user" + assert ( + get_user_info_mock.call_args.kwargs["user_defined_values"]["user_id"] + == "cli-test-user" + ) @pytest.mark.asyncio async def test_cli_sso_callback_rejects_restricted_sso_group(self): @@ -7429,9 +7740,13 @@ class TestCliSsoAttributionMetadata: ) flow_data = mock_cache.set_cache.call_args.kwargs["value"] - assert flow_data["session_data"]["attribution_metadata"] == {"acme_employment_type": "contractor"} + assert flow_data["session_data"]["attribution_metadata"] == { + "acme_employment_type": "contractor" + } mock_prisma.db.litellm_usertable.update_many.assert_awaited_once() - update_data = mock_prisma.db.litellm_usertable.update_many.call_args.kwargs["data"] + update_data = mock_prisma.db.litellm_usertable.update_many.call_args.kwargs[ + "data" + ] assert update_data["metadata"]["acme_employment_type"] == "contractor" assert update_data["metadata"]["auth_provider"] == "generic" @@ -7483,7 +7798,9 @@ class TestValidateReturnTo: def test_returns_false_when_no_control_plane_url_configured(self, monkeypatch): """return_to should be silently ignored if control_plane_url is not in general_settings.""" monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {}) - result = SSOAuthenticationHandler._validate_return_to("https://cp.example.com/ui") + result = SSOAuthenticationHandler._validate_return_to( + "https://cp.example.com/ui" + ) assert result is False def test_allows_matching_origin(self, monkeypatch): @@ -7493,7 +7810,9 @@ class TestValidateReturnTo: {"control_plane_url": "https://cp.example.com"}, ) # Should not raise - SSOAuthenticationHandler._validate_return_to("https://cp.example.com/ui?page=models") + SSOAuthenticationHandler._validate_return_to( + "https://cp.example.com/ui?page=models" + ) def test_allows_matching_origin_with_trailing_slash(self, monkeypatch): """Trailing slash on control_plane_url should not affect origin comparison.""" @@ -7510,7 +7829,9 @@ class TestValidateReturnTo: {"control_plane_url": "https://cp.example.com"}, ) with pytest.raises(HTTPException) as exc_info: - SSOAuthenticationHandler._validate_return_to("https://cp.example.com.evil.com/steal") + SSOAuthenticationHandler._validate_return_to( + "https://cp.example.com.evil.com/steal" + ) assert exc_info.value.status_code == 400 def test_rejects_different_origin(self, monkeypatch): @@ -7549,7 +7870,9 @@ class TestValidateReturnTo: {"control_plane_url": "https://cp.example.com"}, ) with pytest.raises(HTTPException) as exc_info: - SSOAuthenticationHandler._validate_return_to("https://cp.example.com:8443/ui") + SSOAuthenticationHandler._validate_return_to( + "https://cp.example.com:8443/ui" + ) assert exc_info.value.status_code == 400 def test_allows_explicit_default_port(self, monkeypatch): @@ -7675,7 +7998,9 @@ class TestSyncUserRoleFromJwtRoleMap: user_id=user_id, user_role=LitellmUserRoles.INTERNAL_USER_VIEW_ONLY.value, ) - await cache.async_set_cache(key=user_id, value=existing_user.model_dump(), ttl=60) + await cache.async_set_cache( + key=user_id, value=existing_user.model_dump(), ttl=60 + ) sso_values = self._make_sso_values( user_role=LitellmUserRoles.INTERNAL_USER_VIEW_ONLY.value, @@ -7756,7 +8081,9 @@ class TestPKCEStateCookieBinding: SSOAuthenticationHandler, ) - mock_redirect = RedirectResponse(url="https://idp.example.com/authorize?state=test-state-xyz") + mock_redirect = RedirectResponse( + url="https://idp.example.com/authorize?state=test-state-xyz" + ) mock_generic_sso = MagicMock() mock_generic_sso.__enter__ = MagicMock(return_value=mock_generic_sso) mock_generic_sso.__exit__ = MagicMock(return_value=None) @@ -7777,8 +8104,12 @@ class TestPKCEStateCookieBinding: assert response is not None cookie_headers = response.headers.getlist("set-cookie") - cookie_str = next((c for c in cookie_headers if "litellm_oauth_state=" in c), None) - assert cookie_str is not None, f"litellm_oauth_state cookie not set; got: {cookie_headers}" + cookie_str = next( + (c for c in cookie_headers if "litellm_oauth_state=" in c), None + ) + assert ( + cookie_str is not None + ), f"litellm_oauth_state cookie not set; got: {cookie_headers}" assert "test-state-xyz" in cookie_str assert "HttpOnly" in cookie_str assert "SameSite=lax" in cookie_str @@ -7800,7 +8131,9 @@ class TestPKCEStateCookieBinding: SSOAuthenticationHandler, ) - mock_redirect = RedirectResponse(url="https://idp.example.com/authorize?state=test-state-xyz") + mock_redirect = RedirectResponse( + url="https://idp.example.com/authorize?state=test-state-xyz" + ) mock_generic_sso = MagicMock() mock_generic_sso.__enter__ = MagicMock(return_value=mock_generic_sso) mock_generic_sso.__exit__ = MagicMock(return_value=None) @@ -7821,9 +8154,9 @@ class TestPKCEStateCookieBinding: assert response is not None cookie_headers = response.headers.getlist("set-cookie") - assert not any("litellm_oauth_state=" in c for c in cookie_headers), ( - f"litellm_oauth_state cookie set on non-PKCE flow; got: {cookie_headers}" - ) + assert not any( + "litellm_oauth_state=" in c for c in cookie_headers + ), f"litellm_oauth_state cookie set on non-PKCE flow; got: {cookie_headers}" @pytest.mark.asyncio async def test_redirect_response_drops_secure_flag_for_http_dev(self): @@ -7836,7 +8169,9 @@ class TestPKCEStateCookieBinding: SSOAuthenticationHandler, ) - mock_redirect = RedirectResponse(url="http://idp.local/authorize?state=local-dev-state") + mock_redirect = RedirectResponse( + url="http://idp.local/authorize?state=local-dev-state" + ) mock_generic_sso = MagicMock() mock_generic_sso.__enter__ = MagicMock(return_value=mock_generic_sso) mock_generic_sso.__exit__ = MagicMock(return_value=None) @@ -7860,12 +8195,16 @@ class TestPKCEStateCookieBinding: ) cookie_headers = response.headers.getlist("set-cookie") - cookie_str = next((c for c in cookie_headers if "litellm_oauth_state=" in c), None) + cookie_str = next( + (c for c in cookie_headers if "litellm_oauth_state=" in c), None + ) assert cookie_str is not None assert "Secure" not in cookie_str @pytest.mark.asyncio - async def test_redirect_response_sets_secure_flag_behind_trusted_tls_terminating_proxy(self, monkeypatch): + async def test_redirect_response_sets_secure_flag_behind_trusted_tls_terminating_proxy( + self, monkeypatch + ): """Regression: litellm sees a plain-HTTP hop when TLS terminates at a reverse proxy. The Secure flag must still be set when the direct peer is a configured trusted proxy and it reports X-Forwarded-Proto: https -- but NOT from an @@ -7877,7 +8216,9 @@ class TestPKCEStateCookieBinding: SSOAuthenticationHandler, ) - mock_redirect = RedirectResponse(url="http://idp.internal/authorize?state=behind-proxy-state") + mock_redirect = RedirectResponse( + url="http://idp.internal/authorize?state=behind-proxy-state" + ) mock_generic_sso = MagicMock() mock_generic_sso.__enter__ = MagicMock(return_value=mock_generic_sso) mock_generic_sso.__exit__ = MagicMock(return_value=None) @@ -7909,12 +8250,16 @@ class TestPKCEStateCookieBinding: ) cookie_headers = response.headers.getlist("set-cookie") - cookie_str = next((c for c in cookie_headers if "litellm_oauth_state=" in c), None) + cookie_str = next( + (c for c in cookie_headers if "litellm_oauth_state=" in c), None + ) assert cookie_str is not None assert "Secure" in cookie_str @pytest.mark.asyncio - async def test_redirect_response_ignores_spoofed_forwarded_proto_without_trust_config(self, monkeypatch): + async def test_redirect_response_ignores_spoofed_forwarded_proto_without_trust_config( + self, monkeypatch + ): """The same X-Forwarded-Proto: https header must NOT flip Secure on when no trusted-proxy config is present -- honoring it unconditionally would let any client spoof the header and would not itself be the vulnerability the ticket @@ -7925,7 +8270,9 @@ class TestPKCEStateCookieBinding: SSOAuthenticationHandler, ) - mock_redirect = RedirectResponse(url="http://idp.internal/authorize?state=spoofed-state") + mock_redirect = RedirectResponse( + url="http://idp.internal/authorize?state=spoofed-state" + ) mock_generic_sso = MagicMock() mock_generic_sso.__enter__ = MagicMock(return_value=mock_generic_sso) mock_generic_sso.__exit__ = MagicMock(return_value=None) @@ -7954,7 +8301,9 @@ class TestPKCEStateCookieBinding: ) cookie_headers = response.headers.getlist("set-cookie") - cookie_str = next((c for c in cookie_headers if "litellm_oauth_state=" in c), None) + cookie_str = next( + (c for c in cookie_headers if "litellm_oauth_state=" in c), None + ) assert cookie_str is not None assert "Secure" not in cookie_str @@ -8361,15 +8710,9 @@ async def _render_debug_page(provider_env, id_jag_registered, force_inert=False) AsyncMock(return_value=id_jag_registered), ), patch("litellm.proxy.proxy_server.general_settings", {}), # test-quality-ok: debug endpoint reads proxy globals - patch( - "litellm.proxy.proxy_server.prisma_client", MagicMock() - ), # test-quality-ok: debug endpoint reads proxy DB - patch( - "litellm.proxy.proxy_server.user_api_key_cache", MagicMock() - ), # test-quality-ok: debug endpoint reads proxy globals - patch( - "litellm.proxy.proxy_server.jwt_handler", MagicMock(spec=JWTHandler) - ), # test-quality-ok: debug endpoint reads proxy globals + patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), # test-quality-ok: debug endpoint reads proxy DB + patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()), # test-quality-ok: debug endpoint reads proxy globals + patch("litellm.proxy.proxy_server.jwt_handler", MagicMock(spec=JWTHandler)), # test-quality-ok: debug endpoint reads proxy globals ] if force_inert: stack.append( @@ -8407,7 +8750,9 @@ async def test_debug_page_logs_the_capture_gap_but_never_renders_it(caplog): async def test_debug_page_is_byte_identical_when_the_provider_captures(): """A deployment with no gap must get the page it got before this change, to the byte. The comparison is against the endpoint with the diagnostic forced inert, not against a guess.""" - with_feature = await _render_debug_page({"GENERIC_CLIENT_ID": _GENERIC_DEBUG_CLIENT_ID}, id_jag_registered=True) + with_feature = await _render_debug_page( + {"GENERIC_CLIENT_ID": _GENERIC_DEBUG_CLIENT_ID}, id_jag_registered=True + ) pre_change = await _render_debug_page( {"GENERIC_CLIENT_ID": _GENERIC_DEBUG_CLIENT_ID}, id_jag_registered=True, @@ -8422,7 +8767,9 @@ async def test_debug_page_is_byte_identical_when_the_provider_captures(): async def test_debug_page_is_byte_identical_when_no_id_jag_server_is_registered(): """Most deployments run Google SSO and no id_jag server at all; their debug page must not grow an ID-JAG section about a feature they do not use.""" - with_feature = await _render_debug_page({"GOOGLE_CLIENT_ID": _GOOGLE_DEBUG_CLIENT_ID}, id_jag_registered=False) + with_feature = await _render_debug_page( + {"GOOGLE_CLIENT_ID": _GOOGLE_DEBUG_CLIENT_ID}, id_jag_registered=False + ) pre_change = await _render_debug_page( {"GOOGLE_CLIENT_ID": _GOOGLE_DEBUG_CLIENT_ID}, id_jag_registered=False, @@ -8586,19 +8933,16 @@ async def test_saml_callback_enforces_free_sso_user_limit_after_validation(): request_double = SimpleNamespace(cookies={}, headers={}, stream=_stream) - with ( - patch.dict(os.environ, {"DISABLE_ADMIN_UI": "false"}), - patch("litellm.proxy.proxy_server.premium_user", False), - patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), - patch("litellm.proxy.proxy_server.master_key", MASTER_KEY), - patch( - "litellm.proxy.management_endpoints.sso.saml_sso.SAMLAuthHandler.handle_acs", - new=_fake_handle_acs, - ), - patch( - "litellm.repositories.user_repository.UserRepository.count_billable_users", - new=AsyncMock(side_effect=_fake_count_billable_users), - ), + with patch.dict(os.environ, {"DISABLE_ADMIN_UI": "false"}), patch( + "litellm.proxy.proxy_server.premium_user", False + ), patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), patch( + "litellm.proxy.proxy_server.master_key", MASTER_KEY + ), patch( + "litellm.proxy.management_endpoints.sso.saml_sso.SAMLAuthHandler.handle_acs", + new=_fake_handle_acs, + ), patch( + "litellm.repositories.user_repository.UserRepository.count_billable_users", + new=AsyncMock(side_effect=_fake_count_billable_users), ): with pytest.raises(ProxyException) as exc: await saml_callback(request_double) @@ -8814,7 +9158,9 @@ async def test_verify_and_process_arm_captures_sso_assertion(): mock_jwt_handler.get_team_ids_from_jwt.return_value = [] mock_sso_instance = MagicMock() - mock_sso_instance.verify_and_process = AsyncMock(return_value={"sub": "u1", "email": "u@example.com"}) + mock_sso_instance.verify_and_process = AsyncMock( + return_value={"sub": "u1", "email": "u@example.com"} + ) mock_sso_instance.access_token = None mock_sso_instance.id_token = id_token mock_sso_instance.refresh_token = "rt_from_idp" @@ -8830,7 +9176,9 @@ async def test_verify_and_process_arm_captures_sso_assertion(): }, ): with patch("fastapi_sso.sso.base.DiscoveryDocument"): - with patch("fastapi_sso.sso.generic.create_provider", return_value=mock_sso_class): + with patch( + "fastapi_sso.sso.generic.create_provider", return_value=mock_sso_class + ): _, _, _, sso_assertion = await get_generic_sso_response( request=mock_request, jwt_handler=mock_jwt_handler, @@ -8873,7 +9221,9 @@ async def test_redirect_from_openid_persists_assertion_under_canonical_user_id() patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()), patch( "litellm.proxy.proxy_server.generate_key_helper_fn", - AsyncMock(return_value={"token": "sk-ui-key", "user_id": "canonical-user-id"}), + AsyncMock( + return_value={"token": "sk-ui-key", "user_id": "canonical-user-id"} + ), ), patch( "litellm.proxy.management_endpoints.ui_sso.get_user_info_from_db", @@ -8908,7 +9258,9 @@ async def test_redirect_from_openid_persists_assertion_under_canonical_user_id() sso_assertion=assertion, ) - retain_mock.assert_awaited_once_with(user_id="canonical-user-id", assertion=assertion) + retain_mock.assert_awaited_once_with( + user_id="canonical-user-id", assertion=assertion + ) assert response is not None @@ -9105,13 +9457,9 @@ async def test_browser_funnel_reports_an_uncaptured_assertion(monkeypatch, caplo patch("litellm.proxy.proxy_server.general_settings", {}), # test-quality-ok: endpoint reads proxy globals patch("litellm.proxy.proxy_server.premium_user", False), # test-quality-ok: endpoint reads proxy globals patch("litellm.proxy.proxy_server.user_custom_sso", None), # test-quality-ok: endpoint reads proxy globals - patch( - "litellm.proxy.proxy_server.proxy_logging_obj", MagicMock() - ), # test-quality-ok: endpoint reads proxy globals + patch("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock()), # test-quality-ok: endpoint reads proxy globals patch("litellm.proxy.proxy_server.redis_usage_cache", None), # test-quality-ok: endpoint reads proxy globals - patch( - "litellm.proxy.proxy_server.user_api_key_cache", MagicMock() - ), # test-quality-ok: endpoint reads proxy globals + patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()), # test-quality-ok: endpoint reads proxy globals patch( # test-quality-ok: endpoint test stubs key generation at its module boundary "litellm.proxy.proxy_server.generate_key_helper_fn", AsyncMock(return_value={"token": "sk-ui-key", "user_id": "canonical-user-id"}), @@ -9272,7 +9620,9 @@ async def test_cli_completion_drops_teams_whose_rows_no_longer_exist(): _complete_cli_sso_callback_session, ) - live_detail = CliSsoTeamDetail(team_id="team-live", team_alias="Live", team_models=("gpt-4.1",)) + live_detail = CliSsoTeamDetail( + team_id="team-live", team_alias="Live", team_models=("gpt-4.1",) + ) flow = {} with ( patch( @@ -9584,16 +9934,8 @@ async def test_cli_sign_in_enrolls_only_verified_subjects_before_completing( if trusted: table.upsert.assert_awaited_once_with( where={"issuer_tenant_id_oid": {"issuer": "issuer", "tenant_id": "tenant", "oid": "subject"}}, - data={ - "create": { - "issuer": "issuer", - "tenant_id": "tenant", - "oid": "subject", - "user_id": "cli-user-id", - "verified_via": "sso_interactive", - }, - "update": {}, - }, + data={"create": {"issuer": "issuer", "tenant_id": "tenant", "oid": "subject", + "user_id": "cli-user-id", "verified_via": "sso_interactive"}, "update": {}}, ) else: table.upsert.assert_not_awaited() @@ -9612,9 +9954,7 @@ async def _sso_key_generate_on_ui_disabled_node(*, source, key, google_sso_confi discovery document.""" from litellm.proxy.management_endpoints.ui_sso import _get_cli_sso_flow_cache_key, google_login - env_without_sso_providers = { - name: value for name, value in os.environ.items() if name not in _SSO_PROVIDER_ENV_VARS - } + env_without_sso_providers = {name: value for name, value in os.environ.items() if name not in _SSO_PROVIDER_ENV_VARS} env = { **env_without_sso_providers, "DISABLE_ADMIN_UI": "true", From 334e4e35e47293e2ed433a25bb3523bc152f8f15 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang <yuneng@berri.ai> Date: Sat, 10 Oct 2026 19:33:58 -0700 Subject: [PATCH 4/5] fix(sso): preserve admin organization defaults during concurrent login --- litellm/proxy/management_endpoints/ui_sso.py | 4 +++- .../proxy/management_endpoints/test_ui_sso.py | 17 +++++++++++------ 2 files changed, 14 insertions(+), 7 deletions(-) diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index f8e0d3ab98e..cee7f928a8e 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -2703,7 +2703,9 @@ async def _concurrently_inserted_sso_user( getattr(litellm, "default_internal_user_params", None) or {} ) required_organizations: Final = TypeAdapter(tuple[str, ...]).validate_python( - default_parameters.get("organizations") or () + () + 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) diff --git a/tests/unit/proxy/management_endpoints/test_ui_sso.py b/tests/unit/proxy/management_endpoints/test_ui_sso.py index 1b69b511543..b1899e44887 100644 --- a/tests/unit/proxy/management_endpoints/test_ui_sso.py +++ b/tests/unit/proxy/management_endpoints/test_ui_sso.py @@ -1348,17 +1348,22 @@ async def test_concurrent_sso_creation_refuses_partial_membership_across_multipl @pytest.mark.asyncio -async def test_public_sso_upsert_recovers_an_actual_duplicate_user_refusal(monkeypatch: pytest.MonkeyPatch) -> None: +@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": "internal_user"} + 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", None) - values: Final[SSOUserDefinedValues] = {**RACE_VALUES} + 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="internal_user", team_ids=[]), + 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, @@ -1368,7 +1373,7 @@ async def test_public_sso_upsert_recovers_an_actual_duplicate_user_refusal(monke assert (recovered.user_id, recovered.sso_user_id, recovered.user_role) == ( RACE_SUBJECT, RACE_SUBJECT, - "internal_user", + user_role, ) From e08a3f8f3aae40a55e4cec78a72328459e0a3278 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang <yuneng@berri.ai> Date: Sat, 10 Oct 2026 21:13:12 -0700 Subject: [PATCH 5/5] test: hold replica SSO prechecks before user creation --- .../authorization/test_cli_sso_login_ui_disabled.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) 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 6aba820ef96..3fd7018eaae 100644 --- a/tests/integration/authorization/test_cli_sso_login_ui_disabled.py +++ b/tests/integration/authorization/test_cli_sso_login_ui_disabled.py @@ -471,15 +471,17 @@ def test_cli_link_is_reentrant_until_the_poll_consumes_the_session(ui_disabled: 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(writer_url, DUPLICATE_USER_ID_CHECK) as (relay, relayed_url), + 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_URL": relayed_url, + database_setting: relayed_url, "PRISMA_HEALTH_WATCHDOG_ENABLED": "false", **_sso_environment(idp.wire.url), },