mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-13 08:52:02 +00:00
fix(sso): keep CLI SSO and Claude Code gateway team lists in the user record's order (#49033)
* fix(sso): bind CLI SSO and Claude Code gateway logins to the first team on the user record * fix(proxy): keep SCIM and JWT team sync in record order, and look up SSO groups once SCIM PATCH and JWT team sync saved a user's teams from a set, so the first team on the record, the one a Claude Code gateway sign-in binds, could change on any sync. Both now keep the record's order and append new teams in the order the IdP sent them. SSO group sync now drops claim groups that are not LiteLLM teams with one lookup before adding the missing teams one at a time, so a claim full of non-team groups no longer costs one failed add per group on every sign-in. * test(sso): expect the gateway team dropdown in the user record's order * fix(sso): read team existence from the writer, keep new-user and SCIM replace team order * fix(scim): keep the teams a user already has in record order on PATCH and PUT * fix(scim): keep record team order when a SCIM POST matches an existing email * fix(scim): join new teams in request order so a failed roster write keeps it * fix(scim): let a groups replace win over later group ops in a PATCH again Applying group adds and removes after a replace made a PATCH that deactivates a user and adds a team that does not exist answer 404 and leave the user active, where it answered 200 before. Keep the record's team order, but ignore group ops after the last replace as the base did. * test(sso): type the CLI team detail row helper --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Co-authored-by: rj-cairn <349020452+rj-cairn@users.noreply.github.com>
This commit is contained in:
zoveki
GitHub
parent
02ef44c388
commit
87fec5467f
10 changed files with 661 additions and 191 deletions
|
|
@ -2199,15 +2199,19 @@ class JWTAuthManager:
|
|||
# not mistaken for claimless and left with stale DB memberships the fallback
|
||||
# could later attribute. With the flag off, keep the upstream plural-only
|
||||
# reconciliation so existing deployments are unchanged.
|
||||
jwt_team_ids: Final = set(
|
||||
jwt_handler.get_all_jwt_team_ids(jwt_valid_token)
|
||||
if jwt_handler.litellm_jwtauth.fallback_to_db_teams
|
||||
else jwt_handler.get_team_ids_from_jwt(jwt_valid_token)
|
||||
claim_team_ids: Final = tuple(
|
||||
dict.fromkeys(
|
||||
jwt_handler.get_all_jwt_team_ids(jwt_valid_token)
|
||||
if jwt_handler.litellm_jwtauth.fallback_to_db_teams
|
||||
else jwt_handler.get_team_ids_from_jwt(jwt_valid_token)
|
||||
)
|
||||
)
|
||||
existing_teams: Final = set(user_object.teams or [])
|
||||
teams_to_add: Final = jwt_team_ids - existing_teams
|
||||
jwt_team_ids: Final = frozenset(claim_team_ids)
|
||||
existing_team_order: Final = tuple(dict.fromkeys(user_object.teams or []))
|
||||
existing_teams: Final = frozenset(existing_team_order)
|
||||
teams_to_add: Final = tuple(team_id for team_id in claim_team_ids if team_id not in existing_teams)
|
||||
preserve_db_teams_without_claims: Final = jwt_handler.litellm_jwtauth.fallback_to_db_teams and not jwt_team_ids
|
||||
teams_to_remove: Final = set() if preserve_db_teams_without_claims else existing_teams - jwt_team_ids
|
||||
teams_to_remove: Final = frozenset() if preserve_db_teams_without_claims else existing_teams - jwt_team_ids
|
||||
if teams_to_add or teams_to_remove:
|
||||
from litellm.proxy.management_endpoints.scim.scim_v2 import (
|
||||
patch_team_membership,
|
||||
|
|
@ -2218,7 +2222,10 @@ class JWTAuthManager:
|
|||
teams_ids_to_add_user_to=list(teams_to_add),
|
||||
teams_ids_to_remove_user_from=list(teams_to_remove),
|
||||
)
|
||||
user_object.teams = list(jwt_team_ids)
|
||||
user_object.teams = [
|
||||
*(team_id for team_id in existing_team_order if team_id not in teams_to_remove),
|
||||
*teams_to_add,
|
||||
]
|
||||
if user_api_key_cache is not None:
|
||||
await user_api_key_cache.async_set_cache(
|
||||
key=user_object.user_id,
|
||||
|
|
|
|||
|
|
@ -13,6 +13,7 @@ These are members of a Team on LiteLLM
|
|||
"""
|
||||
|
||||
import asyncio
|
||||
import contextlib
|
||||
import json
|
||||
import traceback
|
||||
from collections.abc import Awaitable, Mapping, Sequence
|
||||
|
|
@ -443,30 +444,25 @@ async def add_new_user_to_default_team(
|
|||
teams: Sequence[str | NewUserRequestTeam],
|
||||
prisma_client: "PrismaClient",
|
||||
):
|
||||
tasks: Final[list[Awaitable[object]]] = []
|
||||
for team in teams:
|
||||
user_role: Literal["user", "admin"] = "user"
|
||||
max_budget_in_team: float | None = None
|
||||
if isinstance(team, str):
|
||||
team_id = team
|
||||
elif isinstance(team, NewUserRequestTeam):
|
||||
team_id = team.team_id
|
||||
user_role = team.user_role
|
||||
max_budget_in_team = team.max_budget_in_team
|
||||
else:
|
||||
raise ValueError(f"Invalid team type: {type(team)}")
|
||||
|
||||
tasks.append(
|
||||
_add_user_to_team(
|
||||
memberships: Final = tuple(_default_team_membership(team) for team in teams)
|
||||
for membership in memberships:
|
||||
with contextlib.suppress(Exception):
|
||||
await _add_user_to_team(
|
||||
user_id=user_id,
|
||||
team_id=team_id,
|
||||
team_id=membership.team_id,
|
||||
user_email=user_email,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
max_budget_in_team=max_budget_in_team,
|
||||
user_role=user_role,
|
||||
max_budget_in_team=membership.max_budget_in_team,
|
||||
user_role=membership.user_role,
|
||||
)
|
||||
)
|
||||
await asyncio.gather(*tasks, return_exceptions=True)
|
||||
|
||||
|
||||
def _default_team_membership(team: str | NewUserRequestTeam) -> NewUserRequestTeam:
|
||||
if isinstance(team, str):
|
||||
return NewUserRequestTeam(team_id=team)
|
||||
if isinstance(team, NewUserRequestTeam):
|
||||
return team
|
||||
raise ValueError(f"Invalid team type: {type(team)}")
|
||||
|
||||
|
||||
async def _fetch_user_team_ids(user_id: str, prisma_client: "PrismaClient") -> tuple[str, ...]:
|
||||
|
|
|
|||
|
|
@ -229,7 +229,11 @@ class UserProvisionerHelpers:
|
|||
return None
|
||||
|
||||
requested_teams: Final = list(dict.fromkeys(new_user_request.teams or []))
|
||||
new_teams: Final = requested_teams if requested_teams else list(existing_user.teams or [])
|
||||
new_teams: Final = (
|
||||
list(_keep_record_team_order(existing_user.teams or [], requested_teams))
|
||||
if requested_teams
|
||||
else list(existing_user.teams or [])
|
||||
)
|
||||
|
||||
if new_user_request.user_id != existing_user.user_id:
|
||||
verbose_proxy_logger.info(
|
||||
|
|
@ -1165,7 +1169,7 @@ async def _handle_team_membership_changes(
|
|||
existing_teams_set: Final = set(existing_teams)
|
||||
new_teams_set: Final = set(new_teams)
|
||||
|
||||
teams_to_add: Final = new_teams_set - existing_teams_set
|
||||
teams_to_add: Final = tuple(team for team in dict.fromkeys(new_teams) if team not in existing_teams_set)
|
||||
teams_to_remove: Final = existing_teams_set - new_teams_set
|
||||
|
||||
if teams_to_add or teams_to_remove:
|
||||
|
|
@ -1948,7 +1952,11 @@ async def update_user(
|
|||
|
||||
# SCIM User.groups is readOnly (RFC 7643 4.1.2): IdPs sync membership via /Groups and send
|
||||
# no groups or `groups: []` on profile PUTs, so empty means unspecified, not "remove from every team"
|
||||
target_teams: Final = user_data["teams"] or existing_user.teams
|
||||
target_teams: Final = (
|
||||
list(_keep_record_team_order(existing_user.teams, user_data["teams"]))
|
||||
if user_data["teams"]
|
||||
else existing_user.teams
|
||||
)
|
||||
await _handle_team_membership_changes(
|
||||
user_id=user_id,
|
||||
existing_teams=existing_user.teams,
|
||||
|
|
@ -2241,19 +2249,29 @@ def _handle_name_update(path: str, op_type: str, value: object, scim_metadata: d
|
|||
scim_metadata["familyName"] = str(value)
|
||||
|
||||
|
||||
def _handle_group_operations(op_type: str, value: object, teams_set: set[str], path: str | None) -> set[str] | None:
|
||||
"""Handle group/team membership operations."""
|
||||
group_values = _extract_group_values(value)
|
||||
if not group_values and value is None:
|
||||
group_values = _extract_ids_from_path_filter(path, "groups")
|
||||
if op_type == "replace":
|
||||
return set(group_values)
|
||||
elif op_type == "add":
|
||||
teams_set.update(group_values)
|
||||
elif op_type == "remove":
|
||||
for gid in group_values:
|
||||
teams_set.discard(gid)
|
||||
return None
|
||||
def _is_group_path(path: str | None) -> bool:
|
||||
return isinstance(path, str) and path.lower().startswith("groups")
|
||||
|
||||
|
||||
def _apply_group_operation(teams: tuple[str, ...], op: SCIMPatchOperation) -> tuple[str, ...]:
|
||||
listed: Final = _extract_group_values(op.value)
|
||||
group_values: Final = listed if listed or op.value is not None else _extract_ids_from_path_filter(op.path, "groups")
|
||||
if op.op == "replace":
|
||||
return tuple(dict.fromkeys(group_values))
|
||||
if op.op == "add":
|
||||
return tuple(dict.fromkeys((*teams, *group_values)))
|
||||
removed: Final = frozenset(group_values)
|
||||
return tuple(team for team in teams if team not in removed)
|
||||
|
||||
|
||||
def _keep_record_team_order(record_teams: Sequence[str], target_teams: Sequence[str]) -> tuple[str, ...]:
|
||||
target: Final = frozenset(target_teams)
|
||||
record: Final = tuple(dict.fromkeys(record_teams))
|
||||
known: Final = frozenset(record)
|
||||
return (
|
||||
*(team for team in record if team in target),
|
||||
*(team for team in dict.fromkeys(target_teams) if team not in known),
|
||||
)
|
||||
|
||||
|
||||
def _multi_valued_attribute_base(path: str) -> str:
|
||||
|
|
@ -2520,8 +2538,8 @@ def _handle_generic_metadata(path: str, op_type: str, value: object, metadata: d
|
|||
def _apply_patch_ops(
|
||||
existing_user: LiteLLM_UserTable,
|
||||
patch_ops: SCIMPatchOp,
|
||||
) -> tuple[dict[str, object], set[str]]:
|
||||
"""Apply patch operations and return update data and final team set."""
|
||||
) -> tuple[dict[str, object], tuple[str, ...]]:
|
||||
"""Apply patch operations and return update data and the final teams in order."""
|
||||
update_data: Final[dict[str, object]] = {}
|
||||
stored_metadata: Final[object] = TypeAdapter(object).validate_python(existing_user.model_dump().get("metadata"))
|
||||
metadata: Final = _JSON_OBJECT_ADAPTER.validate_python(stored_metadata) if isinstance(stored_metadata, dict) else {}
|
||||
|
|
@ -2541,9 +2559,6 @@ def _apply_patch_ops(
|
|||
if isinstance(enterprise, _InvalidEnterpriseValue):
|
||||
_raise_invalid_enterprise_value(enterprise)
|
||||
|
||||
teams_set: Final[set[str]] = set(existing_user.teams or [])
|
||||
replace_team_set: set[str] | None = None
|
||||
|
||||
for op in patch_ops.Operations:
|
||||
path = op.path.lower() if isinstance(op.path, str) else ""
|
||||
value = op.value
|
||||
|
|
@ -2585,16 +2600,20 @@ def _apply_patch_ops(
|
|||
_handle_name_update(path, op_type, value, scim_metadata)
|
||||
elif _multi_valued_attribute_base(path) in SCIM_MULTI_VALUED_ATTRIBUTE_METADATA_KEYS:
|
||||
_handle_multi_valued_attribute_update(path, op_type, value, metadata)
|
||||
elif path.startswith("groups"):
|
||||
new_replace_set = _handle_group_operations(op_type, value, teams_set, op.path)
|
||||
if new_replace_set is not None:
|
||||
replace_team_set = new_replace_set
|
||||
elif _enterprise_extension_sub_path(path) is not None:
|
||||
elif _is_group_path(path) or _enterprise_extension_sub_path(path) is not None:
|
||||
continue
|
||||
else:
|
||||
_handle_generic_metadata(path, op_type, value, metadata)
|
||||
|
||||
final_team_set: Final = replace_team_set if replace_team_set is not None else teams_set
|
||||
record_teams: Final = tuple(existing_user.teams or [])
|
||||
group_ops: Final = tuple(op for op in patch_ops.Operations if _is_group_path(op.path))
|
||||
last_replace: Final = next((op for op in reversed(group_ops) if op.op == "replace"), None)
|
||||
final_teams: Final = _keep_record_team_order(
|
||||
record_teams,
|
||||
_apply_group_operation(record_teams, last_replace)
|
||||
if last_replace is not None
|
||||
else reduce(_apply_group_operation, group_ops, record_teams),
|
||||
)
|
||||
final_metadata: Final = (
|
||||
{
|
||||
**{key: metadata_value for key, metadata_value in metadata.items() if key != SCIM_ENTERPRISE_METADATA_KEY},
|
||||
|
|
@ -2605,7 +2624,7 @@ def _apply_patch_ops(
|
|||
else {**metadata, "scim_metadata": scim_metadata}
|
||||
)
|
||||
update_data["metadata"] = final_metadata
|
||||
return update_data, final_team_set
|
||||
return update_data, final_teams
|
||||
|
||||
|
||||
def _is_user_not_in_team_error(exc: HTTPException) -> bool:
|
||||
|
|
@ -2758,7 +2777,7 @@ async def patch_user(
|
|||
|
||||
prev_active: Final = _user_scim_active(existing_user)
|
||||
|
||||
update_data, final_team_set = _apply_patch_ops(
|
||||
update_data, final_teams = _apply_patch_ops(
|
||||
existing_user=existing_user,
|
||||
patch_ops=patch_ops,
|
||||
)
|
||||
|
|
@ -2770,15 +2789,15 @@ async def patch_user(
|
|||
await _handle_team_membership_changes(
|
||||
user_id=user_id,
|
||||
existing_teams=existing_user.teams or [],
|
||||
new_teams=list(final_team_set),
|
||||
new_teams=list(final_teams),
|
||||
)
|
||||
|
||||
update_data["teams"] = list(final_team_set)
|
||||
update_data["teams"] = list(final_teams)
|
||||
|
||||
admin_group: Final = await _get_scim_admin_group()
|
||||
if admin_group is not None:
|
||||
update_data["user_role"] = _resolve_scim_user_role(
|
||||
await _scim_groups_from_team_ids(prisma_client, list(final_team_set)),
|
||||
await _scim_groups_from_team_ids(prisma_client, list(final_teams)),
|
||||
admin_group,
|
||||
_default_scim_user_role(),
|
||||
)
|
||||
|
|
|
|||
|
|
@ -8,7 +8,6 @@ Has all /sso/* routes
|
|||
/sso/debug/callback - returns the OpenID object returned by the SSO provider
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import base64
|
||||
import hashlib
|
||||
import inspect
|
||||
|
|
@ -151,6 +150,7 @@ from litellm.proxy.utils import (
|
|||
get_custom_url,
|
||||
get_server_root_path,
|
||||
)
|
||||
from litellm.repositories.chunked_in import find_many_in
|
||||
from litellm.repositories.prisma_protocols import TableActions
|
||||
from litellm.repositories.table_repositories import SSOConfigRepository
|
||||
from litellm.repositories.team_repository import TeamRepository
|
||||
|
|
@ -1855,7 +1855,7 @@ 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):
|
||||
"""Create a task for adding a member to a team."""
|
||||
try:
|
||||
member: Final = Member(user_id=user_info.user_id, role="user")
|
||||
|
|
@ -1871,22 +1871,43 @@ async def create_team_member_add_task(team_id, user_info):
|
|||
verbose_proxy_logger.debug("[Non-Blocking] Error trying to add sso user to db: %s", e)
|
||||
|
||||
|
||||
async def add_missing_team_member(user_info: NewUserResponse | LiteLLM_UserTable, sso_teams: list[str]):
|
||||
async def find_existing_team_ids(team_ids: Sequence[str]) -> frozenset[str] | None:
|
||||
from litellm.proxy.proxy_server import prisma_client
|
||||
|
||||
if prisma_client is None:
|
||||
return None
|
||||
try:
|
||||
rows: Final = await find_many_in(TeamRepository(prisma_client, use_writer=True).table, "team_id", team_ids)
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.debug("[Non-Blocking] Error looking up SSO teams: %s", e)
|
||||
return None
|
||||
return frozenset(row.team_id for row in rows)
|
||||
|
||||
|
||||
async def add_missing_team_member(
|
||||
user_info: NewUserResponse | LiteLLM_UserTable,
|
||||
sso_teams: Sequence[str],
|
||||
*,
|
||||
add_team_member: Callable[
|
||||
[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,
|
||||
) -> None:
|
||||
"""
|
||||
- Get missing teams (diff b/w user_info.team_ids and sso_teams)
|
||||
- Add missing user to missing teams
|
||||
- Drop the IdP groups that are not teams with one lookup, since group claims often carry many of them
|
||||
- Add missing user to missing teams, one at a time in sso_teams order, so user.teams keeps the claim's order
|
||||
"""
|
||||
# Handle None as empty list for new users
|
||||
user_teams: Final = user_info.teams if user_info.teams is not None else []
|
||||
missing_teams: Final = set(sso_teams) - set(user_teams)
|
||||
missing_teams_list: Final = list(missing_teams)
|
||||
tasks = []
|
||||
tasks = [create_team_member_add_task(team_id, user_info) for team_id in missing_teams_list]
|
||||
|
||||
try:
|
||||
await asyncio.gather(*tasks)
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.debug("[Non-Blocking] Error trying to add sso user to db: %s", e)
|
||||
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)
|
||||
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)
|
||||
)
|
||||
for team_id in teams_to_add:
|
||||
await add_team_member(team_id, user_info)
|
||||
|
||||
|
||||
def get_disabled_non_admin_personal_key_creation():
|
||||
|
|
@ -2451,7 +2472,9 @@ async def fetch_cli_sso_team_details(
|
|||
except Exception as e:
|
||||
verbose_proxy_logger.error("Error fetching team details for CLI SSO session: %s", e)
|
||||
return None
|
||||
return tuple(_cli_sso_team_detail(team_row.model_dump()) for team_row in prisma_teams)
|
||||
fetched_details: Final = tuple(_cli_sso_team_detail(team_row.model_dump()) for team_row in prisma_teams)
|
||||
details_by_team_id: Final = MappingProxyType({detail.team_id: detail for detail in fetched_details})
|
||||
return tuple(details_by_team_id[team_id] for team_id in dict.fromkeys(teams) if team_id in details_by_team_id)
|
||||
|
||||
|
||||
def _cli_sso_session_teams(team_details: Sequence[CliSsoTeamDetail]) -> list[str]:
|
||||
|
|
|
|||
|
|
@ -636,3 +636,55 @@ def test_device_flow_mints_the_session_for_the_team_the_developer_picks_in_the_b
|
|||
served_models: Final = {model.id for model in sdk.models.list()}
|
||||
assert served_models == {team_models[chosen_team]}
|
||||
assert len(identity_provider.drain()) == 3
|
||||
|
||||
|
||||
@pytest.mark.timeout(300)
|
||||
def test_gateway_login_lists_and_defaults_to_the_first_team_on_the_user_record_after_a_write_to_its_row(
|
||||
enabled_gateway: Gateway, identity_provider: Wire
|
||||
) -> None:
|
||||
run_id: Final = uuid.uuid4().hex
|
||||
with enabled_gateway.scenario() as scenario:
|
||||
user_id: Final = scenario.user(user_id=_USER_ID, user_email=_USER_EMAIL, user_role="internal_user")
|
||||
first_team: Final = scenario.team(team_id=f"zeta-{run_id}")
|
||||
second_team: Final = scenario.team(team_id=f"alpha-{run_id}")
|
||||
for team_id in (first_team, second_team):
|
||||
enabled_gateway.post(
|
||||
"/team/member_add", {"team_id": team_id, "member": {"role": "user", "user_id": user_id}}
|
||||
)
|
||||
enabled_gateway.post("/team/update", {"team_id": first_team, "metadata": {"updated_by_run": run_id}})
|
||||
assert enabled_gateway.rows('SELECT teams FROM "LiteLLM_UserTable" WHERE user_id = %s', (user_id,)) == [
|
||||
{"teams": [first_team, second_team]}
|
||||
]
|
||||
|
||||
authorization_response: Final = enabled_gateway.client.post(
|
||||
"/claude_code_gateway/oauth/device_authorization", data={"client_id": _CLIENT_ID}
|
||||
)
|
||||
assert authorization_response.status_code == 200, authorization_response.text
|
||||
authorization: Final = _DeviceAuthorization.model_validate_json(authorization_response.content)
|
||||
verification: Final = enabled_gateway.client.get(authorization.verification_uri, follow_redirects=False)
|
||||
assert verification.status_code == 303, verification.text
|
||||
identity_provider_redirect: Final = enabled_gateway.client.get(
|
||||
verification.headers["location"], follow_redirects=False
|
||||
)
|
||||
assert identity_provider_redirect.status_code == 302, identity_provider_redirect.text
|
||||
callback: Final = enabled_gateway.client.get(
|
||||
identity_provider_redirect.headers["location"], follow_redirects=False
|
||||
)
|
||||
assert callback.status_code == 200, callback.text
|
||||
assert tuple(_TEAM_OPTION.findall(callback.text)) == (first_team, second_team), callback.text
|
||||
browser_token: Final = _BROWSER_TOKEN.search(callback.text)
|
||||
assert browser_token is not None, callback.text
|
||||
completion_without_a_pick: Final = enabled_gateway.client.post(
|
||||
f"/sso/cli/complete/{authorization.device_code.partition('.')[0]}",
|
||||
data={
|
||||
"user_code": authorization.user_code,
|
||||
"browser_complete_token": html.unescape(browser_token.group(1)),
|
||||
},
|
||||
)
|
||||
assert completion_without_a_pick.status_code == 200, completion_without_a_pick.text
|
||||
token_response: Final = _device_code_grant(enabled_gateway, authorization.device_code)
|
||||
assert token_response.status_code == 200, token_response.text
|
||||
token: Final = _AccessToken.model_validate_json(token_response.content)
|
||||
identity_provider.drain()
|
||||
|
||||
assert (token.user_id, token.team_id) == (user_id, first_team)
|
||||
|
|
|
|||
|
|
@ -913,6 +913,31 @@ async def test_sync_user_role_and_teams_cache_invalidation_on_team_change():
|
|||
assert call_kwargs.kwargs["model_type"] == LiteLLM_UserTable
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_sync_user_role_and_teams_keeps_record_order_and_adds_claim_teams_in_claim_order():
|
||||
jwt_handler: Final = JWTHandler()
|
||||
jwt_handler.update_environment(
|
||||
prisma_client=None,
|
||||
user_api_key_cache=AsyncMock(),
|
||||
litellm_jwtauth=LiteLLM_JWTAuth(team_ids_jwt_field="groups", sync_user_role_and_teams=True),
|
||||
)
|
||||
cache: Final = AsyncMock()
|
||||
user: Final = LiteLLM_UserTable(user_id="u1", teams=["team-zeta", "team-kappa", "team-alpha"])
|
||||
token: Final = {"groups": ["team-mu", "team-alpha", "team-beta", "team-mu", "team-zeta", "team-omega"]}
|
||||
|
||||
with patch(
|
||||
"litellm.proxy.management_endpoints.scim.scim_v2.patch_team_membership",
|
||||
new_callable=AsyncMock,
|
||||
) as mock_patch:
|
||||
await JWTAuthManager.sync_user_role_and_teams(jwt_handler, token, user, AsyncMock(), user_api_key_cache=cache)
|
||||
|
||||
assert mock_patch.await_args.kwargs["teams_ids_to_add_user_to"] == ["team-mu", "team-beta", "team-omega"]
|
||||
assert mock_patch.await_args.kwargs["teams_ids_to_remove_user_from"] == ["team-kappa"]
|
||||
expected_teams: Final = ["team-zeta", "team-alpha", "team-mu", "team-beta", "team-omega"]
|
||||
assert user.teams == expected_teams
|
||||
assert cache.async_set_cache.await_args.kwargs["value"].teams == expected_teams
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_sync_user_role_and_teams_no_cache_write_when_nothing_changes():
|
||||
"""Test that cache is NOT written when role and teams already match."""
|
||||
|
|
|
|||
|
|
@ -502,9 +502,9 @@ def test_apply_patch_ops_remove_group_filtered_path_without_value():
|
|||
Operations=[SCIMPatchOperation(op="remove", path='groups[value eq "team-1"]')]
|
||||
)
|
||||
|
||||
_, final_team_set = _apply_patch_ops(existing_user=user, patch_ops=patch_ops)
|
||||
_, final_teams = _apply_patch_ops(existing_user=user, patch_ops=patch_ops)
|
||||
|
||||
assert final_team_set == {"team-2"}
|
||||
assert final_teams == ("team-2",)
|
||||
|
||||
|
||||
def test_apply_patch_ops_add_group_filtered_path_without_value():
|
||||
|
|
@ -519,9 +519,9 @@ def test_apply_patch_ops_add_group_filtered_path_without_value():
|
|||
Operations=[SCIMPatchOperation(op="add", path="groups[value eq 'team-3']")]
|
||||
)
|
||||
|
||||
_, final_team_set = _apply_patch_ops(existing_user=user, patch_ops=patch_ops)
|
||||
_, final_teams = _apply_patch_ops(existing_user=user, patch_ops=patch_ops)
|
||||
|
||||
assert final_team_set == {"team-1", "team-3"}
|
||||
assert final_teams == ("team-1", "team-3")
|
||||
|
||||
|
||||
def test_apply_patch_ops_replace_groups_empty_value_does_not_use_path_filter():
|
||||
|
|
@ -539,9 +539,70 @@ def test_apply_patch_ops_replace_groups_empty_value_does_not_use_path_filter():
|
|||
]
|
||||
)
|
||||
|
||||
_, final_team_set = _apply_patch_ops(existing_user=user, patch_ops=patch_ops)
|
||||
_, final_teams = _apply_patch_ops(existing_user=user, patch_ops=patch_ops)
|
||||
|
||||
assert final_team_set == set()
|
||||
assert final_teams == ()
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"operations,expected_teams",
|
||||
[
|
||||
pytest.param(
|
||||
[
|
||||
SCIMPatchOperation(op="add", path="groups", value=[{"value": "team-mu"}, {"value": "team-beta"}]),
|
||||
SCIMPatchOperation(op="remove", path='groups[value eq "team-kappa"]'),
|
||||
SCIMPatchOperation(op="add", path="groups", value=[{"value": "team-alpha"}, {"value": "team-omega"}]),
|
||||
],
|
||||
("team-zeta", "team-alpha", "team-delta", "team-mu", "team-beta", "team-omega"),
|
||||
id="adds_follow_the_record_in_op_order",
|
||||
),
|
||||
pytest.param(
|
||||
[
|
||||
SCIMPatchOperation(
|
||||
op="replace",
|
||||
path="groups",
|
||||
value=[{"value": "team-omega"}, {"value": "team-delta"}, {"value": "team-beta"}],
|
||||
)
|
||||
],
|
||||
("team-delta", "team-omega", "team-beta"),
|
||||
id="replace_keeps_teams_the_user_has_in_record_order_and_adds_new_ones_in_request_order",
|
||||
),
|
||||
pytest.param(
|
||||
[
|
||||
SCIMPatchOperation(
|
||||
op="replace",
|
||||
path="groups",
|
||||
value=[{"value": "team-omega"}, {"value": "team-delta"}, {"value": "team-beta"}],
|
||||
),
|
||||
SCIMPatchOperation(op="add", path="groups", value=[{"value": "team-mu"}]),
|
||||
SCIMPatchOperation(op="remove", path='groups[value eq "team-delta"]'),
|
||||
],
|
||||
("team-delta", "team-omega", "team-beta"),
|
||||
id="a_replace_wins_over_later_group_ops",
|
||||
),
|
||||
pytest.param(
|
||||
[
|
||||
SCIMPatchOperation(op="remove", path='groups[value eq "team-zeta"]'),
|
||||
SCIMPatchOperation(op="add", path="groups", value=[{"value": "team-zeta"}]),
|
||||
],
|
||||
("team-zeta", "team-alpha", "team-kappa", "team-delta"),
|
||||
id="removing_and_re_adding_a_team_keeps_its_place",
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_apply_patch_ops_keeps_team_order_so_the_first_team_stays_first(
|
||||
operations: list[SCIMPatchOperation], expected_teams: tuple[str, ...]
|
||||
):
|
||||
user: Final = LiteLLM_UserTable(
|
||||
user_id="user-order",
|
||||
user_email="order@example.com",
|
||||
teams=["team-zeta", "team-alpha", "team-kappa", "team-delta"],
|
||||
metadata={},
|
||||
)
|
||||
|
||||
_, final_teams = _apply_patch_ops(existing_user=user, patch_ops=SCIMPatchOp(Operations=operations))
|
||||
|
||||
assert final_teams == expected_teams
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
|
|||
|
|
@ -25,6 +25,7 @@ from litellm.proxy._types import (
|
|||
NewUserResponse,
|
||||
ProxyErrorTypes,
|
||||
ProxyException,
|
||||
TeamMemberAddRequest,
|
||||
UserAPIKeyAuth,
|
||||
)
|
||||
from litellm.proxy.management_endpoints.scim.scim_v2 import (
|
||||
|
|
@ -941,6 +942,57 @@ async def test_handle_existing_user_by_email_syncs_roster_and_dedups_teams(mocke
|
|||
assert update_calls[0].kwargs["data"]["teams"] == ["team-a", "team-b"]
|
||||
|
||||
|
||||
class _EmailMatchedUserRow:
|
||||
def __init__(self, user: LiteLLM_UserTable) -> None:
|
||||
self.user = user
|
||||
|
||||
async def find_first(self, *, where: dict[str, object]) -> LiteLLM_UserTable:
|
||||
return self.user
|
||||
|
||||
async def update(self, *, where: dict[str, object], data: dict[str, object]) -> LiteLLM_UserTable:
|
||||
self.user = self.user.model_copy(update=data)
|
||||
return self.user
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_handle_existing_user_by_email_keeps_record_team_order(mocker: MockerFixture):
|
||||
user_rows: Final = _EmailMatchedUserRow(
|
||||
LiteLLM_UserTable(
|
||||
user_id="same-id",
|
||||
user_email="member@example.com",
|
||||
user_alias="Member",
|
||||
teams=["team-zeta", "team-alpha", "team-kappa"],
|
||||
metadata={},
|
||||
)
|
||||
)
|
||||
mock_prisma_client: Final = mocker.MagicMock()
|
||||
mock_prisma_client.db = mocker.MagicMock()
|
||||
mock_prisma_client.db.litellm_usertable = user_rows
|
||||
|
||||
mocker.patch( # test-quality-ok: roster helpers are module-level, not injectable into the helper
|
||||
"litellm.proxy.management_endpoints.scim.scim_v2.ScimTransformations.transform_litellm_user_to_scim_user",
|
||||
AsyncMock(return_value=None),
|
||||
)
|
||||
mocker.patch( # test-quality-ok: roster helpers are module-level, not injectable into the helper
|
||||
"litellm.proxy.management_endpoints.scim.scim_v2._handle_team_membership_changes",
|
||||
AsyncMock(),
|
||||
)
|
||||
|
||||
await UserProvisionerHelpers.handle_existing_user_by_email(
|
||||
prisma_client=mock_prisma_client,
|
||||
new_user_request=NewUserRequest(
|
||||
user_id="same-id",
|
||||
user_email="member@example.com",
|
||||
user_alias="Member",
|
||||
teams=["team-mu", "team-kappa", "team-beta", "team-zeta"],
|
||||
metadata={},
|
||||
auto_create_key=False,
|
||||
),
|
||||
)
|
||||
|
||||
assert user_rows.user.teams == ["team-zeta", "team-kappa", "team-mu", "team-beta"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_handle_existing_user_by_email_without_teams_preserves_memberships(mocker):
|
||||
"""Adoption via POST /Users without ``groups`` must keep the user's existing teams.
|
||||
|
|
@ -1290,6 +1342,49 @@ async def test_handle_team_membership_changes_add_and_remove(mocker):
|
|||
assert call_args[1]["teams_ids_to_remove_user_from"] == ["team1"]
|
||||
|
||||
|
||||
class _TeamRosterRecord:
|
||||
def __init__(self, missing_team: str) -> None:
|
||||
self.missing_team = missing_team
|
||||
self.teams: tuple[str, ...] = ()
|
||||
|
||||
async def team_member_add(self, *, data: TeamMemberAddRequest, user_api_key_dict: UserAPIKeyAuth) -> None:
|
||||
if data.team_id == self.missing_team:
|
||||
raise HTTPException(status_code=404, detail={"error": "Team not found"})
|
||||
self.teams = (*self.teams, data.team_id)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_handle_team_membership_changes_adds_teams_in_request_order_when_one_fails(mocker: MockerFixture):
|
||||
requested_teams: Final = (
|
||||
"team-zeta",
|
||||
"team-alpha",
|
||||
"team-kappa",
|
||||
"team-kept",
|
||||
"team-mu",
|
||||
"team-missing",
|
||||
"team-beta",
|
||||
"team-omega",
|
||||
"team-delta",
|
||||
"team-sigma",
|
||||
"team-gamma",
|
||||
"team-tau",
|
||||
)
|
||||
roster: Final = _TeamRosterRecord(missing_team="team-missing")
|
||||
mocker.patch( # test-quality-ok: team_member_add is module-level, not injectable into the helper
|
||||
"litellm.proxy.management_endpoints.scim.scim_v2.team_member_add",
|
||||
roster.team_member_add,
|
||||
)
|
||||
|
||||
with pytest.raises(SCIMRosterSyncError):
|
||||
await _handle_team_membership_changes(
|
||||
user_id="test-user",
|
||||
existing_teams=["team-kept"],
|
||||
new_teams=list(requested_teams),
|
||||
)
|
||||
|
||||
assert roster.teams == tuple(team for team in requested_teams if team not in {"team-kept", "team-missing"})
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_user_success(mocker):
|
||||
"""Should successfully update user with PUT request"""
|
||||
|
|
@ -1492,6 +1587,57 @@ async def test_update_user_without_groups_preserves_memberships_and_role(mocker,
|
|||
assert "user_role" not in update_data
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_user_put_keeps_the_record_order_of_teams_the_user_already_has(
|
||||
mocker: MockerFixture, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
from litellm.proxy.proxy_server import proxy_config
|
||||
|
||||
async def no_scim_admin_group() -> dict[str, object]:
|
||||
return {}
|
||||
|
||||
monkeypatch.setattr(proxy_config, "get_config", no_scim_admin_group)
|
||||
monkeypatch.setattr("litellm.default_internal_user_params", None, raising=False)
|
||||
|
||||
existing_user: Final = mocker.MagicMock()
|
||||
existing_user.teams = ["team-zeta", "team-alpha", "team-kappa"]
|
||||
existing_user.metadata = {}
|
||||
scim_user: Final = SCIMUser(
|
||||
schemas=["urn:ietf:params:scim:schemas:core:2.0:User"],
|
||||
userName="order-user",
|
||||
emails=[SCIMUserEmail(value="order@example.com")],
|
||||
groups=[
|
||||
SCIMUserGroup(value="team-mu"),
|
||||
SCIMUserGroup(value="team-kappa"),
|
||||
SCIMUserGroup(value="team-beta"),
|
||||
SCIMUserGroup(value="team-zeta"),
|
||||
],
|
||||
)
|
||||
mock_prisma_client: Final = mocker.MagicMock()
|
||||
mock_prisma_client.db.litellm_usertable.update = AsyncMock(return_value={"user_id": "order-user"})
|
||||
mocker.patch( # test-quality-ok: update_user's collaborators are module-level, not injectable
|
||||
"litellm.proxy.management_endpoints.scim.scim_v2._get_prisma_client_or_raise_exception",
|
||||
AsyncMock(return_value=mock_prisma_client),
|
||||
)
|
||||
mocker.patch( # test-quality-ok: update_user's collaborators are module-level, not injectable
|
||||
"litellm.proxy.management_endpoints.scim.scim_v2._check_user_exists",
|
||||
AsyncMock(return_value=existing_user),
|
||||
)
|
||||
mocker.patch( # test-quality-ok: update_user's collaborators are module-level, not injectable
|
||||
"litellm.proxy.management_endpoints.scim.scim_v2.ScimTransformations.transform_litellm_user_to_scim_user",
|
||||
AsyncMock(return_value=scim_user),
|
||||
)
|
||||
mocker.patch( # test-quality-ok: roster writes are module-level, not injectable
|
||||
"litellm.proxy.management_endpoints.scim.scim_v2.patch_team_membership",
|
||||
AsyncMock(),
|
||||
)
|
||||
|
||||
await update_user(user_id="order-user", user=scim_user)
|
||||
|
||||
update_data: Final = mock_prisma_client.db.litellm_usertable.update.call_args.kwargs["data"]
|
||||
assert update_data["teams"] == ["team-zeta", "team-kappa", "team-mu", "team-beta"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_user_not_found(mocker):
|
||||
"""Should raise 404 when user doesn't exist"""
|
||||
|
|
|
|||
|
|
@ -4195,6 +4195,65 @@ async def test_add_new_user_to_default_team_string_teams_have_no_member_budget(m
|
|||
assert mock_add.call_args.kwargs["team_id"] == "string-team"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_add_new_user_to_default_team_joins_teams_in_the_order_given_and_past_a_failing_team(
|
||||
mocker: MockerFixture,
|
||||
) -> None:
|
||||
from litellm.proxy.management_endpoints.internal_user_endpoints import (
|
||||
add_new_user_to_default_team,
|
||||
)
|
||||
|
||||
requested: Final = ("team-zeta", "team-broken", "team-mu", "team-alpha")
|
||||
seconds_per_team: Final = {"team-zeta": 0.03, "team-broken": 0.0, "team-mu": 0.02, "team-alpha": 0.0}
|
||||
finished: Final = mocker.MagicMock()
|
||||
|
||||
async def join_team(*, team_id: str, **_: object) -> None:
|
||||
await asyncio.sleep(seconds_per_team[team_id])
|
||||
if team_id == "team-broken":
|
||||
raise RuntimeError("team row locked")
|
||||
finished(team_id)
|
||||
|
||||
mocker.patch(
|
||||
"litellm.proxy.management_endpoints.internal_user_endpoints._add_user_to_team",
|
||||
side_effect=join_team,
|
||||
)
|
||||
|
||||
await add_new_user_to_default_team(
|
||||
user_id="scim-user",
|
||||
user_email=None,
|
||||
user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN),
|
||||
teams=list(requested),
|
||||
prisma_client=mocker.MagicMock(),
|
||||
)
|
||||
|
||||
assert [call.args[0] for call in finished.call_args_list] == ["team-zeta", "team-mu", "team-alpha"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_add_new_user_to_default_team_rejects_an_invalid_team_before_joining_any(
|
||||
mocker: MockerFixture,
|
||||
) -> None:
|
||||
from litellm.proxy.management_endpoints.internal_user_endpoints import (
|
||||
add_new_user_to_default_team,
|
||||
)
|
||||
|
||||
join_team: Final = mocker.patch(
|
||||
"litellm.proxy.management_endpoints.internal_user_endpoints._add_user_to_team",
|
||||
new_callable=mocker.AsyncMock,
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match="Invalid team type"):
|
||||
await add_new_user_to_default_team(
|
||||
user_id="new-user",
|
||||
user_email=None,
|
||||
user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN),
|
||||
teams=["team-zeta", {"team_id": "team-mu"}],
|
||||
prisma_client=mocker.MagicMock(),
|
||||
)
|
||||
|
||||
join_team.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_add_user_to_team_logs_unknown_team_at_error(mocker, caplog):
|
||||
"""A default team that no longer exists makes every membership write 404.
|
||||
|
|
|
|||
|
|
@ -4049,8 +4049,8 @@ class TestCLIKeyRegenerationFlow:
|
|||
assert (
|
||||
'<select id="team_id" name="team_id" required>'
|
||||
'<option value="" disabled selected>Choose a team</option>'
|
||||
'<option value="team-ops">Ops <eu> (team-ops)</option>'
|
||||
'<option value="team-claude">team-claude</option>'
|
||||
'<option value="team-ops">Ops <eu> (team-ops)</option>'
|
||||
'<option value="team-ops-us">Ops <eu> (team-ops-us)</option>'
|
||||
"</select>"
|
||||
) in page
|
||||
|
|
@ -4530,6 +4530,39 @@ class TestCLIKeyRegenerationFlow:
|
|||
== ()
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_fetch_cli_sso_team_details_follows_the_user_record_order(self):
|
||||
"""The CLI login binds to the first detail, so it must be the first team on the user record.
|
||||
|
||||
Postgres returns the rows in whatever order the table holds them, which every write to a
|
||||
team row changes, so the bound team used to flip between sign-ins.
|
||||
"""
|
||||
from litellm.proxy.management_endpoints.ui_sso import (
|
||||
fetch_cli_sso_team_details,
|
||||
)
|
||||
|
||||
def team_row(team_id: str) -> MagicMock:
|
||||
row: Final = MagicMock()
|
||||
row.model_dump.return_value = {
|
||||
"team_id": team_id,
|
||||
"team_alias": None,
|
||||
"models": [],
|
||||
"litellm_model_table": None,
|
||||
}
|
||||
return row
|
||||
|
||||
prisma_client: Final = MagicMock()
|
||||
prisma_client.db.litellm_teamtable.find_many = AsyncMock(
|
||||
return_value=[team_row("team-alpha"), team_row("team-zeta")]
|
||||
)
|
||||
|
||||
details: Final = await fetch_cli_sso_team_details(
|
||||
prisma_client=prisma_client,
|
||||
teams=["team-zeta", "team-gone", "team-alpha", "team-zeta"],
|
||||
)
|
||||
|
||||
assert [detail.team_id for detail in details] == ["team-zeta", "team-alpha"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cli_poll_key_mints_jwt_with_selected_team_grants(self):
|
||||
"""The selected team's grants must reach the mint, not just its alias."""
|
||||
|
|
@ -6692,72 +6725,57 @@ class TestPKCEFunctionality:
|
|||
class TestAddMissingTeamMember:
|
||||
"""Tests for the add_missing_team_member function"""
|
||||
|
||||
@staticmethod
|
||||
def _recording_add() -> tuple[
|
||||
list[tuple[str, str | None]],
|
||||
Callable[[str, NewUserResponse | LiteLLM_UserTable], Awaitable[None]],
|
||||
]:
|
||||
added: Final[list[tuple[str, str | None]]] = []
|
||||
|
||||
async def add(team_id: str, user_info: NewUserResponse | LiteLLM_UserTable) -> None:
|
||||
added.append((team_id, user_info.user_id))
|
||||
|
||||
return added, add
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_add_missing_team_member_with_new_user_response_teams_none(self):
|
||||
"""
|
||||
Bug reproduction: When a NewUserResponse has teams=None (new SSO user),
|
||||
add_missing_team_member() should still add the user to the SSO teams.
|
||||
|
||||
Currently FAILS: The function returns early when teams is None.
|
||||
"""
|
||||
"""A new SSO user's NewUserResponse has teams=None, and the user must still be added to the SSO teams."""
|
||||
from litellm.proxy._types import NewUserResponse
|
||||
from litellm.proxy.management_endpoints.ui_sso import add_missing_team_member
|
||||
|
||||
# Simulate a new SSO user - NewUserResponse has teams=None by default
|
||||
new_user = NewUserResponse(
|
||||
user_id="new-sso-user-123",
|
||||
key="sk-xxxxx",
|
||||
teams=None, # This is the default for NewUserResponse
|
||||
new_user = NewUserResponse(user_id="new-sso-user-123", key="sk-xxxxx", teams=None)
|
||||
added, add = self._recording_add()
|
||||
|
||||
await add_missing_team_member(
|
||||
user_info=new_user,
|
||||
sso_teams=["team-from-entra-1", "team-from-entra-2"],
|
||||
add_team_member=add,
|
||||
)
|
||||
|
||||
sso_teams = ["team-from-entra-1", "team-from-entra-2"]
|
||||
|
||||
with patch(
|
||||
"litellm.proxy.management_endpoints.ui_sso.create_team_member_add_task"
|
||||
) as mock_add_task:
|
||||
mock_add_task.return_value = AsyncMock()
|
||||
|
||||
await add_missing_team_member(user_info=new_user, sso_teams=sso_teams)
|
||||
|
||||
# Bug: This assertion currently FAILS - no teams are added
|
||||
# because function returns early when teams is None
|
||||
assert (
|
||||
mock_add_task.call_count == 2
|
||||
), f"Expected 2 calls to add user to teams, but got {mock_add_task.call_count}"
|
||||
called_team_ids = [call.args[0] for call in mock_add_task.call_args_list]
|
||||
assert set(called_team_ids) == {
|
||||
"team-from-entra-1",
|
||||
"team-from-entra-2",
|
||||
}
|
||||
assert added == [
|
||||
("team-from-entra-1", "new-sso-user-123"),
|
||||
("team-from-entra-2", "new-sso-user-123"),
|
||||
]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_add_missing_team_member_with_litellm_user_table_empty_teams(self):
|
||||
"""
|
||||
Control test: When a LiteLLM_UserTable has teams=[] (existing user, no teams),
|
||||
add_missing_team_member() should add the user to SSO teams.
|
||||
|
||||
This test PASSES because LiteLLM_UserTable defaults teams to [] not None.
|
||||
"""
|
||||
"""An existing user with teams=[] is added to every SSO team."""
|
||||
from litellm.proxy._types import LiteLLM_UserTable
|
||||
from litellm.proxy.management_endpoints.ui_sso import add_missing_team_member
|
||||
|
||||
# Existing user has teams=[] by default (not None)
|
||||
existing_user = LiteLLM_UserTable(
|
||||
user_id="existing-user-456",
|
||||
teams=[], # Empty list, not None
|
||||
existing_user = LiteLLM_UserTable(user_id="existing-user-456", teams=[])
|
||||
added, add = self._recording_add()
|
||||
|
||||
await add_missing_team_member(
|
||||
user_info=existing_user,
|
||||
sso_teams=["team-from-entra-1", "team-from-entra-2"],
|
||||
add_team_member=add,
|
||||
)
|
||||
|
||||
sso_teams = ["team-from-entra-1", "team-from-entra-2"]
|
||||
|
||||
with patch(
|
||||
"litellm.proxy.management_endpoints.ui_sso.create_team_member_add_task"
|
||||
) as mock_add_task:
|
||||
mock_add_task.return_value = AsyncMock()
|
||||
|
||||
await add_missing_team_member(user_info=existing_user, sso_teams=sso_teams)
|
||||
|
||||
# This PASSES - teams are added because teams=[] not None
|
||||
assert mock_add_task.call_count == 2
|
||||
assert added == [
|
||||
("team-from-entra-1", "existing-user-456"),
|
||||
("team-from-entra-2", "existing-user-456"),
|
||||
]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_add_user_to_teams_from_sso_response_new_user(self):
|
||||
|
|
@ -6800,96 +6818,160 @@ class TestAddMissingTeamMember:
|
|||
|
||||
@pytest.mark.asyncio
|
||||
async def test_sso_first_login_full_flow_adds_user_to_teams(self):
|
||||
"""
|
||||
End-to-end test: Simulates complete first-time SSO login with Entra groups.
|
||||
Verifies teams are created AND user is added as a member.
|
||||
"""
|
||||
"""A first-time SSO login with Entra groups adds the new user as a member of every group's team."""
|
||||
from litellm.proxy._types import NewUserResponse
|
||||
from litellm.proxy.management_endpoints.ui_sso import add_missing_team_member
|
||||
|
||||
team_member_calls = []
|
||||
new_user = NewUserResponse(user_id="first-time-sso-user", key="sk-xxxxx", teams=None)
|
||||
added, add = self._recording_add()
|
||||
|
||||
async def track_team_member_add(team_id, user_info):
|
||||
team_member_calls.append({"team_id": team_id, "user_id": user_info.user_id})
|
||||
|
||||
# New SSO user with Entra groups
|
||||
new_user = NewUserResponse(
|
||||
user_id="first-time-sso-user",
|
||||
key="sk-xxxxx",
|
||||
teams=None, # The problematic default
|
||||
await add_missing_team_member(
|
||||
user_info=new_user,
|
||||
sso_teams=["entra-team-alpha", "entra-team-beta"],
|
||||
add_team_member=add,
|
||||
)
|
||||
|
||||
sso_teams = ["entra-team-alpha", "entra-team-beta"]
|
||||
|
||||
with patch(
|
||||
"litellm.proxy.management_endpoints.ui_sso.create_team_member_add_task",
|
||||
side_effect=track_team_member_add,
|
||||
):
|
||||
await add_missing_team_member(user_info=new_user, sso_teams=sso_teams)
|
||||
|
||||
# Bug: With current code, team_member_calls will be empty
|
||||
# After fix: Should have 2 entries
|
||||
assert (
|
||||
len(team_member_calls) == 2
|
||||
), f"Expected 2 teams to be added, but got {len(team_member_calls)}"
|
||||
assert {c["team_id"] for c in team_member_calls} == {
|
||||
"entra-team-alpha",
|
||||
"entra-team-beta",
|
||||
}
|
||||
assert all(c["user_id"] == "first-time-sso-user" for c in team_member_calls)
|
||||
assert added == [
|
||||
("entra-team-alpha", "first-time-sso-user"),
|
||||
("entra-team-beta", "first-time-sso-user"),
|
||||
]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"user_info_factory,teams_value,expected_teams_added",
|
||||
"user_info_factory,expected_teams_added",
|
||||
[
|
||||
# Bug case: NewUserResponse with teams=None
|
||||
pytest.param(
|
||||
lambda uid: NewUserResponse(user_id=uid, key="sk-xxx", teams=None),
|
||||
None,
|
||||
["team-1", "team-2"], # Should still add teams
|
||||
["team-1", "team-2"],
|
||||
id="new_user_teams_none",
|
||||
),
|
||||
# Working case: LiteLLM_UserTable with teams=[]
|
||||
pytest.param(
|
||||
lambda uid: LiteLLM_UserTable(user_id=uid, teams=[]),
|
||||
[],
|
||||
["team-1", "team-2"],
|
||||
id="existing_user_empty_teams",
|
||||
),
|
||||
# Existing user with some teams already
|
||||
pytest.param(
|
||||
lambda uid: LiteLLM_UserTable(user_id=uid, teams=["team-1"]),
|
||||
["team-1"],
|
||||
["team-2"], # Only missing team should be added
|
||||
["team-2"],
|
||||
id="existing_user_partial_teams",
|
||||
),
|
||||
pytest.param(
|
||||
lambda uid: LiteLLM_UserTable(user_id=uid, teams=["team-2", "team-1"]),
|
||||
[],
|
||||
id="existing_user_all_teams",
|
||||
),
|
||||
],
|
||||
)
|
||||
async def test_add_missing_team_member_handles_all_user_types(
|
||||
self, user_info_factory, teams_value, expected_teams_added
|
||||
self, user_info_factory, expected_teams_added
|
||||
):
|
||||
"""
|
||||
Parametrized test ensuring add_missing_team_member works for all user types.
|
||||
"""
|
||||
"""Only the SSO teams the user is not already on are added."""
|
||||
from litellm.proxy.management_endpoints.ui_sso import add_missing_team_member
|
||||
|
||||
user_info = user_info_factory("test-user-id")
|
||||
sso_teams = ["team-1", "team-2"]
|
||||
added, add = self._recording_add()
|
||||
|
||||
added_teams = []
|
||||
await add_missing_team_member(
|
||||
user_info=user_info_factory("test-user-id"),
|
||||
sso_teams=["team-1", "team-2"],
|
||||
add_team_member=add,
|
||||
)
|
||||
|
||||
async def mock_create_task(team_id, user):
|
||||
added_teams.append(team_id)
|
||||
assert added == [(team_id, "test-user-id") for team_id in expected_teams_added]
|
||||
|
||||
with patch(
|
||||
"litellm.proxy.management_endpoints.ui_sso.create_team_member_add_task",
|
||||
side_effect=mock_create_task,
|
||||
):
|
||||
await add_missing_team_member(user_info=user_info, sso_teams=sso_teams)
|
||||
@pytest.mark.asyncio
|
||||
async def test_add_missing_team_member_lands_teams_in_claim_order(self):
|
||||
"""user.teams is appended to as each add finishes, and the first team on it is the one a CLI
|
||||
login binds to, so the adds must finish in the claim's order even when an earlier add is slower.
|
||||
|
||||
assert set(added_teams) == set(
|
||||
expected_teams_added
|
||||
), f"Expected teams {expected_teams_added}, but got {added_teams}"
|
||||
The adds used to run concurrently from a set, so they finished in an arbitrary order.
|
||||
"""
|
||||
from litellm.proxy._types import LiteLLM_UserTable
|
||||
from litellm.proxy.management_endpoints.ui_sso import add_missing_team_member
|
||||
|
||||
claim: Final = ["team-zeta", "team-alpha", "team-held", "team-zeta", "team-mu"]
|
||||
yields_before_finishing: Final = {"team-zeta": 3, "team-alpha": 2, "team-mu": 1}
|
||||
finished: Final[list[str]] = []
|
||||
|
||||
async def slower_when_earlier(team_id: str, user_info: NewUserResponse | LiteLLM_UserTable) -> None:
|
||||
for _ in range(yields_before_finishing[team_id]):
|
||||
await asyncio.sleep(0)
|
||||
finished.append(team_id)
|
||||
|
||||
await add_missing_team_member(
|
||||
user_info=LiteLLM_UserTable(user_id="group-sync-user", teams=["team-held"]),
|
||||
sso_teams=claim,
|
||||
add_team_member=slower_when_earlier,
|
||||
)
|
||||
|
||||
assert finished == ["team-zeta", "team-alpha", "team-mu"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"existing,expected_teams_added",
|
||||
[
|
||||
pytest.param(frozenset({"team-zeta", "team-alpha"}), ["team-zeta", "team-alpha"], id="only_teams"),
|
||||
pytest.param(None, ["grp-everyone", "team-zeta", "grp-vpn", "team-alpha"], id="lookup_failed"),
|
||||
],
|
||||
)
|
||||
async def test_add_missing_team_member_looks_up_missing_groups_once_and_adds_only_teams(
|
||||
self, existing: frozenset[str] | None, expected_teams_added: list[str]
|
||||
):
|
||||
from litellm.proxy.management_endpoints.ui_sso import add_missing_team_member
|
||||
|
||||
added, add = self._recording_add()
|
||||
lookups: Final[list[tuple[str, ...]]] = []
|
||||
|
||||
async def existing_team_ids(team_ids: Sequence[str]) -> frozenset[str] | None:
|
||||
lookups.append(tuple(team_ids))
|
||||
return existing
|
||||
|
||||
await add_missing_team_member(
|
||||
user_info=LiteLLM_UserTable(user_id="group-sync-user", teams=["team-held"]),
|
||||
sso_teams=["grp-everyone", "team-zeta", "team-held", "grp-vpn", "team-alpha"],
|
||||
add_team_member=add,
|
||||
existing_team_ids=existing_team_ids,
|
||||
)
|
||||
|
||||
assert lookups == [("grp-everyone", "team-zeta", "grp-vpn", "team-alpha")]
|
||||
assert added == [(team_id, "group-sync-user") for team_id in expected_teams_added]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_add_missing_team_member_skips_the_lookup_when_every_group_is_held(self):
|
||||
from litellm.proxy.management_endpoints.ui_sso import add_missing_team_member
|
||||
|
||||
added, add = self._recording_add()
|
||||
lookup: Final = AsyncMock(return_value=frozenset())
|
||||
|
||||
await add_missing_team_member(
|
||||
user_info=LiteLLM_UserTable(user_id="group-sync-user", teams=["team-alpha", "team-zeta"]),
|
||||
sso_teams=["team-zeta", "team-alpha"],
|
||||
add_team_member=add,
|
||||
existing_team_ids=lookup,
|
||||
)
|
||||
|
||||
lookup.assert_not_awaited()
|
||||
assert added == []
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_find_existing_team_ids_reads_the_writer_and_returns_none_when_the_db_cannot_answer(self):
|
||||
from litellm.proxy.management_endpoints.ui_sso import find_existing_team_ids
|
||||
|
||||
stored_on_writer: Final = ("team-zeta", "team-alpha")
|
||||
|
||||
async def find_many(*, where: Mapping[str, Mapping[str, Sequence[str]]]) -> list[SimpleNamespace]:
|
||||
return [
|
||||
SimpleNamespace(team_id=team_id) for team_id in stored_on_writer if team_id in where["team_id"]["in"]
|
||||
]
|
||||
|
||||
prisma: Final = MagicMock()
|
||||
prisma.writer_db.litellm_teamtable.find_many = AsyncMock(side_effect=find_many)
|
||||
prisma.db.litellm_teamtable.find_many = AsyncMock(return_value=[])
|
||||
with patch("litellm.proxy.proxy_server.prisma_client", prisma):
|
||||
assert await find_existing_team_ids(("grp-vpn", "team-alpha")) == frozenset({"team-alpha"})
|
||||
prisma.writer_db.litellm_teamtable.find_many.side_effect = RuntimeError("db unavailable")
|
||||
assert await find_existing_team_ids(("team-alpha",)) is None
|
||||
with patch("litellm.proxy.proxy_server.prisma_client", None):
|
||||
assert await find_existing_team_ids(("team-alpha",)) is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue