From 66acab51c10150858003678db63a875f274aa219 Mon Sep 17 00:00:00 2001 From: Jesus Velarde Date: Wed, 28 Oct 2026 16:21:31 -0700 Subject: [PATCH 1/9] feat(sso): create the team_id_jwt_field team on UI SSO login when team_id_upsert is set A UI SSO login read team_id_jwt_field but only joined teams that already existed, so with litellm_jwtauth.team_id_upsert set the team was only created on the first API call. The login now creates the single string team_id_jwt_field team the same way the API path does, while team_ids_jwt_field lists and SSO group claims stay join-only. A team the login cannot join is now logged at warning instead of debug --- litellm/proxy/management_endpoints/types.py | 1 + litellm/proxy/management_endpoints/ui_sso.py | 62 ++++++-- .../management/test_sso_sign_in.py | 38 ++++- .../proxy/management_endpoints/test_ui_sso.py | 137 ++++++++++++++---- 4 files changed, 197 insertions(+), 41 deletions(-) diff --git a/litellm/proxy/management_endpoints/types.py b/litellm/proxy/management_endpoints/types.py index 221a6221fab..d4fd7a31d13 100644 --- a/litellm/proxy/management_endpoints/types.py +++ b/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_to_upsert: str | None = None user_role: LitellmUserRoles | None = None extra_fields: dict[str, object] | None = None diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index 22c4f27d94a..ecd6994a2c1 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 functools import partial from html import escape from types import MappingProxyType, TracebackType from typing import ( @@ -87,7 +88,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, @@ -1238,14 +1239,15 @@ def generic_response_convertor( generic_user_email_attribute_name, ) + token: Final = cast(dict[str, object], response) # cast-ok: the IdP userinfo JSON object fastapi-sso hands back all_teams: Final = [] 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=[], ) @@ -1257,8 +1259,14 @@ 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 + upsert_candidate: Final[object] = ( + get_nested_value(data=token, key_path=team_id_jwt_field) + if jwt_handler.litellm_jwtauth.team_id_upsert and team_id_jwt_field is not None + else None + ) # Determine user role based on role_mappings if available # Only apply role_mappings for GENERIC SSO provider @@ -1327,6 +1335,9 @@ 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, user_role=user_role, extra_fields=extra_fields, ) @@ -1807,9 +1818,15 @@ 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): +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.""" 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, @@ -1820,10 +1837,25 @@ async def create_team_member_add_task(team_id, user_info): 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) + verbose_proxy_logger.warning("SSO login could not add user %s to team %s: %s", user_info.user_id, team_id, e) -async def add_missing_team_member(user_info: NewUserResponse | LiteLLM_UserTable, sso_teams: list[str]): +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], + team_id_to_upsert: str | None = None, + upsert_team: Callable[[str], Awaitable[object]] | None = None, +): """ - Get missing teams (diff b/w user_info.team_ids and sso_teams) - Add missing user to missing teams @@ -1833,7 +1865,10 @@ async def add_missing_team_member(user_info: NewUserResponse | LiteLLM_UserTable 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] + 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) @@ -1941,6 +1976,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 @@ -3567,6 +3604,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 @@ -3578,7 +3617,12 @@ 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) + await add_missing_team_member( + user_info=user_info, + sso_teams=sso_teams, + team_id_to_upsert=result.team_id_to_upsert if isinstance(result, CustomOpenID) else None, + upsert_team=partial(_upsert_sso_team, prisma_client=prisma_client, user_api_key_cache=user_api_key_cache), + ) @staticmethod def verify_user_in_restricted_sso_group( diff --git a/tests/integration/management/test_sso_sign_in.py b/tests/integration/management/test_sso_sign_in.py index d0b6e2579c8..1af3e5443dc 100644 --- a/tests/integration/management/test_sso_sign_in.py +++ b/tests/integration/management/test_sso_sign_in.py @@ -52,14 +52,23 @@ def _idp_reply(request: Request) -> Reply: "preferred_username": signed_in, "email": f"{signed_in}@example.com", "provider": IDP_PROVIDER, + "dept": _dept(signed_in), } ).encode() ) return Reply(status=404, body=b'{"error": "not_found"}') +def _dept(subject: str) -> str: + return f"dept-{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, +) -> Iterator[Gateway]: config: Final = directory / "sso_proxy_config.yaml" config.write_text( yaml.safe_dump( @@ -68,6 +77,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), } @@ -165,3 +175,29 @@ 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 []}]) diff --git a/tests/unit/proxy/management_endpoints/test_ui_sso.py b/tests/unit/proxy/management_endpoints/test_ui_sso.py index 4135e9be6f6..2ae67472149 100644 --- a/tests/unit/proxy/management_endpoints/test_ui_sso.py +++ b/tests/unit/proxy/management_endpoints/test_ui_sso.py @@ -6,7 +6,7 @@ from collections.abc import Awaitable, Callable, Mapping, Sequence from contextlib import ExitStack, asynccontextmanager from types import SimpleNamespace from typing import Final -from unittest.mock import AsyncMock, MagicMock, patch +from unittest.mock import ANY, AsyncMock, MagicMock, patch import httpx import pytest @@ -17,7 +17,7 @@ from pydantic import SecretStr import litellm 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 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 +287,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 +1165,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 +2233,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 +2306,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 +2366,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 +2442,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 +2585,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 +2607,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", @@ -4796,7 +4803,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 @@ -5183,7 +5190,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 @@ -5254,7 +5261,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 = { @@ -5284,7 +5291,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 = { @@ -6659,11 +6666,16 @@ 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"], + team_id_to_upsert=None, + upsert_team=ANY, ) @pytest.mark.asyncio @@ -6677,7 +6689,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 @@ -6746,7 +6758,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( @@ -7683,7 +7695,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( @@ -7717,7 +7729,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( @@ -7745,7 +7757,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( @@ -7772,7 +7784,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( @@ -8651,7 +8663,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, @@ -8703,7 +8715,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, @@ -8766,7 +8778,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, @@ -8873,7 +8885,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. @@ -8959,7 +8971,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) @@ -9010,7 +9022,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( @@ -9375,7 +9387,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, @@ -9401,7 +9413,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() @@ -9466,7 +9478,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 @@ -10930,3 +10942,66 @@ 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"], "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": ["a", "b"]}, ["a", "b"], None), + ( + {"team_id_jwt_field": "dept", "team_id_upsert": True, "team_id_default": "everyone"}, + {}, + [], + None, + ), + ( + {"team_ids_jwt_field": "groups", "team_id_upsert": True}, + {"groups": ["g1", "g2"]}, + ["g1", "g2"], + None, + ), + ( + {"team_ids_jwt_field": "groups", "team_id_jwt_field": "dept", "team_id_upsert": True}, + {"groups": ["g1"], "dept": "finance"}, + ["g1", "finance"], + "finance", + ), + ], + ids=[ + "upsert-on", + "upsert-off", + "claim-missing", + "list-valued-claim-join-only", + "default-team-not-upserted", + "plural-claim-join-only", + "mixed", + ], +) +def test_generic_sso_login_upserts_only_the_team_id_jwt_field_claim( + jwtauth: dict[str, object], + userinfo: dict[str, object], + expected_team_ids: list[str], + expected_upsert: str | None, +) -> 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_to_upsert) == (expected_team_ids, expected_upsert) From 4887f13035328a397bf139b54e53e9344623c039 Mon Sep 17 00:00:00 2001 From: jesus Date: Thu, 29 Oct 2026 01:37:12 +0000 Subject: [PATCH 2/9] 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 | 79 ++++++++----- .../management/test_sso_sign_in.py | 61 +++++++++- .../proxy/management_endpoints/test_ui_sso.py | 108 +++++++++++++++++- 3 files changed, 216 insertions(+), 32 deletions(-) diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index ecd6994a2c1..634b9231328 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/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(): diff --git a/tests/integration/management/test_sso_sign_in.py b/tests/integration/management/test_sso_sign_in.py index 1af3e5443dc..f83fa43e404 100644 --- a/tests/integration/management/test_sso_sign_in.py +++ b/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]}]) diff --git a/tests/unit/proxy/management_endpoints/test_ui_sso.py b/tests/unit/proxy/management_endpoints/test_ui_sso.py index 2ae67472149..a0190190db0 100644 --- a/tests/unit/proxy/management_endpoints/test_ui_sso.py +++ b/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] From 683c20c13f1e4878193ca0b369a3fd0630d833ad Mon Sep 17 00:00:00 2001 From: jesus Date: Thu, 29 Oct 2026 03:21:13 +0000 Subject: [PATCH 3/9] fix(sso): build the upserted team list without mutation to satisfy the LIT001 gate --- litellm/proxy/management_endpoints/ui_sso.py | 13 ++++++------- 1 file changed, 6 insertions(+), 7 deletions(-) diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index 634b9231328..e00713bb0d7 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/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[list[str]] = [] + all_teams: Final = [] if sso_jwt_handler is not None: team_ids = sso_jwt_handler.get_all_jwt_team_ids(token) all_teams.extend(team_ids) @@ -1268,8 +1268,9 @@ def generic_response_convertor( 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) + 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] + ) # Determine user role based on role_mappings if available # Only apply role_mappings for GENERIC SSO provider @@ -1337,7 +1338,7 @@ 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_with_upsert, team_id_to_upsert=team_id_to_upsert, user_role=user_role, extra_fields=extra_fields, @@ -1872,9 +1873,7 @@ async def add_missing_team_member( - 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 = sorted(set(sso_teams) - set(user_teams)) + 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 if team_id == team_id_to_upsert else None) From 0f19833c23eea036cd1b3d8b7c2acb6f45565a0f Mon Sep 17 00:00:00 2001 From: jesus Date: Thu, 29 Oct 2026 03:35:50 +0000 Subject: [PATCH 4/9] fix(sso): ignore blank team_id_jwt_field claims and warn when creating the team fails --- litellm/proxy/management_endpoints/ui_sso.py | 30 ++++++++---- .../proxy/management_endpoints/test_ui_sso.py | 47 +++++++++++++++++-- 2 files changed, 65 insertions(+), 12 deletions(-) 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 From 9924cdb433c2d0a25d9c3869edd3a652a9b3f8a5 Mon Sep 17 00:00:00 2001 From: jesus Date: Thu, 29 Oct 2026 19:51:40 +0000 Subject: [PATCH 5/9] feat(sso): create every claimed team on login when team_id_upsert is set, on the UI SSO and API JWT paths --- litellm/proxy/auth/handle_jwt.py | 26 ++++- litellm/proxy/management_endpoints/types.py | 2 +- litellm/proxy/management_endpoints/ui_sso.py | 29 ++--- .../test_jwt_team_claim_upsert.py | 79 +++++++++++++ .../management/test_sso_sign_in.py | 7 +- tests/unit/proxy/auth/test_handle_jwt.py | 108 +++++++++++++++++- .../proxy/management_endpoints/test_ui_sso.py | 78 ++++++++----- 7 files changed, 280 insertions(+), 49 deletions(-) create mode 100644 tests/integration/authorization/test_jwt_team_claim_upsert.py diff --git a/litellm/proxy/auth/handle_jwt.py b/litellm/proxy/auth/handle_jwt.py index 9d800a1fd63..a96ffacf9e2 100644 --- a/litellm/proxy/auth/handle_jwt.py +++ b/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,23 @@ 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: + 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 +2834,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 diff --git a/litellm/proxy/management_endpoints/types.py b/litellm/proxy/management_endpoints/types.py index d4fd7a31d13..8c52875c71c 100644 --- a/litellm/proxy/management_endpoints/types.py +++ b/litellm/proxy/management_endpoints/types.py @@ -93,6 +93,6 @@ def get_litellm_user_role(role_str: object) -> LitellmUserRoles | None: class CustomOpenID(OpenID): team_ids: list[str] - team_id_to_upsert: str | None = None + team_id_upsert: bool = False user_role: LitellmUserRoles | None = None extra_fields: dict[str, object] | None = None diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index 0171a59b8e8..43067b80940 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -1263,17 +1263,18 @@ def generic_response_convertor( 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 - upsert_candidate: Final[object] = ( + 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 jwt_handler.litellm_jwtauth.team_id_upsert and team_id_jwt_field is not None + if 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) 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] + 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 not (isinstance(team_id, str) and not team_id.strip()) + ] # Determine user role based on role_mappings if available # Only apply role_mappings for GENERIC SSO provider @@ -1341,8 +1342,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=team_ids_with_upsert, - team_id_to_upsert=team_id_to_upsert, + team_ids=team_ids_to_join, + team_id_upsert=team_id_upsert, user_role=user_role, extra_fields=extra_fields, ) @@ -1880,17 +1881,17 @@ async def _upsert_sso_team(team_id: str, prisma_client: PrismaClient, user_api_k async def add_missing_team_member( user_info: NewUserResponse | LiteLLM_UserTable, 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) + - Create each missing team first when ``upsert_team`` is given (``team_id_upsert``) - Add missing user to missing teams """ 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 if team_id == team_id_to_upsert else None) + create_team_member_add_task(team_id, user_info, upsert_team) for team_id in missing_teams ) ) @@ -3651,11 +3652,13 @@ 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 [] + upsert_teams: Final = isinstance(result, CustomOpenID) and result.team_id_upsert await add_missing_team_member( user_info=user_info, sso_teams=sso_teams, - team_id_to_upsert=result.team_id_to_upsert if isinstance(result, CustomOpenID) else None, - upsert_team=partial(_upsert_sso_team, prisma_client=prisma_client, user_api_key_cache=user_api_key_cache), + upsert_team=partial(_upsert_sso_team, prisma_client=prisma_client, user_api_key_cache=user_api_key_cache) + if upsert_teams + else None, ) @staticmethod diff --git a/tests/integration/authorization/test_jwt_team_claim_upsert.py b/tests/integration/authorization/test_jwt_team_claim_upsert.py new file mode 100644 index 00000000000..d3e9fd4952f --- /dev/null +++ b/tests/integration/authorization/test_jwt_team_claim_upsert.py @@ -0,0 +1,79 @@ +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-skips-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 = 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(scenario.delete_user, subject) + 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, + ) + assert response.status_code == 200, response.text + expected_teams: Final = claimed_teams if team_id_upsert else [] + user: Final = eventually( + lambda: read_rows('SELECT teams FROM "LiteLLM_UserTable" WHERE user_id = %s', (subject,)), + lambda rows: rows == [{"teams": expected_teams}], + ) + teams: Final = read_rows( + 'SELECT team_id FROM "LiteLLM_TeamTable" WHERE team_id = ANY(%s) ORDER BY team_id', (claimed_teams,) + ) + assert (user, [row["team_id"] for row in teams]) == ([{"teams": expected_teams}], expected_teams) diff --git a/tests/integration/management/test_sso_sign_in.py b/tests/integration/management/test_sso_sign_in.py index f83fa43e404..4e6983f9b37 100644 --- a/tests/integration/management/test_sso_sign_in.py +++ b/tests/integration/management/test_sso_sign_in.py @@ -235,7 +235,7 @@ def test_sso_sign_in_upserts_the_team_id_jwt_field_team_only_when_team_id_upsert 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( +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]}" @@ -257,4 +257,7 @@ def test_sso_sign_in_with_db_team_mappings_creates_the_team_id_jwt_field_team_an '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]}]) + assert (sorted(row["team_id"] for row in teams), user) == ( + sorted([team_id, group_team_id]), + [{"teams": sorted([team_id, group_team_id])}], + ) diff --git a/tests/unit/proxy/auth/test_handle_jwt.py b/tests/unit/proxy/auth/test_handle_jwt.py index a0ea6ce0403..8071bdb33ab 100644 --- a/tests/unit/proxy/auth/test_handle_jwt.py +++ b/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 diff --git a/tests/unit/proxy/management_endpoints/test_ui_sso.py b/tests/unit/proxy/management_endpoints/test_ui_sso.py index 0f8cf3b65bf..cf42fbe1c57 100644 --- a/tests/unit/proxy/management_endpoints/test_ui_sso.py +++ b/tests/unit/proxy/management_endpoints/test_ui_sso.py @@ -6,7 +6,7 @@ from collections.abc import Awaitable, Callable, Mapping, Sequence from contextlib import ExitStack, asynccontextmanager from types import SimpleNamespace from typing import Final -from unittest.mock import ANY, AsyncMock, MagicMock, patch +from unittest.mock import AsyncMock, MagicMock, patch import httpx import pytest @@ -6683,8 +6683,7 @@ class TestAddMissingTeamMember: mock_add_member.assert_called_once_with( user_info=new_user_info, sso_teams=["entra-group-1", "entra-group-2"], - team_id_to_upsert=None, - upsert_team=ANY, + upsert_team=None, ) @pytest.mark.asyncio @@ -10967,48 +10966,55 @@ def _upsert_jwt_handler(**jwtauth: object) -> JWTHandler: @pytest.mark.parametrize( ("jwtauth", "userinfo", "expected_team_ids", "expected_upsert"), [ - ({"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}, {"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"}, {}, [], - None, + True, ), ( {"team_ids_jwt_field": "groups", "team_id_upsert": True}, {"groups": ["g1", "g2"]}, ["g1", "g2"], - None, + 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"], - "finance", + True, ), ], ids=[ "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", + "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_upserts_only_the_team_id_jwt_field_claim( +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: str | None, + expected_upsert: bool, ) -> None: from litellm.proxy.management_endpoints.ui_sso import generic_response_convertor @@ -11017,10 +11023,10 @@ def test_generic_sso_login_upserts_only_the_team_id_jwt_field_claim( jwt_handler=_upsert_jwt_handler(**jwtauth), ) - assert (result.team_ids, result.team_id_to_upsert) == (expected_team_ids, expected_upsert) + assert (result.team_ids, result.team_id_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: +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( @@ -11029,7 +11035,7 @@ def test_generic_sso_login_upserts_the_team_id_jwt_field_team_when_db_team_mappi team_mappings=TeamMappings(team_ids_jwt_field="groups"), ) - assert (result.team_ids, result.team_id_to_upsert) == (["g1", "finance"], "finance") + 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]]]: @@ -11044,21 +11050,35 @@ def _recording_team_calls(events: list[tuple[str, str]]) -> tuple[AsyncMock, Cal @pytest.mark.asyncio -async def test_sso_login_creates_only_the_upsert_team_and_does_so_before_joining_it() -> None: +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=None) + 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, ["g1", "finance", "g2"], team_id_to_upsert="finance", upsert_team=upsert_team) + 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 creates == ["finance"] + assert sorted(creates) == ["finance", "g1", "g2"] assert sorted(joins) == ["finance", "g1", "g2"] - assert events.index(("create", "finance")) < events.index(("join", "finance")) + 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 @@ -11076,7 +11096,7 @@ async def test_sso_login_still_joins_the_team_when_a_concurrent_login_created_it 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) + await add_missing_team_member(user, ["finance"], upsert_team=lost_the_create_race) assert joined == ["finance"] @@ -11127,7 +11147,7 @@ async def test_sso_login_warns_when_creating_the_team_fails_instead_of_calling_i AsyncMock(side_effect=TeamNotFoundError(team_id="finance")), ), ): - await add_missing_team_member(user, ["finance"], team_id_to_upsert="finance", upsert_team=failed_create) + 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] From ac0880cfaf7bb5c1abc479b4aacb19895cc83e5e Mon Sep 17 00:00:00 2001 From: jesus Date: Thu, 29 Oct 2026 19:58:02 +0000 Subject: [PATCH 6/9] test(jwt): assert the teams array claim is rejected without team_id_upsert and joined with it --- .../test_jwt_team_claim_upsert.py | 27 ++++++++++++------- 1 file changed, 17 insertions(+), 10 deletions(-) diff --git a/tests/integration/authorization/test_jwt_team_claim_upsert.py b/tests/integration/authorization/test_jwt_team_claim_upsert.py index d3e9fd4952f..713c2c5adc3 100644 --- a/tests/integration/authorization/test_jwt_team_claim_upsert.py +++ b/tests/integration/authorization/test_jwt_team_claim_upsert.py @@ -23,7 +23,9 @@ 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-skips-them"]) +@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: @@ -59,7 +61,6 @@ def test_jwt_teams_array_claim_creates_and_joins_every_team_only_when_team_id_up ) with owned_proxy(gateway, tmp_path, {"JWT_PUBLIC_KEY_URL": jwks.url}, config=path) as candidate: model: Final = scenario.model() - scenario.cleanups.callback(scenario.delete_user, subject) scenario.cleanups.callback(gateway.request, "POST", "/team/delete", {"team_ids": claimed_teams}) response: Final = candidate.request( "POST", @@ -67,13 +68,19 @@ def test_jwt_teams_array_claim_creates_and_joins_every_team_only_when_team_id_up {"model": model, "messages": [{"role": "user", "content": "teams array claim"}]}, key=token, ) - assert response.status_code == 200, response.text - expected_teams: Final = claimed_teams if team_id_upsert else [] - user: Final = eventually( - lambda: read_rows('SELECT teams FROM "LiteLLM_UserTable" WHERE user_id = %s', (subject,)), - lambda rows: rows == [{"teams": expected_teams}], - ) - teams: Final = read_rows( + team_rows: Final = read_rows( 'SELECT team_id FROM "LiteLLM_TeamTable" WHERE team_id = ANY(%s) ORDER BY team_id', (claimed_teams,) ) - assert (user, [row["team_id"] for row in teams]) == ([{"teams": expected_teams}], expected_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, + ) From 58fa60ab6a9c5edc4c6a4ccf11f00c22392216bb Mon Sep 17 00:00:00 2001 From: jesus Date: Thu, 29 Oct 2026 20:12:21 +0000 Subject: [PATCH 7/9] style(sso): ruff format the new upsert blocks --- litellm/proxy/auth/handle_jwt.py | 4 +++- litellm/proxy/management_endpoints/ui_sso.py | 9 ++++----- 2 files changed, 7 insertions(+), 6 deletions(-) diff --git a/litellm/proxy/auth/handle_jwt.py b/litellm/proxy/auth/handle_jwt.py index a96ffacf9e2..66f9a96bece 100644 --- a/litellm/proxy/auth/handle_jwt.py +++ b/litellm/proxy/auth/handle_jwt.py @@ -2247,7 +2247,9 @@ class JWTAuthManager: team_id_upsert=True, ) except Exception as e: - verbose_proxy_logger.debug("JWT auth could not create team %s, trying to join it anyway: %s", team_id, e) + 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( diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index 43067b80940..c5c8ad461b8 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -1270,7 +1270,9 @@ def generic_response_convertor( 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] + 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 not (isinstance(team_id, str) and not team_id.strip()) @@ -1890,10 +1892,7 @@ async def add_missing_team_member( """ 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 - ) + *(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) From f0fcd73aa051ce7f791bbd6ff22d3c2d3f706ffc Mon Sep 17 00:00:00 2001 From: jesus Date: Thu, 29 Oct 2026 20:33:09 +0000 Subject: [PATCH 8/9] fix(jwt): name the BLE001 suppression on the fail-safe team create in _upsert_claimed_teams --- litellm/proxy/auth/handle_jwt.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/proxy/auth/handle_jwt.py b/litellm/proxy/auth/handle_jwt.py index 66f9a96bece..e8585ff83a3 100644 --- a/litellm/proxy/auth/handle_jwt.py +++ b/litellm/proxy/auth/handle_jwt.py @@ -2246,7 +2246,7 @@ class JWTAuthManager: check_db_only=True, team_id_upsert=True, ) - except Exception as e: + 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 ) From caf4c043611beaf3426a77e2d5ce1cfa55b39cb2 Mon Sep 17 00:00:00 2001 From: jesus Date: Thu, 29 Oct 2026 21:07:49 +0000 Subject: [PATCH 9/9] fix(sso): type the sso team list so the basedpyright gate stays at its ceiling --- litellm/proxy/management_endpoints/ui_sso.py | 6 ++---- .../integration/authorization/test_jwt_team_claim_upsert.py | 2 +- 2 files changed, 3 insertions(+), 5 deletions(-) diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index c5c8ad461b8..20f6ad8bed0 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -1241,7 +1241,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) @@ -1274,9 +1274,7 @@ def generic_response_convertor( 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 not (isinstance(team_id, str) and not team_id.strip()) - ] + 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 diff --git a/tests/integration/authorization/test_jwt_team_claim_upsert.py b/tests/integration/authorization/test_jwt_team_claim_upsert.py index 713c2c5adc3..9cc197858dd 100644 --- a/tests/integration/authorization/test_jwt_team_claim_upsert.py +++ b/tests/integration/authorization/test_jwt_team_claim_upsert.py @@ -30,7 +30,7 @@ def test_jwt_teams_array_claim_creates_and_joins_every_team_only_when_team_id_up 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 = jwt.algorithms.RSAAlgorithm.to_jwk(private_key.public_key()) + public_jwk: Final[str] = jwt.algorithms.RSAAlgorithm.to_jwk(private_key.public_key()) def respond(request: Request) -> Reply: assert request.method == "GET", request