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

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

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

* fix(sso): keep the upserted team_id_jwt_field team out of result.team_ids so restricted_sso_group and the access-token group fallback are unchanged

* fix(sso): only upsert the team_id_jwt_field team on UI SSO login when enable_jwt_auth is on

The API JWT path only reaches team_id_upsert inside the enable_jwt_auth branch, so a deployment
with a leftover litellm_jwtauth block and JWT auth off never created teams. Gate the UI SSO upsert
the same way

* fix(sso): read enable_jwt_auth through the proxy_server module so the basedpyright gate stays at its ceiling
This commit is contained in:
jesus-berri 2026-10-30 14:39:25 -07:00 • committed by GitHub
parent ec4f66bb94
commit e8f572b39b
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 448 additions and 39 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

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

@ -87,7 +87,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,
@ -1266,6 +1266,7 @@ def generic_response_convertor(
sso_jwt_handler: JWTHandler | None = None,
role_mappings: Optional["RoleMappings"] = None,
team_mappings: Optional["TeamMappings"] = None,
jwt_auth_enabled: bool = False,
) -> CustomOpenID:
generic_user_id_attribute_name: Final = os.getenv("GENERIC_USER_ID_ATTRIBUTE", "preferred_username")
generic_user_display_name_attribute_name: Final = os.getenv("GENERIC_USER_DISPLAY_NAME_ATTRIBUTE", "sub")
@ -1286,14 +1287,15 @@ def generic_response_convertor(
generic_user_email_attribute_name,
)
all_teams: Final = []
token: Final = cast(dict[str, object], response) # cast-ok: the IdP userinfo JSON object fastapi-sso hands back
all_teams: Final = list[str]()
if sso_jwt_handler is not None:
team_ids = sso_jwt_handler.get_all_jwt_team_ids(cast(dict, response))
team_ids = sso_jwt_handler.get_all_jwt_team_ids(token)
all_teams.extend(team_ids)
if team_mappings is not None and team_mappings.team_ids_jwt_field is not None:
team_ids_from_db_mapping: Final[list[str] | None] = get_nested_value(
data=cast(dict, response),
data=token,
key_path=team_mappings.team_ids_jwt_field,
default=[],
)
@ -1305,8 +1307,17 @@ 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_auth_enabled and 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) and upsert_candidate.strip() else None
)
# Determine user role based on role_mappings if available
# Only apply role_mappings for GENERIC SSO provider
@ -1375,6 +1386,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=team_id_to_upsert,
user_role=user_role,
extra_fields=extra_fields,
)
@ -1614,6 +1626,8 @@ async def get_generic_sso_response(
# make generic sso provider
from fastapi_sso.sso.base import DiscoveryDocument
from litellm.proxy import proxy_server
received_response: dict | None = None
sso_assertion: SSOIdentityAssertion | None = None
@ -1665,6 +1679,7 @@ async def get_generic_sso_response(
sso_jwt_handler=sso_jwt_handler,
role_mappings=role_mappings,
team_mappings=team_mappings,
jwt_auth_enabled=proxy_server.general_settings.get("enable_jwt_auth", False) is True,
)
SSOProvider: Final = create_generic_sso_provider(
@ -1884,6 +1899,29 @@ async def find_existing_team_ids(team_ids: Sequence[str]) -> frozenset[str] | No
return frozenset(row.team_id for row in rows)
async def create_sso_team(team_id: str) -> None:
"""Create the ``team_id_jwt_field`` team on UI SSO login the way the API JWT path does under ``team_id_upsert``."""
from litellm.proxy.proxy_server import prisma_client, user_api_key_cache
if prisma_client is None:
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 _create_sso_team_or_warn(team_id: str, create_team: Callable[[str], Awaitable[object]]) -> None:
"""A concurrent login may create the team first, so a failed create is not fatal: the join still runs."""
try:
await create_team(team_id)
except Exception as e: # noqa: BLE001 # fail-safe: a failed team create must not fail the login or skip the join
verbose_proxy_logger.warning("SSO login could not create team %s, trying to join it anyway: %s", team_id, e)
async def add_missing_team_member(
user_info: NewUserResponse | LiteLLM_UserTable,
sso_teams: Sequence[str],
@ -1892,21 +1930,30 @@ async def add_missing_team_member(
[str, NewUserResponse | LiteLLM_UserTable], Awaitable[object]
] = create_team_member_add_task,
existing_team_ids: Callable[[Sequence[str]], Awaitable[frozenset[str] | None]] = find_existing_team_ids,
team_id_to_upsert: str | None = None,
create_team: Callable[[str], Awaitable[object]] = create_sso_team,
) -> None:
"""
- Get missing teams (diff b/w user_info.team_ids and sso_teams)
- Drop the IdP groups that are not teams with one lookup, since group claims often carry many of them
- Join the ``team_id_to_upsert`` team too, creating it when it is missing (``team_id_upsert``), the way the API JWT path does;
it stays out of ``result.team_ids`` so restricted_sso_group and the access-token group fallback see only the claimed groups
- Add missing user to missing teams, one at a time in sso_teams order, so user.teams keeps the claim's order
"""
held_teams: Final = frozenset(user_info.teams or ())
missing_teams: Final = tuple(team_id for team_id in dict.fromkeys(sso_teams) if team_id not in held_teams)
claimed_teams: Final = sso_teams if team_id_to_upsert is None else (*sso_teams, team_id_to_upsert)
missing_teams: Final = tuple(team_id for team_id in dict.fromkeys(claimed_teams) if team_id not in held_teams)
if not missing_teams:
return
existing: Final = await existing_team_ids(missing_teams)
teams_to_add: Final = (
missing_teams if existing is None else tuple(team_id for team_id in missing_teams if team_id in existing)
missing_teams
if existing is None
else tuple(team_id for team_id in missing_teams if team_id in existing or team_id == team_id_to_upsert)
)
for team_id in teams_to_add:
if team_id == team_id_to_upsert and (existing is None or team_id not in existing):
await _create_sso_team_or_warn(team_id, create_team)
await add_team_member(team_id, user_info)
@ -3663,7 +3710,11 @@ 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,
)
@staticmethod
def verify_user_in_restricted_sso_group(

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

@ -16,7 +16,7 @@ import yaml
from pydantic import JsonValue
from tests.integration._support.client import Gateway, gateway_from_environment
from tests.integration._support.database import read_rows
from tests.integration._support.database import read_rows, write_rows
from tests.integration._support.process import OWNED_STATE_VARIABLE, OwnedState, owned_proxy
from tests.integration._support.wire import Reply, Request, wire_server
@ -52,14 +52,29 @@ def _idp_reply(request: Request) -> Reply:
"preferred_username": signed_in,
"email": f"{signed_in}@example.com",
"provider": IDP_PROVIDER,
"dept": _dept(signed_in),
"groups": [_group(signed_in)],
}
).encode()
)
return Reply(status=404, body=b'{"error": "not_found"}')
def _dept(subject: str) -> str:
return f"dept-{subject}"
def _group(subject: str) -> str:
return f"group-{subject}"
@contextmanager
def _sso_gateway(directory: Path, litellm_settings: Mapping[str, JsonValue]) -> Iterator[Gateway]:
def _sso_gateway(
directory: Path,
litellm_settings: Mapping[str, JsonValue],
general_settings: Mapping[str, JsonValue] | None = None,
db_sso_settings: Mapping[str, JsonValue] | None = None,
) -> Iterator[Gateway]:
config: Final = directory / "sso_proxy_config.yaml"
config.write_text(
yaml.safe_dump(
@ -68,6 +83,7 @@ def _sso_gateway(directory: Path, litellm_settings: Mapping[str, JsonValue]) ->
"general_settings": {
"master_key": "os.environ/LITELLM_MASTER_KEY",
"database_url": "os.environ/DATABASE_URL",
**(general_settings or {}),
},
"litellm_settings": dict(litellm_settings),
}
@ -83,7 +99,33 @@ def _sso_gateway(directory: Path, litellm_settings: Mapping[str, JsonValue]) ->
"OAUTHLIB_INSECURE_TRANSPORT": "1",
}
with owned_proxy(rig, directory, environment, config=config, remove_environment=("PROXY_BASE_URL",)) as gw:
yield gw
if db_sso_settings is None:
yield gw
return
with contextlib.ExitStack() as cleanups:
cleanups.callback(
write_rows,
'DELETE FROM "LiteLLM_SSOConfig" WHERE id = %s',
("sso_config",),
database_url=gw.database_url,
)
stored: Final = _store_sso_settings(gw, {**_generic_sso_settings(environment), **db_sso_settings})
assert stored.status_code == 200, stored.text
yield gw
def _generic_sso_settings(environment: Mapping[str, str]) -> Mapping[str, JsonValue]:
return {
"generic_client_id": environment["GENERIC_CLIENT_ID"],
"generic_client_secret": environment["GENERIC_CLIENT_SECRET"],
"generic_authorization_endpoint": environment["GENERIC_AUTHORIZATION_ENDPOINT"],
"generic_token_endpoint": environment["GENERIC_TOKEN_ENDPOINT"],
"generic_userinfo_endpoint": environment["GENERIC_USERINFO_ENDPOINT"],
}
def _store_sso_settings(gateway: Gateway, settings: Mapping[str, JsonValue]) -> httpx.Response:
return gateway.request("PATCH", "/update/sso_settings", settings)
def _sign_in(gateway: Gateway, subject: str) -> httpx.Response:
@ -165,3 +207,65 @@ 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(
("enable_jwt_auth", "team_id_upsert", "team_exists_before", "expect_member"),
[(True, True, False, True), (False, True, False, False), (True, False, False, False), (True, False, True, True)],
ids=[
"upsert-creates-missing-team",
"upsert-without-jwt-auth-skips-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, enable_jwt_auth: bool, 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 = {
"enable_jwt_auth": enable_jwt_auth,
"litellm_jwtauth": {"team_id_jwt_field": "dept", "team_id_upsert": team_id_upsert},
}
with _sso_gateway(tmp_path, {}, jwtauth) as gateway, contextlib.ExitStack() as cleanups:
cleanups.callback(gateway.request, "POST", "/team/delete", {"team_ids": [team_id]})
cleanups.callback(_delete_user, gateway, subject)
if team_exists_before:
created: Final = gateway.request("POST", "/team/new", {"team_id": team_id})
assert created.status_code == 200, created.text
callback: Final = _sign_in(gateway, subject)
assert callback.status_code == 303, f"{callback.status_code} {callback.text}"
teams: Final = gateway.rows('SELECT team_id FROM "LiteLLM_TeamTable" WHERE team_id = %s', (team_id,))
user: Final = gateway.rows('SELECT teams FROM "LiteLLM_UserTable" WHERE user_id = %s', (subject,))
assert (len(teams), user) == (int(expect_member), [{"teams": [team_id] if expect_member else []}])
def test_sso_sign_in_with_db_team_mappings_creates_the_team_id_jwt_field_team_and_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 = {
"enable_jwt_auth": True,
"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]}])

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

@ -15,9 +15,15 @@ from fastapi import HTTPException, Request
from pydantic import SecretStr
import litellm
from litellm._logging import verbose_proxy_logger
from litellm._uuid import uuid
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
from litellm.proxy._types import LiteLLM_UserTable, NewUserResponse, ProxyException
from litellm.proxy._types import (
LiteLLM_JWTAuth,
LiteLLM_UserTable,
NewUserResponse,
ProxyException,
)
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 +293,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 +1171,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 +2239,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 +2312,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 +2372,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 +2448,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 +2591,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 +2613,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",
@ -4961,7 +4974,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
@ -5348,7 +5361,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
@ -5419,7 +5432,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 = {
@ -5449,7 +5462,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 = {
@ -6813,7 +6826,9 @@ class TestAddMissingTeamMember:
# 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,
)
@pytest.mark.asyncio
@ -7897,7 +7912,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(
@ -7931,7 +7946,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(
@ -7959,7 +7974,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(
@ -7986,7 +8001,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(
@ -8865,7 +8880,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,
@ -8917,7 +8932,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,
@ -8980,7 +8995,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,
@ -9087,7 +9102,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.
@ -9173,7 +9188,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)
@ -9224,7 +9239,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(
@ -9589,7 +9604,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,
@ -9615,7 +9630,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()
@ -9680,7 +9695,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
@ -11144,3 +11159,241 @@ 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": ""}, [], 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"},
{},
[],
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",
"blank-claim-not-upserted",
"whitespace-claim-not-upserted",
"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),
jwt_auth_enabled=True,
)
assert (result.team_ids, result.team_id_to_upsert) == (expected_team_ids, expected_upsert)
def test_generic_sso_login_does_not_upsert_the_team_id_jwt_field_claim_when_jwt_auth_is_disabled() -> 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", "dept": "finance"},
jwt_handler=_upsert_jwt_handler(team_id_jwt_field="dept", team_id_upsert=True),
jwt_auth_enabled=False,
)
assert (result.team_ids, result.team_id_to_upsert) == (["finance"], None)
def test_generic_sso_login_marks_the_team_id_jwt_field_team_for_upsert_without_adding_it_to_team_ids() -> 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"),
jwt_auth_enabled=True,
)
assert (result.team_ids, result.team_id_to_upsert) == (["g1"], "finance")
def test_generic_sso_login_keeps_team_ids_empty_for_a_dept_only_user_so_the_access_token_group_fallback_runs() -> 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", "dept": "d1"},
jwt_handler=_upsert_jwt_handler(team_id_jwt_field="dept", team_id_upsert=True),
team_mappings=TeamMappings(team_ids_jwt_field="groups"),
jwt_auth_enabled=True,
)
assert (result.team_ids, result.team_id_to_upsert) == ([], "d1")
def test_restricted_sso_group_still_rejects_a_user_whose_only_team_is_the_upserted_dept() -> None:
from litellm.proxy._types import ProxyException
from litellm.proxy.management_endpoints.ui_sso import SSOAuthenticationHandler, generic_response_convertor
response: Final = {"preferred_username": "alice", "email": "alice@example.com", "dept": "X", "groups": []}
result: Final = generic_response_convertor(
response=response,
jwt_handler=_upsert_jwt_handler(team_id_jwt_field="dept", team_id_upsert=True),
team_mappings=TeamMappings(team_ids_jwt_field="groups"),
jwt_auth_enabled=True,
)
with pytest.raises(ProxyException) as denied:
SSOAuthenticationHandler.verify_user_in_restricted_sso_group(
general_settings={"ui_access_mode": {"type": "restricted_sso_group", "restricted_sso_group": "X"}},
result=result,
received_response=response,
)
assert denied.value.code == "403"
def _recording_team_calls(
events: list[tuple[str, str]],
) -> tuple[Callable[[str, NewUserResponse | LiteLLM_UserTable], Awaitable[None]], Callable[[str], Awaitable[None]]]:
async def add(team_id: str, user_info: NewUserResponse | LiteLLM_UserTable) -> None:
events.append(("join", team_id))
async def create(team_id: str) -> None:
events.append(("create", team_id))
return add, create
async def _only_group_teams_exist(team_ids: Sequence[str]) -> frozenset[str] | None:
return frozenset(team_id for team_id in team_ids if team_id.startswith("g"))
@pytest.mark.asyncio
async def test_sso_login_joins_the_upsert_team_from_team_id_to_upsert_and_creates_it_before_the_join() -> 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
add, create = _recording_team_calls(events)
await add_missing_team_member(
NewUserResponse(user_id="alice", key="sk-x", teams=None),
["g1", "g2"],
add_team_member=add,
existing_team_ids=_only_group_teams_exist,
team_id_to_upsert="finance",
create_team=create,
)
assert events == [("join", "g1"), ("join", "g2"), ("create", "finance"), ("join", "finance")]
@pytest.mark.asyncio
async def test_sso_login_joins_only_the_claimed_groups_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
add, create = _recording_team_calls(events)
await add_missing_team_member(
NewUserResponse(user_id="alice", key="sk-x", teams=None),
["g1", "g2"],
add_team_member=add,
existing_team_ids=_only_group_teams_exist,
create_team=create,
)
assert events == [("join", "g1"), ("join", "g2")]
@pytest.mark.asyncio
async def test_sso_login_does_not_create_the_team_id_jwt_field_team_when_it_already_exists() -> 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
add, create = _recording_team_calls(events)
async def every_team_exists(team_ids: Sequence[str]) -> frozenset[str] | None:
return frozenset(team_ids)
await add_missing_team_member(
NewUserResponse(user_id="alice", key="sk-x", teams=None),
["finance"],
add_team_member=add,
existing_team_ids=every_team_exists,
team_id_to_upsert="finance",
create_team=create,
)
assert events == [("join", "finance")]
@pytest.mark.asyncio
@pytest.mark.parametrize(
"create_error",
[
HTTPException(status_code=400, detail={"error": "Team id = finance already exists"}),
HTTPException(status_code=403, detail={"error": "team limit reached for this license"}),
],
ids=["lost-the-create-race", "create-rejected"],
)
async def test_sso_login_still_joins_the_team_and_warns_when_creating_it_fails(
caplog: pytest.LogCaptureFixture, create_error: HTTPException
) -> 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
add, _ = _recording_team_calls(events)
async def failed_create(team_id: str) -> None:
raise create_error
with caplog.at_level(logging.WARNING, logger=verbose_proxy_logger.name):
await add_missing_team_member(
NewUserResponse(user_id="alice", key="sk-x", teams=None),
["finance"],
add_team_member=add,
existing_team_ids=_only_group_teams_exist,
team_id_to_upsert="finance",
create_team=failed_create,
)
warnings: Final = [r.getMessage() for r in caplog.records if r.levelno == logging.WARNING]
assert events == [("join", "finance")]
assert len(warnings) == 1 and "finance" in warnings[0] and str(create_error.detail) in warnings[0]