fix(sso): upsert the team_id_jwt_field team with db team_mappings, join after a lost create race, log skipped teams once

litellm/proxy/management_endpoints/ui_sso.py

@@ -1240,7 +1240,7 @@ def generic_response_convertor(
     )
 
     token: Final = cast(dict[str, object], response)  # cast-ok: the IdP userinfo JSON object fastapi-sso hands back
-    all_teams: Final = []
+    all_teams: Final[list[str]] = []
     if sso_jwt_handler is not None:
         team_ids = sso_jwt_handler.get_all_jwt_team_ids(token)
         all_teams.extend(team_ids)
@@ -1267,6 +1267,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
+    if team_id_to_upsert is not None and team_id_to_upsert not in all_teams:
+        all_teams.append(team_id_to_upsert)
 
     # Determine user role based on role_mappings if available
     # Only apply role_mappings for GENERIC SSO provider
@@ -1335,9 +1338,7 @@ def generic_response_convertor(
         last_name=get_nested_value(response, generic_user_last_name_attribute_name),
         provider=get_nested_value(response, generic_provider_attribute_name),
         team_ids=all_teams,
-        team_id_to_upsert=upsert_candidate
-        if isinstance(upsert_candidate, str) and upsert_candidate in all_teams
-        else None,
+        team_id_to_upsert=team_id_to_upsert,
         user_role=user_role,
         extra_fields=extra_fields,
     )
@@ -1822,22 +1823,33 @@ async def create_team_member_add_task(
     team_id: str,
     user_info: NewUserResponse | LiteLLM_UserTable,
     upsert_team: Callable[[str], Awaitable[object]] | None = None,
-):
-    """Create a task for adding a member to a team."""
+) -> Exception | 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:
-        if upsert_team is not None:
-            await upsert_team(team_id)
-        member: Final = Member(user_id=user_info.user_id, role="user")
-        team_member_add_request: Final = TeamMemberAddRequest(
-            member=member,
-            team_id=team_id,
-        )
-        return await team_member_add(
-            data=team_member_add_request,
+        await team_member_add(
+            data=TeamMemberAddRequest(member=Member(user_id=user_info.user_id, role="user"), team_id=team_id),
             user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN),
         )
     except Exception as e:
-        verbose_proxy_logger.warning("SSO login could not add user %s to team %s: %s", user_info.user_id, team_id, e)
+        return e if upsert_error is None else upsert_error
+    return None
+
+
+async def _upsert_team_error(team_id: str, upsert_team: Callable[[str], Awaitable[object]] | None) -> Exception | None:
+    """A concurrent login may create the team first, so a failed upsert is not fatal: the join below still runs."""
+    if upsert_team is None:
+        return None
+    try:
+        await upsert_team(team_id)
+    except Exception as e:
+        verbose_proxy_logger.debug("SSO login could not create team %s, trying to join it anyway: %s", team_id, e)
+        return e
+    return None
+
+
+def _is_missing_team_error(error: Exception) -> bool:
+    return isinstance(error, HTTPException) and error.status_code == 404
 
 
 async def _upsert_sso_team(team_id: str, prisma_client: PrismaClient, user_api_key_cache: UserApiKeyCache) -> object:
@@ -1855,25 +1867,34 @@ async def add_missing_team_member(
     sso_teams: list[str],
     team_id_to_upsert: str | None = None,
     upsert_team: Callable[[str], Awaitable[object]] | None = None,
-):
+) -> None:
     """
     - Get missing teams (diff b/w user_info.team_ids and sso_teams)
     - Add missing user to missing teams
     """
     # Handle None as empty list for new users
     user_teams: Final = user_info.teams if user_info.teams is not None else []
-    missing_teams: Final = set(sso_teams) - set(user_teams)
-    missing_teams_list: Final = list(missing_teams)
-    tasks = []
-    tasks = [
-        create_team_member_add_task(team_id, user_info, upsert_team if team_id == team_id_to_upsert else None)
-        for team_id in missing_teams_list
-    ]
-
-    try:
-        await asyncio.gather(*tasks)
-    except Exception as e:
-        verbose_proxy_logger.debug("[Non-Blocking] Error trying to add sso user to db: %s", e)
+    missing_teams: Final = sorted(set(sso_teams) - set(user_teams))
+    results: Final = await asyncio.gather(
+        *(
+            create_team_member_add_task(team_id, user_info, upsert_team if team_id == team_id_to_upsert else None)
+            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))
+    if skipped:
+        verbose_proxy_logger.info(
+            "SSO login for user %s skipped %d team(s) that do not exist in the db: %s",
+            user_info.user_id,
+            len(skipped),
+            ", ".join(skipped),
+        )
+    for team_id, error in failures:
+        if not _is_missing_team_error(error):
+            verbose_proxy_logger.warning(
+                "SSO login could not add user %s to team %s: %s", user_info.user_id, team_id, error
+            )
 
 
 def get_disabled_non_admin_personal_key_creation():

tests/integration/management/test_sso_sign_in.py

