Merge 863ede9105 into d8f162ea22

litellm/proxy/auth/handle_jwt.py

@@ -14,7 +14,7 @@ import hashlib
 import os
 import re
 import time
-from collections.abc import Awaitable, Callable, Collection, Mapping, Sequence
+from collections.abc import Awaitable, Callable, Collection, Iterable, Mapping, Sequence
 from dataclasses import dataclass
 from typing import Final, Literal, NoReturn, Protocol, TypeVar, cast
 
@@ -1762,8 +1762,9 @@ class JWTAuthManager:
         parent_otel_span: Span | None,
         proxy_logging_obj: ProxyLogging,
         request_method: str | None = None,
+        team_id_upsert: bool | None = None,
     ) -> tuple[str | None, LiteLLM_TeamTable | None]:
-        """Find first team with access to the requested model"""
+        """Find first team with access to the requested model, creating claimed teams when ``team_id_upsert`` is set"""
         from litellm.proxy.proxy_server import llm_router
 
         denied_pass_through_team: LiteLLM_TeamTable | None = None
@@ -1789,6 +1790,7 @@ class JWTAuthManager:
                     user_api_key_cache=user_api_key_cache,
                     parent_otel_span=parent_otel_span,
                     proxy_logging_obj=proxy_logging_obj,
+                    team_id_upsert=team_id_upsert,
                 )
 
                 if team_object is not None:
@@ -2213,6 +2215,8 @@ class JWTAuthManager:
                 patch_team_membership,
             )
 
+            if jwt_handler.litellm_jwtauth.team_id_upsert and user_api_key_cache is not None:
+                await JWTAuthManager._upsert_claimed_teams(teams_to_add, prisma_client, user_api_key_cache)
             await patch_team_membership(
                 user_id=user_object.user_id,
                 teams_ids_to_add_user_to=list(teams_to_add),
@@ -2228,6 +2232,25 @@ class JWTAuthManager:
                 )
         return
 
+    @staticmethod
+    async def _upsert_claimed_teams(
+        team_ids: Iterable[str], prisma_client: PrismaClient, user_api_key_cache: UserApiKeyCache
+    ) -> None:
+        """Create the teams the user is about to join. A concurrent request may win the create, so the join still runs."""
+        for team_id in sorted(team_ids):
+            try:
+                await get_team_object(
+                    team_id=team_id,
+                    prisma_client=prisma_client,
+                    user_api_key_cache=user_api_key_cache,
+                    check_db_only=True,
+                    team_id_upsert=True,
+                )
+            except Exception as e:  # noqa: BLE001  # fail-safe: a lost create race or a create error must not fail auth, the join below still runs
+                verbose_proxy_logger.debug(
+                    "JWT auth could not create team %s, trying to join it anyway: %s", team_id, e
+                )
+
     @staticmethod
     async def _attach_team_from_header_for_admin(
         admin_result: JWTAuthBuilderResult,
@@ -2813,6 +2836,7 @@ class JWTAuthManager:
                 user_api_key_cache=user_api_key_cache,
                 parent_otel_span=parent_otel_span,
                 proxy_logging_obj=proxy_logging_obj,
+                team_id_upsert=team_id_upsert,
             )
 
         # The RBAC role-claim path (rbac_role == TEAM) sets team_id without

litellm/proxy/management_endpoints/types.py

@@ -93,5 +93,6 @@ def get_litellm_user_role(role_str: object) -> LitellmUserRoles | None:
 
 class CustomOpenID(OpenID):
     team_ids: list[str]
+    team_id_upsert: bool = False
     user_role: LitellmUserRoles | None = None
     extra_fields: dict[str, object] | None = None

litellm/proxy/management_endpoints/ui_sso.py

