scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-12 08:42:18 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 207

feat(sso): create the team_id_jwt_field team on UI SSO login when team_id_upsert is set

Browse source 
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
This commit is contained in:
Jesus Velarde 2026-10-28 16:21:31 -07:00
parent ddb8a14ee9
commit 66acab51c1
4 changed files with 197 additions and 41 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
litellm/proxy/management_endpoints/types.py
View file

@ -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

62
litellm/proxy/management_endpoints/ui_sso.py
View file

@ -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(

38
tests/integration/management/test_sso_sign_in.py
View file

@ -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 []}])

137
tests/unit/proxy/management_endpoints/test_ui_sso.py
View file

@ -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)