@@ -16,7 +16,7 @@ import yaml
 from pydantic import JsonValue
 
 from tests.integration._support.client import Gateway, gateway_from_environment
-from tests.integration._support.database import read_rows
+from tests.integration._support.database import read_rows, write_rows
 from tests.integration._support.process import OWNED_STATE_VARIABLE, OwnedState, owned_proxy
 from tests.integration._support.wire import Reply, Request, wire_server
 
@@ -53,6 +53,7 @@ def _idp_reply(request: Request) -> Reply:
                     "email": f"{signed_in}@example.com",
                     "provider": IDP_PROVIDER,
                     "dept": _dept(signed_in),
+                    "groups": [_group(signed_in)],
                 }
             ).encode()
         )
@@ -63,11 +64,16 @@ def _dept(subject: str) -> str:
     return f"dept-{subject}"
 
 
+def _group(subject: str) -> str:
+    return f"group-{subject}"
+
+
 @contextmanager
 def _sso_gateway(
     directory: Path,
     litellm_settings: Mapping[str, JsonValue],
     general_settings: Mapping[str, JsonValue] | None = None,
+    db_sso_settings: Mapping[str, JsonValue] | None = None,
 ) -> Iterator[Gateway]:
     config: Final = directory / "sso_proxy_config.yaml"
     config.write_text(
@@ -93,7 +99,33 @@ def _sso_gateway(
             "OAUTHLIB_INSECURE_TRANSPORT": "1",
         }
         with owned_proxy(rig, directory, environment, config=config, remove_environment=("PROXY_BASE_URL",)) as gw:
-            yield gw
+            if db_sso_settings is None:
+                yield gw
+                return
+            with contextlib.ExitStack() as cleanups:
+                cleanups.callback(
+                    write_rows,
+                    'DELETE FROM "LiteLLM_SSOConfig" WHERE id = %s',
+                    ("sso_config",),
+                    database_url=gw.database_url,
+                )
+                stored: Final = _store_sso_settings(gw, {**_generic_sso_settings(environment), **db_sso_settings})
+                assert stored.status_code == 200, stored.text
+                yield gw
+
+
+def _generic_sso_settings(environment: Mapping[str, str]) -> Mapping[str, JsonValue]:
+    return {
+        "generic_client_id": environment["GENERIC_CLIENT_ID"],
+        "generic_client_secret": environment["GENERIC_CLIENT_SECRET"],
+        "generic_authorization_endpoint": environment["GENERIC_AUTHORIZATION_ENDPOINT"],
+        "generic_token_endpoint": environment["GENERIC_TOKEN_ENDPOINT"],
+        "generic_userinfo_endpoint": environment["GENERIC_USERINFO_ENDPOINT"],
+    }
+
+
+def _store_sso_settings(gateway: Gateway, settings: Mapping[str, JsonValue]) -> httpx.Response:
+    return gateway.request("PATCH", "/update/sso_settings", settings)
 
 
 def _sign_in(gateway: Gateway, subject: str) -> httpx.Response:
@@ -201,3 +233,28 @@ def test_sso_sign_in_upserts_the_team_id_jwt_field_team_only_when_team_id_upsert
         teams: Final = gateway.rows('SELECT team_id FROM "LiteLLM_TeamTable" WHERE team_id = %s', (team_id,))
         user: Final = gateway.rows('SELECT teams FROM "LiteLLM_UserTable" WHERE user_id = %s', (subject,))
         assert (len(teams), user) == (int(expect_member), [{"teams": [team_id] if expect_member else []}])
+
+
+def test_sso_sign_in_with_db_team_mappings_creates_the_team_id_jwt_field_team_and_only_joins_groups(
+    tmp_path: Path,
+) -> None:
+    subject: Final = f"sso-mapped-user-{uuid.uuid4().hex[:12]}"
+    team_id: Final = _dept(subject)
+    group_team_id: Final = _group(subject)
+    jwtauth: Final = {"litellm_jwtauth": {"team_id_jwt_field": "dept", "team_id_upsert": True}}
+    team_mappings: Final = {"team_mappings": {"team_ids_jwt_field": "groups"}}
+    with (
+        _sso_gateway(tmp_path, {}, jwtauth, db_sso_settings=team_mappings) as gateway,
+        contextlib.ExitStack() as cleanups,
+    ):
+        cleanups.callback(gateway.request, "POST", "/team/delete", {"team_ids": [team_id, group_team_id]})
+        cleanups.callback(_delete_user, gateway, subject)
+
+        callback: Final = _sign_in(gateway, subject)
+
+        assert callback.status_code == 303, f"{callback.status_code} {callback.text}"
+        teams: Final = gateway.rows(
+            'SELECT team_id FROM "LiteLLM_TeamTable" WHERE team_id IN (%s, %s)', (team_id, group_team_id)
+        )
+        user: Final = gateway.rows('SELECT teams FROM "LiteLLM_UserTable" WHERE user_id = %s', (subject,))
+        assert (teams, user) == ([{"team_id": team_id}], [{"teams": [team_id]}])

tests/unit/proxy/management_endpoints/test_ui_sso.py

@@ -15,9 +15,18 @@ from fastapi import HTTPException, Request
 from pydantic import SecretStr
 
 import litellm
+from litellm._logging import verbose_proxy_logger
 from litellm._uuid import uuid
 from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
-from litellm.proxy._types import LiteLLM_JWTAuth, LiteLLM_UserTable, NewUserResponse, ProxyException
+from litellm.proxy._types import (
+    LiteLLM_JWTAuth,
+    LiteLLM_UserTable,
+    NewUserResponse,
+    ProxyException,
+    TeamMemberAddRequest,
+    UserAPIKeyAuth,
+)
+from litellm.proxy.auth.auth_checks import TeamNotFoundError
 from litellm.proxy.auth.handle_jwt import JWTHandler
 from litellm.proxy.management_endpoints.sso import CustomMicrosoftSSO
 from litellm.proxy.management_endpoints.types import CustomOpenID
@@ -11005,3 +11014,100 @@ def test_generic_sso_login_upserts_only_the_team_id_jwt_field_claim(
     )
 
     assert (result.team_ids, result.team_id_to_upsert) == (expected_team_ids, expected_upsert)
+
+
+def test_generic_sso_login_upserts_the_team_id_jwt_field_team_when_db_team_mappings_are_set() -> None:
+    from litellm.proxy.management_endpoints.ui_sso import generic_response_convertor
+
+    result: Final = generic_response_convertor(
+        response={"preferred_username": "alice", "email": "alice@example.com", "groups": ["g1"], "dept": "finance"},
+        jwt_handler=_upsert_jwt_handler(team_id_jwt_field="dept", team_id_upsert=True),
+        team_mappings=TeamMappings(team_ids_jwt_field="groups"),
+    )
+
+    assert (result.team_ids, result.team_id_to_upsert) == (["g1", "finance"], "finance")
+
+
+def _recording_team_calls(events: list[tuple[str, str]]) -> tuple[AsyncMock, Callable[[str], Awaitable[object]]]:
+    async def upsert_team(team_id: str) -> object:
+        events.append(("create", team_id))
+        return None
+
+    async def member_add(data: TeamMemberAddRequest, user_api_key_dict: UserAPIKeyAuth) -> None:
+        events.append(("join", data.team_id))
+
+    return AsyncMock(side_effect=member_add), upsert_team
+
+
+@pytest.mark.asyncio
+async def test_sso_login_creates_only_the_upsert_team_and_does_so_before_joining_it() -> None:
+    from litellm.proxy.management_endpoints.ui_sso import add_missing_team_member
+
+    events: Final[list[tuple[str, str]]] = []  # mutable-ok: records call order across awaited tasks
+    member_add, upsert_team = _recording_team_calls(events)
+    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, ["g1", "finance", "g2"], team_id_to_upsert="finance", upsert_team=upsert_team)
+
+    creates: Final = [team_id for kind, team_id in events if kind == "create"]
+    joins: Final = [team_id for kind, team_id in events if kind == "join"]
+    assert creates == ["finance"]
+    assert sorted(joins) == ["finance", "g1", "g2"]
+    assert events.index(("create", "finance")) < events.index(("join", "finance"))
+
+
+@pytest.mark.asyncio
+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
+
+    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)
+    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"]
+
+
+@pytest.mark.asyncio
+async def test_sso_login_logs_skipped_idp_groups_once_per_login_and_not_as_warnings(
+    caplog: pytest.LogCaptureFixture,
+) -> None:
+    from litellm.proxy.management_endpoints.ui_sso import add_missing_team_member
+
+    async def no_such_team(data: TeamMemberAddRequest, user_api_key_dict: UserAPIKeyAuth) -> None:
+        raise TeamNotFoundError(team_id=data.team_id)
+
+    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=no_such_team)),
+    ):
+        await add_missing_team_member(user, ["idp-group-a", "idp-group-b", "idp-group-c"])
+
+    skip_lines: Final = [r for r in caplog.records if "skipped" in r.getMessage()]
+    assert [r.levelno for r in skip_lines] == [logging.INFO]
+    assert all(team in skip_lines[0].getMessage() for team in ("idp-group-a", "idp-group-b", "idp-group-c"))
+    assert [r for r in caplog.records if r.levelno >= logging.WARNING] == []
+
+
+@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
+
+    user: Final = NewUserResponse(user_id="alice", key="sk-x", teams=None)
+    with (
+        caplog.at_level(logging.WARNING, logger=verbose_proxy_logger.name),
+        patch(
+            "litellm.proxy.management_endpoints.ui_sso.team_member_add",
+            AsyncMock(side_effect=HTTPException(status_code=500, detail="db down")),
+        ),
+    ):
+        await add_missing_team_member(user, ["finance"])
+
+    warnings: Final = [r.getMessage() for r in caplog.records if r.levelno == logging.WARNING]
+    assert len(warnings) == 1 and "finance" in warnings[0] and "db down" in warnings[0]