@@ -19,6 +19,8 @@ import secrets
 from collections import Counter
 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
 from typing import (
@@ -88,7 +90,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, get_team_object, 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,
@@ -1286,14 +1288,15 @@ def generic_response_convertor(
         generic_user_email_attribute_name,
     )
 
-    all_teams: Final = []
+    token: Final = cast(dict[str, object], response)  # cast-ok: the IdP userinfo JSON object fastapi-sso hands back
+    all_teams: Final = list[str]()
     if sso_jwt_handler is not None:
-        team_ids = sso_jwt_handler.get_all_jwt_team_ids(cast(dict, response))
+        team_ids = sso_jwt_handler.get_all_jwt_team_ids(token)
         all_teams.extend(team_ids)
 
     if team_mappings is not None and team_mappings.team_ids_jwt_field is not None:
         team_ids_from_db_mapping: Final[list[str] | None] = get_nested_value(
-            data=cast(dict, response),
+            data=token,
             key_path=team_mappings.team_ids_jwt_field,
             default=[],
         )
@@ -1305,8 +1308,21 @@ def generic_response_convertor(
                 team_ids_from_db_mapping,
             )
     else:
-        team_ids = jwt_handler.get_all_jwt_team_ids(cast(dict, response))
+        team_ids = jwt_handler.get_all_jwt_team_ids(token)
         all_teams.extend(team_ids)
+    team_id_jwt_field: Final = jwt_handler.litellm_jwtauth.team_id_jwt_field
+    team_id_upsert: Final = bool(jwt_handler.litellm_jwtauth.team_id_upsert)
+    scalar_team_id: Final[object] = (
+        get_nested_value(data=token, key_path=team_id_jwt_field)
+        if team_id_upsert and team_id_jwt_field is not None
+        else None
+    )
+    team_ids_with_scalar: Final = (
+        all_teams
+        if not isinstance(scalar_team_id, str) or scalar_team_id in all_teams
+        else [*all_teams, scalar_team_id]
+    )
+    team_ids_to_join: Final = [team_id for team_id in team_ids_with_scalar if team_id.strip()]
 
     # Determine user role based on role_mappings if available
     # Only apply role_mappings for GENERIC SSO provider
@@ -1374,7 +1390,8 @@ def generic_response_convertor(
         first_name=get_nested_value(response, generic_user_first_name_attribute_name),
         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_ids=team_ids_to_join,
+        team_id_upsert=team_id_upsert,
         user_role=user_role,
         extra_fields=extra_fields,
     )
@@ -1855,38 +1872,88 @@ 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)
 
 
-async def create_team_member_add_task(team_id, user_info):
-    """Create a task for adding a member to a team."""
+@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,
+) -> _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:
-        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.debug("[Non-Blocking] Error trying to add sso user to db: %s", e)
+        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
 
 
-async def add_missing_team_member(user_info: NewUserResponse | LiteLLM_UserTable, sso_teams: list[str]):
+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:
+    return await get_team_object(
+        team_id=team_id,
+        prisma_client=prisma_client,
+        user_api_key_cache=user_api_key_cache,
+        check_db_only=True,
+        team_id_upsert=True,
+    )
+
+
+async def add_missing_team_member(
+    user_info: NewUserResponse | LiteLLM_UserTable,
+    sso_teams: list[str],
+    upsert_team: Callable[[str], Awaitable[object]] | None = None,
+) -> None:
     """
     - Get missing teams (diff b/w user_info.team_ids and sso_teams)
+    - Create each missing team first when ``upsert_team`` is given (``team_id_upsert``)
     - 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) 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_info.teams or ()))
+    results: Final = await asyncio.gather(
+        *(create_team_member_add_task(team_id, user_info, upsert_team) for team_id in missing_teams)
+    )
+    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",
+            user_info.user_id,
+            len(skipped),
+            ", ".join(skipped),
+        )
+    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, failure.team_id, failure.error
+            )
 
 
 def get_disabled_non_admin_personal_key_creation():
@@ -1989,6 +2056,8 @@ async def get_user_info_from_db(
         await SSOAuthenticationHandler.add_user_to_teams_from_sso_response(
             result=result,
             user_info=user_info,
+            prisma_client=prisma_client,
+            user_api_key_cache=user_api_key_cache,
         )
 
         return user_info
@@ -3629,6 +3698,8 @@ class SSOAuthenticationHandler:
     async def add_user_to_teams_from_sso_response(
         result: CustomOpenID | OpenID | dict | None,
         user_info: NewUserResponse | LiteLLM_UserTable | None,
+        prisma_client: PrismaClient,
+        user_api_key_cache: UserApiKeyCache,
     ):
         """
         Adds the user as a team member to the teams specified in the SSO responses `team_ids` field
@@ -3640,7 +3711,14 @@ class SSOAuthenticationHandler:
             verbose_proxy_logger.debug("User not found in LiteLLM DB, skipping team member addition")
             return
         sso_teams: Final = result.team_ids if isinstance(result, CustomOpenID) else []
