diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index e00713bb0d7..0171a59b8e8 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 functools import partial from html import escape from types import MappingProxyType, TracebackType @@ -1267,7 +1268,9 @@ def generic_response_convertor( if jwt_handler.litellm_jwtauth.team_id_upsert and team_id_jwt_field is not None else None ) - team_id_to_upsert: Final = upsert_candidate if isinstance(upsert_candidate, str) else None + team_id_to_upsert: Final = ( + upsert_candidate if isinstance(upsert_candidate, str) and upsert_candidate.strip() else None + ) team_ids_with_upsert: Final = ( all_teams if team_id_to_upsert is None or team_id_to_upsert in all_teams else [*all_teams, team_id_to_upsert] ) @@ -1820,11 +1823,18 @@ async def warn_if_id_jag_capture_gap(*, retention_enabled: RetentionCheck | None verbose_proxy_logger.warning("SSO debug callback ran with an oauth2_id_jag capture gap: %s", gap) +@dataclass(frozen=True, slots=True) +class _TeamJoinFailure: + team_id: str + error: Exception + team_does_not_exist: bool + + async def create_team_member_add_task( team_id: str, user_info: NewUserResponse | LiteLLM_UserTable, upsert_team: Callable[[str], Awaitable[object]] | None = None, -) -> Exception | None: +) -> _TeamJoinFailure | None: """Add the user to one team, creating it first when ``upsert_team`` is given. Returns the failure, if any.""" upsert_error: Final = await _upsert_team_error(team_id, upsert_team) try: @@ -1833,7 +1843,11 @@ async def create_team_member_add_task( user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), ) except Exception as e: - return e if upsert_error is None else upsert_error + return _TeamJoinFailure( + team_id=team_id, + error=e if upsert_error is None else upsert_error, + team_does_not_exist=upsert_team is None and _is_missing_team_error(e), + ) return None @@ -1880,8 +1894,8 @@ async def add_missing_team_member( for team_id in missing_teams ) ) - failures: Final = tuple((team_id, error) for team_id, error in zip(missing_teams, results) if error is not None) - skipped: Final = tuple(team_id for team_id, error in failures if _is_missing_team_error(error)) + failures: Final = tuple(failure for failure in results if failure is not None) + skipped: Final = tuple(failure.team_id for failure in failures if failure.team_does_not_exist) if skipped: verbose_proxy_logger.info( "SSO login for user %s skipped %d team(s) that do not exist in the db: %s", @@ -1889,10 +1903,10 @@ async def add_missing_team_member( len(skipped), ", ".join(skipped), ) - for team_id, error in failures: - if not _is_missing_team_error(error): + for failure in failures: + if not failure.team_does_not_exist: verbose_proxy_logger.warning( - "SSO login could not add user %s to team %s: %s", user_info.user_id, team_id, error + "SSO login could not add user %s to team %s: %s", user_info.user_id, failure.team_id, failure.error ) diff --git a/tests/unit/proxy/management_endpoints/test_ui_sso.py b/tests/unit/proxy/management_endpoints/test_ui_sso.py index a0190190db0..0f8cf3b65bf 100644 --- a/tests/unit/proxy/management_endpoints/test_ui_sso.py +++ b/tests/unit/proxy/management_endpoints/test_ui_sso.py @@ -6599,7 +6599,7 @@ class TestAddMissingTeamMember: with patch( "litellm.proxy.management_endpoints.ui_sso.create_team_member_add_task" ) as mock_add_task: - mock_add_task.return_value = AsyncMock() + mock_add_task.return_value = None await add_missing_team_member(user_info=new_user, sso_teams=sso_teams) @@ -6636,7 +6636,7 @@ class TestAddMissingTeamMember: with patch( "litellm.proxy.management_endpoints.ui_sso.create_team_member_add_task" ) as mock_add_task: - mock_add_task.return_value = AsyncMock() + mock_add_task.return_value = None await add_missing_team_member(user_info=existing_user, sso_teams=sso_teams) @@ -10970,6 +10970,8 @@ def _upsert_jwt_handler(**jwtauth: object) -> JWTHandler: ({"team_id_jwt_field": "dept", "team_id_upsert": True}, {"dept": "finance"}, ["finance"], "finance"), ({"team_id_jwt_field": "dept", "team_id_upsert": False}, {"dept": "finance"}, ["finance"], None), ({"team_id_jwt_field": "dept", "team_id_upsert": True}, {}, [], None), + ({"team_id_jwt_field": "dept", "team_id_upsert": True}, {"dept": ""}, [], None), + ({"team_id_jwt_field": "dept", "team_id_upsert": True}, {"dept": " "}, [" "], None), ({"team_id_jwt_field": "dept", "team_id_upsert": True}, {"dept": ["a", "b"]}, ["a", "b"], None), ( {"team_id_jwt_field": "dept", "team_id_upsert": True, "team_id_default": "everyone"}, @@ -10994,6 +10996,8 @@ def _upsert_jwt_handler(**jwtauth: object) -> JWTHandler: "upsert-on", "upsert-off", "claim-missing", + "blank-claim-not-upserted", + "whitespace-claim-not-upserted", "list-valued-claim-join-only", "default-team-not-upserted", "plural-claim-join-only", @@ -11061,16 +11065,20 @@ async def test_sso_login_creates_only_the_upsert_team_and_does_so_before_joining async def test_sso_login_still_joins_the_team_when_a_concurrent_login_created_it_first() -> None: from litellm.proxy.management_endpoints.ui_sso import add_missing_team_member + joined: Final[list[str]] = [] # mutable-ok: records the joins the proxy would persist + async def lost_the_create_race(team_id: str) -> object: raise HTTPException(status_code=400, detail={"error": f"Team id = {team_id} already exists"}) - member_add: Final = AsyncMock(return_value=None) + async def member_add(data: TeamMemberAddRequest, user_api_key_dict: UserAPIKeyAuth) -> None: + joined.append(data.team_id) + user: Final = NewUserResponse(user_id="alice", key="sk-x", teams=None) with patch("litellm.proxy.management_endpoints.ui_sso.team_member_add", member_add): await add_missing_team_member(user, ["finance"], team_id_to_upsert="finance", upsert_team=lost_the_create_race) - assert [call.kwargs["data"].team_id for call in member_add.call_args_list] == ["finance"] + assert joined == ["finance"] @pytest.mark.asyncio @@ -11095,6 +11103,37 @@ async def test_sso_login_logs_skipped_idp_groups_once_per_login_and_not_as_warni assert [r for r in caplog.records if r.levelno >= logging.WARNING] == [] +@pytest.mark.asyncio +@pytest.mark.parametrize( + "create_error", + [ + HTTPException(status_code=403, detail={"error": "team limit reached for this license"}), + HTTPException(status_code=404, detail={"error": "Team doesn't exist in db. Team=finance"}), + ], +) +async def test_sso_login_warns_when_creating_the_team_fails_instead_of_calling_it_missing( + caplog: pytest.LogCaptureFixture, create_error: HTTPException +) -> None: + from litellm.proxy.management_endpoints.ui_sso import add_missing_team_member + + async def failed_create(team_id: str) -> object: + raise create_error + + user: Final = NewUserResponse(user_id="alice", key="sk-x", teams=None) + with ( + caplog.at_level(logging.INFO, logger=verbose_proxy_logger.name), + patch( + "litellm.proxy.management_endpoints.ui_sso.team_member_add", + AsyncMock(side_effect=TeamNotFoundError(team_id="finance")), + ), + ): + await add_missing_team_member(user, ["finance"], team_id_to_upsert="finance", upsert_team=failed_create) + + warnings: Final = [r.getMessage() for r in caplog.records if r.levelno == logging.WARNING] + assert len(warnings) == 1 and "finance" in warnings[0] and str(create_error.detail) in warnings[0] + assert [r for r in caplog.records if "skipped" in r.getMessage()] == [] + + @pytest.mark.asyncio async def test_sso_login_warns_when_joining_an_existing_team_fails(caplog: pytest.LogCaptureFixture) -> None: from litellm.proxy.management_endpoints.ui_sso import add_missing_team_member