fix(sso): ignore blank team_id_jwt_field claims and warn when creating the team fails

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
             )
 
 

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