-        await add_missing_team_member(user_info=user_info, sso_teams=sso_teams)
+        upsert_teams: Final = isinstance(result, CustomOpenID) and result.team_id_upsert
+        await add_missing_team_member(
+            user_info=user_info,
+            sso_teams=sso_teams,
+            upsert_team=partial(_upsert_sso_team, prisma_client=prisma_client, user_api_key_cache=user_api_key_cache)
+            if upsert_teams
+            else None,
+        )
 
     @staticmethod
     def verify_user_in_restricted_sso_group(

tests/integration/authorization/test_jwt_team_claim_upsert.py

@@ -0,0 +1,86 @@
+import json
+import time
+import uuid
+from pathlib import Path
+from typing import Final
+
+import jwt
+import pytest
+import yaml
+from cryptography.hazmat.primitives.asymmetric import rsa
+
+from tests.integration._support.client import Gateway, eventually
+from tests.integration._support.database import read_rows
+from tests.integration._support.process import owned_proxy
+from tests.integration._support.wire import Reply, Request, wire_server
+
+pytestmark: Final = pytest.mark.shares_worker_state
+
+KEY_ID: Final = "integration-jwt-team-claim-signing-key"
+
+
+def _jwks_reply(public_jwk: str) -> Reply:
+    return Reply(body=json.dumps({"keys": [{**json.loads(public_jwk), "kid": KEY_ID}]}).encode())
+
+
+@pytest.mark.parametrize(
+    "team_id_upsert", [True, False], ids=["upsert-creates-claimed-teams", "no-upsert-rejects-them"]
+)
+def test_jwt_teams_array_claim_creates_and_joins_every_team_only_when_team_id_upsert_is_on(
+    gateway: Gateway, tmp_path: Path, team_id_upsert: bool
+) -> None:
+    private_key: Final = rsa.generate_private_key(public_exponent=65537, key_size=2048)
+    public_jwk: Final[str] = jwt.algorithms.RSAAlgorithm.to_jwk(private_key.public_key())
+
+    def respond(request: Request) -> Reply:
+        assert request.method == "GET", request
+        return _jwks_reply(public_jwk)
+
+    with wire_server(respond) as jwks, gateway.scenario() as scenario:
+        config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
+        config["general_settings"] = {
+            **config["general_settings"],
+            "enable_jwt_auth": True,
+            "litellm_jwtauth": {
+                "user_id_jwt_field": "sub",
+                "user_id_upsert": True,
+                "team_ids_jwt_field": "teams",
+                "team_id_upsert": team_id_upsert,
+                "sync_user_role_and_teams": True,
+            },
+        }
+        path: Final = tmp_path / "jwt_team_claim_upsert.yaml"
+        path.write_text(yaml.safe_dump(config))
+        subject: Final = f"integration-jwt-{uuid.uuid4().hex}"
+        claimed_teams: Final = sorted(f"jwt-team-{subject}-{suffix}" for suffix in ("a", "b"))
+        token: Final = jwt.encode(
+            {"sub": subject, "teams": claimed_teams, "iat": int(time.time()), "exp": int(time.time()) + 300},
+            private_key,
+            algorithm="RS256",
+            headers={"kid": KEY_ID},
+        )
+        with owned_proxy(gateway, tmp_path, {"JWT_PUBLIC_KEY_URL": jwks.url}, config=path) as candidate:
+            model: Final = scenario.model()
+            scenario.cleanups.callback(gateway.request, "POST", "/team/delete", {"team_ids": claimed_teams})
+            response: Final = candidate.request(
+                "POST",
+                "/v1/chat/completions",
+                {"model": model, "messages": [{"role": "user", "content": "teams array claim"}]},
+                key=token,
+            )
+            team_rows: Final = read_rows(
+                'SELECT team_id FROM "LiteLLM_TeamTable" WHERE team_id = ANY(%s) ORDER BY team_id', (claimed_teams,)
+            )
+            if not team_id_upsert:
+                assert (response.status_code, team_rows) == (403, []), response.text
+                return
+            scenario.cleanups.callback(scenario.delete_user, subject)
+            assert response.status_code == 200, response.text
+            user_teams: Final = eventually(
+                lambda: read_rows('SELECT teams FROM "LiteLLM_UserTable" WHERE user_id = %s', (subject,)),
+                lambda rows: [sorted(row["teams"]) for row in rows] == [claimed_teams],
+            )
+            assert ([sorted(row["teams"]) for row in user_teams], [row["team_id"] for row in team_rows]) == (
+                [claimed_teams],
+                claimed_teams,
+            )

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
 
@@ -52,14 +52,29 @@ def _idp_reply(request: Request) -> Reply:
                     "preferred_username": signed_in,
                     "email": f"{signed_in}@example.com",
                     "provider": IDP_PROVIDER,
+                    "dept": _dept(signed_in),
+                    "groups": [_group(signed_in)],
                 }
             ).encode()
         )
     return Reply(status=404, body=b'{"error": "not_found"}')
 
 
+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]) -> Iterator[Gateway]:
+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(
         yaml.safe_dump(
@@ -68,6 +83,7 @@ def _sso_gateway(directory: Path, litellm_settings: Mapping[str, JsonValue]) ->
                 "general_settings": {
                     "master_key": "os.environ/LITELLM_MASTER_KEY",
                     "database_url": "os.environ/DATABASE_URL",
+                    **(general_settings or {}),
                 },
                 "litellm_settings": dict(litellm_settings),
             }
@@ -83,7 +99,33 @@ def _sso_gateway(directory: Path, litellm_settings: Mapping[str, JsonValue]) ->
             "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:
@@ -165,3 +207,57 @@ def test_free_sso_admission_preserves_existing_users_and_serializes_the_last_sea
         refused: Final = _sign_in(gateway, "free-sso-extra")
         assert (above_capacity.status_code, refused.status_code) == (303, 403)
         assert gateway.rows('SELECT COUNT(*) AS count FROM "LiteLLM_UserTable"', ()) == [{"count": 6}]
+
+
+@pytest.mark.parametrize(
+    ("team_id_upsert", "team_exists_before", "expect_member"),
+    [(True, False, True), (False, False, False), (False, True, True)],
+    ids=["upsert-creates-missing-team", "no-upsert-skips-missing-team", "no-upsert-joins-existing-team"],
+)
+def test_sso_sign_in_upserts_the_team_id_jwt_field_team_only_when_team_id_upsert_is_on(
+    tmp_path: Path, team_id_upsert: bool, team_exists_before: bool, expect_member: bool
+) -> None:
+    subject: Final = f"sso-dept-user-{uuid.uuid4().hex[:12]}"
+    team_id: Final = _dept(subject)
+    jwtauth: Final = {"litellm_jwtauth": {"team_id_jwt_field": "dept", "team_id_upsert": team_id_upsert}}
+    with _sso_gateway(tmp_path, {}, jwtauth) as gateway, contextlib.ExitStack() as cleanups:
+        cleanups.callback(gateway.request, "POST", "/team/delete", {"team_ids": [team_id]})
+        cleanups.callback(_delete_user, gateway, subject)
+        if team_exists_before:
+            created: Final = gateway.request("POST", "/team/new", {"team_id": team_id})
+            assert created.status_code == 200, created.text
+
+        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 = %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_the_group_teams(
+    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 (sorted(row["team_id"] for row in teams), user) == (
+            sorted([team_id, group_team_id]),
+            [{"teams": sorted([team_id, group_team_id])}],
+        )

tests/unit/proxy/auth/test_handle_jwt.py

@@ -14,14 +14,14 @@ from fastapi import HTTPException
 from litellm.caching.dual_cache import DualCache
 from litellm.proxy._types import (
     DEFAULT_JWKS_STALE_TTL,
-    JWTLiteLLMRoleMap,
     JWTIssuerConfig,
-    LiteLLMRoutes,
+    JWTLiteLLMRoleMap,
     LiteLLM_JWTAuth,
     LiteLLM_ModelTable,
     LiteLLM_TeamMembership,
     LiteLLM_TeamTable,
     LiteLLM_UserTable,
+    LiteLLMRoutes,
     LitellmUserRoles,
     Member,
     ProxyErrorTypes,
@@ -8637,3 +8637,107 @@ async def test_decode_trusted_assertion_runs_the_configured_custom_validate_on_t
     else:
         assert isinstance(verified, AssertionRejection)
         assert "custom_validate" in verified.reason
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(("team_id_upsert", "expected_created"), [(True, ["team1", "team3"]), (False, [])])
+async def test_sync_user_role_and_teams_creates_the_claimed_array_teams_only_when_team_id_upsert_is_on(
+    team_id_upsert: bool, expected_created: list[str]
+) -> None:
+    jwt_handler = JWTHandler()
+    jwt_handler.update_environment(
+        prisma_client=None,
+        user_api_key_cache=AsyncMock(),
+        litellm_jwtauth=LiteLLM_JWTAuth(
+            team_ids_jwt_field="my_id_teams", sync_user_role_and_teams=True, team_id_upsert=team_id_upsert
+        ),
+    )
+    token = {"my_id_teams": ["team3", "team1", "team2"]}
+    user = LiteLLM_UserTable(user_id="u1", user_role=LitellmUserRoles.INTERNAL_USER.value, teams=["team2"])
+    created: list[str] = []  # mutable-ok: records the teams the proxy would create
+
+    async def create_team(team_id: str, **kwargs: object) -> LiteLLM_TeamTable:
+        created.append(team_id)
+        return LiteLLM_TeamTable(team_id=team_id)
+
+    with (
+        patch("litellm.proxy.auth.handle_jwt.get_team_object", create_team),
+        patch(
+            "litellm.proxy.management_endpoints.scim.scim_v2.patch_team_membership", new_callable=AsyncMock
+        ) as membership,
+    ):
+        await JWTAuthManager.sync_user_role_and_teams(
+            jwt_handler, token, user, AsyncMock(), user_api_key_cache=AsyncMock()
+        )
+
+    assert created == expected_created
+    assert sorted(membership.call_args.kwargs["teams_ids_to_add_user_to"]) == ["team1", "team3"]
+
+
+@pytest.mark.asyncio
+async def test_sync_user_role_and_teams_still_joins_the_team_when_a_concurrent_request_created_it_first() -> None:
+    jwt_handler = JWTHandler()
+    jwt_handler.update_environment(
+        prisma_client=None,
+        user_api_key_cache=AsyncMock(),
+        litellm_jwtauth=LiteLLM_JWTAuth(
+            team_ids_jwt_field="my_id_teams", sync_user_role_and_teams=True, team_id_upsert=True
+        ),
+    )
+    user = LiteLLM_UserTable(user_id="u1", user_role=LitellmUserRoles.INTERNAL_USER.value, teams=[])
+
+    async def lost_the_create_race(team_id: str, **kwargs: object) -> LiteLLM_TeamTable:
+        raise HTTPException(status_code=404, detail={"error": f"Team id = {team_id} already exists"})
+
+    with (
+        patch("litellm.proxy.auth.handle_jwt.get_team_object", lost_the_create_race),
+        patch(
+            "litellm.proxy.management_endpoints.scim.scim_v2.patch_team_membership", new_callable=AsyncMock
+        ) as membership,
+    ):
+        await JWTAuthManager.sync_user_role_and_teams(
+            jwt_handler, {"my_id_teams": ["team1"]}, user, AsyncMock(), user_api_key_cache=AsyncMock()
+        )
+
+    assert membership.call_args.kwargs["teams_ids_to_add_user_to"] == ["team1"]
+    assert user.teams == ["team1"]
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(("team_id_upsert", "expected_team_id"), [(True, "team-1"), (False, None)])
+async def test_find_team_with_model_access_creates_the_claimed_team_only_when_team_id_upsert_is_on(
+    monkeypatch: pytest.MonkeyPatch, team_id_upsert: bool, expected_team_id: str | None
+) -> None:
+    import sys
+    import types
+
+    from litellm.caching import DualCache
+    from litellm.proxy.utils import ProxyLogging
+
+    proxy_server_module = types.ModuleType("proxy_server")
+    proxy_server_module.llm_router = None
+    monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", proxy_server_module)
+
+    async def team_table_with_upsert(team_id: str, **kwargs: object) -> LiteLLM_TeamTable:
+        if not kwargs.get("team_id_upsert"):
+            raise HTTPException(status_code=404, detail={"error": f"Team doesn't exist in db. Team={team_id}"})
+        return LiteLLM_TeamTable(team_id=team_id)
+
+    monkeypatch.setattr("litellm.proxy.auth.handle_jwt.get_team_object", team_table_with_upsert)
+    jwt_handler = JWTHandler()
+    jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(team_id_upsert=team_id_upsert)
+    user_api_key_cache = DualCache()
+
+    team_id, _ = await JWTAuthManager.find_team_with_model_access(
+        team_ids={"team-1"},
+        requested_model=None,
+        route="/chat/completions",
+        jwt_handler=jwt_handler,
+        prisma_client=None,
+        user_api_key_cache=user_api_key_cache,
+        parent_otel_span=None,
+        proxy_logging_obj=ProxyLogging(user_api_key_cache=user_api_key_cache),
+        team_id_upsert=team_id_upsert,
+    )
+
+    assert team_id == expected_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_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
@@ -287,6 +296,13 @@ def test_get_microsoft_callback_response(stubbed_graph_api):
     assert sso_assertion.issuer == "https://idp.example.com"
 
 
+
+def _jwt_handler_double() -> MagicMock:
+    handler: Final = MagicMock(spec=JWTHandler)
+    handler.litellm_jwtauth = LiteLLM_JWTAuth()
+    return handler
+
+
 GRAPH_ERROR_BODIES: Final = (
     {"error": {"code": "Authentication_MissingOrMalformed", "message": "Access Token missing or malformed."}},
     {"error": {"code": "InvalidAuthenticationToken", "message": "IDX14100: JWT is not well formed."}},
@@ -1158,7 +1174,7 @@ def test_generic_response_convertor_normalizes_email():
     }
 
     # Mock JWT handler
-    mock_jwt_handler = MagicMock(spec=JWTHandler)
+    mock_jwt_handler = _jwt_handler_double()
     mock_jwt_handler.get_team_ids_from_jwt.return_value = []
 
     result = generic_response_convertor(
@@ -2226,7 +2242,7 @@ async def test_get_generic_sso_response_with_additional_headers():
 
     # Arrange
     mock_request = MagicMock(spec=Request)
-    mock_jwt_handler = MagicMock(spec=JWTHandler)
+    mock_jwt_handler = _jwt_handler_double()
     mock_jwt_handler.get_team_ids_from_jwt.return_value = []
 
     generic_client_id = "test_client_id"
@@ -2299,7 +2315,7 @@ async def test_get_generic_sso_response_with_empty_headers():
 
     # Arrange
     mock_request = MagicMock(spec=Request)
-    mock_jwt_handler = MagicMock(spec=JWTHandler)
+    mock_jwt_handler = _jwt_handler_double()
     mock_jwt_handler.get_team_ids_from_jwt.return_value = []
 
     generic_client_id = "test_client_id"
@@ -2359,8 +2375,8 @@ async def test_get_generic_sso_response_includes_token_claims_when_enabled(monke
     from litellm.proxy.management_endpoints.ui_sso import get_generic_sso_response
 
     mock_request = MagicMock(spec=Request)
-    mock_jwt_handler = MagicMock(spec=JWTHandler)
-    mock_sso_jwt_handler = MagicMock(spec=JWTHandler)
+    mock_jwt_handler = _jwt_handler_double()
+    mock_sso_jwt_handler = _jwt_handler_double()
     mock_sso_jwt_handler.get_all_jwt_team_ids.return_value = ["team-from-userinfo"]
     mock_sso_jwt_handler.get_team_ids_from_jwt.return_value = []
 
@@ -2435,8 +2451,8 @@ async def test_get_generic_sso_response_does_not_include_token_claims_when_disab
     from litellm.proxy.management_endpoints.ui_sso import get_generic_sso_response
 
     mock_request = MagicMock(spec=Request)
-    mock_jwt_handler = MagicMock(spec=JWTHandler)
-    mock_sso_jwt_handler = MagicMock(spec=JWTHandler)
+    mock_jwt_handler = _jwt_handler_double()
+    mock_sso_jwt_handler = _jwt_handler_double()
     mock_sso_jwt_handler.get_all_jwt_team_ids.return_value = ["team-from-userinfo"]
     mock_sso_jwt_handler.get_team_ids_from_jwt.return_value = []
     access_token = pyjwt.encode({"upn": "token-user@example.com"}, "test-secret", algorithm="HS256")
@@ -2578,7 +2594,7 @@ async def test_get_generic_sso_response_pkce_merges_token_claims_and_excludes_cr
     ):
         result, received_response, _, _ = await get_generic_sso_response(
             request=request,
-            jwt_handler=MagicMock(spec=JWTHandler),
+            jwt_handler=_jwt_handler_double(),
             generic_client_id="test-client",
             redirect_url="http://test.com/callback",
             sso_jwt_handler=None,
@@ -2600,7 +2616,7 @@ async def test_get_generic_sso_response_ignores_opaque_and_empty_token_claims(mo
     from litellm.proxy.management_endpoints.ui_sso import get_generic_sso_response
 
     mock_request = MagicMock(spec=Request)
-    mock_jwt_handler = MagicMock(spec=JWTHandler)
+    mock_jwt_handler = _jwt_handler_double()
     userinfo = {
         "preferred_username": "userinfo-user",
         "email": "userinfo@example.com",
@@ -4928,7 +4944,7 @@ class TestProcessSSOJWTAccessToken:
     @pytest.fixture
     def mock_jwt_handler(self):
         """Create a mock JWT handler for testing"""
-        mock_handler = MagicMock(spec=JWTHandler)
+        mock_handler = _jwt_handler_double()
         mock_handler.get_team_ids_from_jwt.return_value = ["team1", "team2", "team3"]
         return mock_handler
 
@@ -5315,7 +5331,7 @@ class TestGenericResponseConvertorNestedAttributes:
         from litellm.proxy.management_endpoints.ui_sso import generic_response_convertor
 
         # Mock JWT handler
-        mock_jwt_handler = MagicMock(spec=JWTHandler)
+        mock_jwt_handler = _jwt_handler_double()
         mock_jwt_handler.get_team_ids_from_jwt.return_value = []
 
         # Payload with nested attributes structure
@@ -5386,7 +5402,7 @@ class TestGenericResponseConvertorUserRole:
         from litellm.proxy._types import LitellmUserRoles
         from litellm.proxy.management_endpoints.ui_sso import generic_response_convertor
 
-        mock_jwt_handler = MagicMock(spec=JWTHandler)
+        mock_jwt_handler = _jwt_handler_double()
         mock_jwt_handler.get_team_ids_from_jwt.return_value = []
 
         sso_response = {
@@ -5416,7 +5432,7 @@ class TestGenericResponseConvertorUserRole:
         """
         from litellm.proxy.management_endpoints.ui_sso import generic_response_convertor
 
-        mock_jwt_handler = MagicMock(spec=JWTHandler)
+        mock_jwt_handler = _jwt_handler_double()
         mock_jwt_handler.get_team_ids_from_jwt.return_value = []
 
         sso_response = {
@@ -6715,7 +6731,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)
 
@@ -6752,7 +6768,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)
 
@@ -6791,11 +6807,15 @@ class TestAddMissingTeamMember:
             await SSOAuthenticationHandler.add_user_to_teams_from_sso_response(
                 result=sso_result,
                 user_info=new_user_info,
+                prisma_client=MagicMock(),
+                user_api_key_cache=MagicMock(),
             )
 
             # Verify add_missing_team_member was called with correct args
             mock_add_member.assert_called_once_with(
-                user_info=new_user_info, sso_teams=["entra-group-1", "entra-group-2"]
+                user_info=new_user_info,
+                sso_teams=["entra-group-1", "entra-group-2"],
+                upsert_team=None,
             )
 
     @pytest.mark.asyncio
@@ -6809,7 +6829,7 @@ class TestAddMissingTeamMember:
 
         team_member_calls = []
 
-        async def track_team_member_add(team_id, user_info):
+        async def track_team_member_add(team_id, user_info, upsert_team=None):
             team_member_calls.append({"team_id": team_id, "user_id": user_info.user_id})
 
         # New SSO user with Entra groups
@@ -6878,7 +6898,7 @@ class TestAddMissingTeamMember:
 
         added_teams = []
 
-        async def mock_create_task(team_id, user):
+        async def mock_create_task(team_id, user, upsert_team=None):
             added_teams.append(team_id)
 
         with patch(
@@ -7815,7 +7835,7 @@ def test_generic_response_convertor_with_extra_attributes(monkeypatch):
         "custom_field3": {"nested": "data"},
     }
 
-    mock_jwt_handler = MagicMock(spec=JWTHandler)
+    mock_jwt_handler = _jwt_handler_double()
     mock_jwt_handler.get_team_ids_from_jwt.return_value = []
 
     result = generic_response_convertor(
@@ -7849,7 +7869,7 @@ def test_generic_response_convertor_without_extra_attributes(monkeypatch):
         "custom_field2": "value2",
     }
 
-    mock_jwt_handler = MagicMock(spec=JWTHandler)
+    mock_jwt_handler = _jwt_handler_double()
     mock_jwt_handler.get_team_ids_from_jwt.return_value = []
 
     result = generic_response_convertor(
@@ -7877,7 +7897,7 @@ def test_generic_response_convertor_extra_attributes_with_nested_paths(monkeypat
         "org_info": {"department": "Engineering", "manager": "Jane Smith"},
     }
 
-    mock_jwt_handler = MagicMock(spec=JWTHandler)
+    mock_jwt_handler = _jwt_handler_double()
     mock_jwt_handler.get_team_ids_from_jwt.return_value = []
 
     result = generic_response_convertor(
@@ -7904,7 +7924,7 @@ def test_generic_response_convertor_extra_attributes_missing_field(monkeypatch):
         "email": "user@example.com",
     }
 
-    mock_jwt_handler = MagicMock(spec=JWTHandler)
+    mock_jwt_handler = _jwt_handler_double()
     mock_jwt_handler.get_team_ids_from_jwt.return_value = []
 
     result = generic_response_convertor(
@@ -8783,7 +8803,7 @@ class TestPKCEStateCookieBinding:
         ):
             await get_generic_sso_response(
                 request=mock_request,
-                jwt_handler=MagicMock(spec=JWTHandler),
+                jwt_handler=_jwt_handler_double(),
                 generic_client_id="cid",
                 redirect_url="https://proxy.example.com/sso/callback",
                 sso_jwt_handler=None,
@@ -8835,7 +8855,7 @@ class TestPKCEStateCookieBinding:
         ):
             await get_generic_sso_response(
                 request=mock_request,
-                jwt_handler=MagicMock(spec=JWTHandler),
+                jwt_handler=_jwt_handler_double(),
                 generic_client_id="cid",
                 redirect_url="https://proxy.example.com/sso/callback",
                 sso_jwt_handler=None,
@@ -8898,7 +8918,7 @@ class TestPKCEStateCookieBinding:
                 },
             ),
         ):
-            jwt_handler = MagicMock(spec=JWTHandler)
+            jwt_handler = _jwt_handler_double()
             jwt_handler.get_team_ids_from_jwt.return_value = []
             result, _, _, _ = await get_generic_sso_response(
                 request=mock_request,
@@ -9005,7 +9025,7 @@ async def test_debug_sso_callback_renders_full_jwt_claims():
         patch("litellm.proxy.proxy_server.general_settings", {}),
         patch("litellm.proxy.proxy_server.prisma_client", MagicMock()),
         patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()),
-        patch("litellm.proxy.proxy_server.jwt_handler", MagicMock(spec=JWTHandler)),
+        patch("litellm.proxy.proxy_server.jwt_handler", _jwt_handler_double()),
     ):
         # Microsoft / Google envs may leak in from other tests — ensure only
         # the generic path runs.
@@ -9091,7 +9111,7 @@ async def test_debug_sso_callback_handles_missing_raw_response():
         patch("litellm.proxy.proxy_server.general_settings", {}),
         patch("litellm.proxy.proxy_server.prisma_client", MagicMock()),
         patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()),
-        patch("litellm.proxy.proxy_server.jwt_handler", MagicMock(spec=JWTHandler)),
+        patch("litellm.proxy.proxy_server.jwt_handler", _jwt_handler_double()),
     ):
         for var in ("GENERIC_CLIENT_ID", "GOOGLE_CLIENT_ID"):
             os.environ.pop(var, None)
@@ -9142,7 +9162,7 @@ async def _render_debug_page(provider_env, id_jag_registered, force_inert=False)
         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.jwt_handler", _jwt_handler_double()),  # test-quality-ok: debug endpoint reads proxy globals
     ]
     if force_inert:
         stack.append(
@@ -9507,7 +9527,7 @@ async def test_pkce_arm_captures_sso_assertion():
             },
         ),
     ):
-        jwt_handler = MagicMock(spec=JWTHandler)
+        jwt_handler = _jwt_handler_double()
         jwt_handler.get_team_ids_from_jwt.return_value = []
         result, received_response, _, sso_assertion = await get_generic_sso_response(
             request=mock_request,
@@ -9533,7 +9553,7 @@ async def test_verify_and_process_arm_captures_sso_assertion():
 
     id_token = _ema_id_token()
     mock_request = MagicMock(spec=Request)
-    mock_jwt_handler = MagicMock(spec=JWTHandler)
+    mock_jwt_handler = _jwt_handler_double()
     mock_jwt_handler.get_team_ids_from_jwt.return_value = []
 
     mock_sso_instance = MagicMock()
@@ -9598,7 +9618,7 @@ def _generic_sso_callback_request() -> Request:
 
 
 def _generic_sso_jwt_handler() -> MagicMock:
-    jwt_handler = MagicMock(spec=JWTHandler)
+    jwt_handler = _jwt_handler_double()
     jwt_handler.get_team_ids_from_jwt.return_value = []
     return jwt_handler
 
@@ -11062,3 +11082,223 @@ async def test_cli_sso_login_never_serves_the_admin_login_form_on_a_ui_disabled_
     body = response.body.decode()
     assert "Admin UI is Disabled" in body
     assert 'name="username"' not in body
+
+
+def _upsert_jwt_handler(**jwtauth: object) -> JWTHandler:
+    from litellm.caching.caching import DualCache
+    from litellm.proxy._types import LiteLLM_JWTAuth
+
+    handler: Final = JWTHandler()
+    handler.update_environment(
+        prisma_client=None, user_api_key_cache=DualCache(), litellm_jwtauth=LiteLLM_JWTAuth(**jwtauth)
+    )
+    return handler
+
+
+@pytest.mark.parametrize(
+    ("jwtauth", "userinfo", "expected_team_ids", "expected_upsert"),
+    [
+        ({"team_id_jwt_field": "dept", "team_id_upsert": True}, {"dept": "finance"}, ["finance"], True),
+        ({"team_id_jwt_field": "dept", "team_id_upsert": False}, {"dept": "finance"}, ["finance"], False),
+        ({"team_id_jwt_field": "dept", "team_id_upsert": True}, {}, [], True),
+        ({"team_id_jwt_field": "dept", "team_id_upsert": True}, {"dept": ""}, [], True),
+        ({"team_id_jwt_field": "dept", "team_id_upsert": True}, {"dept": "  "}, [], True),
+        ({"team_id_jwt_field": "dept", "team_id_upsert": True}, {"dept": ["a", "b"]}, ["a", "b"], True),
+        (
+            {"team_id_jwt_field": "dept", "team_id_upsert": True, "team_id_default": "everyone"},
+            {},
+            [],
+            True,
+        ),
+        (
+            {"team_ids_jwt_field": "groups", "team_id_upsert": True},
+            {"groups": ["g1", "g2"]},
+            ["g1", "g2"],
+            True,
+        ),
+        (
+            {"team_ids_jwt_field": "groups", "team_id_upsert": False},
+            {"groups": ["g1", "g2"]},
+            ["g1", "g2"],
+            False,
+        ),
+        (
+            {"team_ids_jwt_field": "groups", "team_id_jwt_field": "dept", "team_id_upsert": True},
+            {"groups": ["g1"], "dept": "finance"},
+            ["g1", "finance"],
+            True,
+        ),
+    ],
+    ids=[
+        "upsert-on",
+        "upsert-off",
+        "claim-missing",
+        "blank-claim-dropped",
+        "whitespace-claim-dropped",
+        "list-valued-claim",
+        "default-team-not-joined",
+        "plural-claim-upsert-on",
+        "plural-claim-upsert-off",
+        "mixed",
+    ],
+)
+def test_generic_sso_login_marks_every_claimed_team_for_upsert_only_when_team_id_upsert_is_on(
+    jwtauth: dict[str, object],
+    userinfo: dict[str, object],
+    expected_team_ids: list[str],
+    expected_upsert: bool,
+) -> 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", **userinfo},
+        jwt_handler=_upsert_jwt_handler(**jwtauth),
+    )
+
+    assert (result.team_ids, result.team_id_upsert) == (expected_team_ids, expected_upsert)
+
+
+def test_generic_sso_login_keeps_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_upsert) == (["g1", "finance"], True)
+
+
+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_every_missing_team_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=["g0"])
+
+    with patch("litellm.proxy.management_endpoints.ui_sso.team_member_add", member_add):
+        await add_missing_team_member(user, ["g0", "g1", "finance", "g2"], 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 sorted(creates) == ["finance", "g1", "g2"]
+    assert sorted(joins) == ["finance", "g1", "g2"]
+    assert all(events.index(("create", team_id)) < events.index(("join", team_id)) for team_id in creates)
+
+
+@pytest.mark.asyncio
+async def test_sso_login_creates_no_team_when_team_id_upsert_is_off() -> 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, _ = _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"])
+
+    assert events == [("join", "finance"), ("join", "g1")]
+
+
+@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
+
+    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"})
+
+    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"], upsert_team=lost_the_create_race)
+
+    assert joined == ["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
+@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"], 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
+
+    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]