mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-06 07:48:06 +00:00
feat(sso): create every claimed team on login when team_id_upsert is set, on the UI SSO and API JWT paths
This commit is contained in:
parent
0f19833c23
commit
9924cdb433
7 changed files with 280 additions and 49 deletions
|
|
@ -14,7 +14,7 @@ import hashlib
|
|||
import os
|
||||
import re
|
||||
import time
|
||||
from collections.abc import Awaitable, Callable, Collection, Mapping, Sequence
|
||||
from collections.abc import Awaitable, Callable, Collection, Iterable, Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from typing import Final, Literal, NoReturn, Protocol, TypeVar, cast
|
||||
|
||||
|
|
@ -1762,8 +1762,9 @@ class JWTAuthManager:
|
|||
parent_otel_span: Span | None,
|
||||
proxy_logging_obj: ProxyLogging,
|
||||
request_method: str | None = None,
|
||||
team_id_upsert: bool | None = None,
|
||||
) -> tuple[str | None, LiteLLM_TeamTable | None]:
|
||||
"""Find first team with access to the requested model"""
|
||||
"""Find first team with access to the requested model, creating claimed teams when ``team_id_upsert`` is set"""
|
||||
from litellm.proxy.proxy_server import llm_router
|
||||
|
||||
denied_pass_through_team: LiteLLM_TeamTable | None = None
|
||||
|
|
@ -1789,6 +1790,7 @@ class JWTAuthManager:
|
|||
user_api_key_cache=user_api_key_cache,
|
||||
parent_otel_span=parent_otel_span,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
team_id_upsert=team_id_upsert,
|
||||
)
|
||||
|
||||
if team_object is not None:
|
||||
|
|
@ -2213,6 +2215,8 @@ class JWTAuthManager:
|
|||
patch_team_membership,
|
||||
)
|
||||
|
||||
if jwt_handler.litellm_jwtauth.team_id_upsert and user_api_key_cache is not None:
|
||||
await JWTAuthManager._upsert_claimed_teams(teams_to_add, prisma_client, user_api_key_cache)
|
||||
await patch_team_membership(
|
||||
user_id=user_object.user_id,
|
||||
teams_ids_to_add_user_to=list(teams_to_add),
|
||||
|
|
@ -2228,6 +2232,23 @@ class JWTAuthManager:
|
|||
)
|
||||
return
|
||||
|
||||
@staticmethod
|
||||
async def _upsert_claimed_teams(
|
||||
team_ids: Iterable[str], prisma_client: PrismaClient, user_api_key_cache: UserApiKeyCache
|
||||
) -> None:
|
||||
"""Create the teams the user is about to join. A concurrent request may win the create, so the join still runs."""
|
||||
for team_id in sorted(team_ids):
|
||||
try:
|
||||
await get_team_object(
|
||||
team_id=team_id,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
check_db_only=True,
|
||||
team_id_upsert=True,
|
||||
)
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.debug("JWT auth could not create team %s, trying to join it anyway: %s", team_id, e)
|
||||
|
||||
@staticmethod
|
||||
async def _attach_team_from_header_for_admin(
|
||||
admin_result: JWTAuthBuilderResult,
|
||||
|
|
@ -2813,6 +2834,7 @@ class JWTAuthManager:
|
|||
user_api_key_cache=user_api_key_cache,
|
||||
parent_otel_span=parent_otel_span,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
team_id_upsert=team_id_upsert,
|
||||
)
|
||||
|
||||
# The RBAC role-claim path (rbac_role == TEAM) sets team_id without
|
||||
|
|
|
|||
|
|
@ -93,6 +93,6 @@ def get_litellm_user_role(role_str: object) -> LitellmUserRoles | None:
|
|||
|
||||
class CustomOpenID(OpenID):
|
||||
team_ids: list[str]
|
||||
team_id_to_upsert: str | None = None
|
||||
team_id_upsert: bool = False
|
||||
user_role: LitellmUserRoles | None = None
|
||||
extra_fields: dict[str, object] | None = None
|
||||
|
|
|
|||
|
|
@ -1263,17 +1263,18 @@ def generic_response_convertor(
|
|||
team_ids = jwt_handler.get_all_jwt_team_ids(token)
|
||||
all_teams.extend(team_ids)
|
||||
team_id_jwt_field: Final = jwt_handler.litellm_jwtauth.team_id_jwt_field
|
||||
upsert_candidate: Final[object] = (
|
||||
team_id_upsert: Final = bool(jwt_handler.litellm_jwtauth.team_id_upsert)
|
||||
scalar_team_id: Final[object] = (
|
||||
get_nested_value(data=token, key_path=team_id_jwt_field)
|
||||
if jwt_handler.litellm_jwtauth.team_id_upsert and team_id_jwt_field is not None
|
||||
if team_id_upsert and team_id_jwt_field is not None
|
||||
else None
|
||||
)
|
||||
team_id_to_upsert: Final = (
|
||||
upsert_candidate if isinstance(upsert_candidate, str) and upsert_candidate.strip() else None
|
||||
)
|
||||
team_ids_with_upsert: Final = (
|
||||
all_teams if team_id_to_upsert is None or team_id_to_upsert in all_teams else [*all_teams, team_id_to_upsert]
|
||||
team_ids_with_scalar: Final = (
|
||||
all_teams if not isinstance(scalar_team_id, str) or scalar_team_id in all_teams else [*all_teams, scalar_team_id]
|
||||
)
|
||||
team_ids_to_join: Final = [
|
||||
team_id for team_id in team_ids_with_scalar if not (isinstance(team_id, str) and not team_id.strip())
|
||||
]
|
||||
|
||||
# Determine user role based on role_mappings if available
|
||||
# Only apply role_mappings for GENERIC SSO provider
|
||||
|
|
@ -1341,8 +1342,8 @@ def generic_response_convertor(
|
|||
first_name=get_nested_value(response, generic_user_first_name_attribute_name),
|
||||
last_name=get_nested_value(response, generic_user_last_name_attribute_name),
|
||||
provider=get_nested_value(response, generic_provider_attribute_name),
|
||||
team_ids=team_ids_with_upsert,
|
||||
team_id_to_upsert=team_id_to_upsert,
|
||||
team_ids=team_ids_to_join,
|
||||
team_id_upsert=team_id_upsert,
|
||||
user_role=user_role,
|
||||
extra_fields=extra_fields,
|
||||
)
|
||||
|
|
@ -1880,17 +1881,17 @@ async def _upsert_sso_team(team_id: str, prisma_client: PrismaClient, user_api_k
|
|||
async def add_missing_team_member(
|
||||
user_info: NewUserResponse | LiteLLM_UserTable,
|
||||
sso_teams: list[str],
|
||||
team_id_to_upsert: str | None = None,
|
||||
upsert_team: Callable[[str], Awaitable[object]] | None = None,
|
||||
) -> None:
|
||||
"""
|
||||
- Get missing teams (diff b/w user_info.team_ids and sso_teams)
|
||||
- Create each missing team first when ``upsert_team`` is given (``team_id_upsert``)
|
||||
- Add missing user to missing teams
|
||||
"""
|
||||
missing_teams: Final = sorted(set(sso_teams) - set(user_info.teams or ()))
|
||||
results: Final = await asyncio.gather(
|
||||
*(
|
||||
create_team_member_add_task(team_id, user_info, upsert_team if team_id == team_id_to_upsert else None)
|
||||
create_team_member_add_task(team_id, user_info, upsert_team)
|
||||
for team_id in missing_teams
|
||||
)
|
||||
)
|
||||
|
|
@ -3651,11 +3652,13 @@ class SSOAuthenticationHandler:
|
|||
verbose_proxy_logger.debug("User not found in LiteLLM DB, skipping team member addition")
|
||||
return
|
||||
sso_teams: Final = result.team_ids if isinstance(result, CustomOpenID) else []
|
||||
upsert_teams: Final = isinstance(result, CustomOpenID) and result.team_id_upsert
|
||||
await add_missing_team_member(
|
||||
user_info=user_info,
|
||||
sso_teams=sso_teams,
|
||||
team_id_to_upsert=result.team_id_to_upsert if isinstance(result, CustomOpenID) else None,
|
||||
upsert_team=partial(_upsert_sso_team, prisma_client=prisma_client, user_api_key_cache=user_api_key_cache),
|
||||
upsert_team=partial(_upsert_sso_team, prisma_client=prisma_client, user_api_key_cache=user_api_key_cache)
|
||||
if upsert_teams
|
||||
else None,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
|
|
|
|||
|
|
@ -0,0 +1,79 @@
|
|||
import json
|
||||
import time
|
||||
import uuid
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
|
||||
import jwt
|
||||
import pytest
|
||||
import yaml
|
||||
from cryptography.hazmat.primitives.asymmetric import rsa
|
||||
|
||||
from tests.integration._support.client import Gateway, eventually
|
||||
from tests.integration._support.database import read_rows
|
||||
from tests.integration._support.process import owned_proxy
|
||||
from tests.integration._support.wire import Reply, Request, wire_server
|
||||
|
||||
pytestmark: Final = pytest.mark.shares_worker_state
|
||||
|
||||
KEY_ID: Final = "integration-jwt-team-claim-signing-key"
|
||||
|
||||
|
||||
def _jwks_reply(public_jwk: str) -> Reply:
|
||||
return Reply(body=json.dumps({"keys": [{**json.loads(public_jwk), "kid": KEY_ID}]}).encode())
|
||||
|
||||
|
||||
@pytest.mark.parametrize("team_id_upsert", [True, False], ids=["upsert-creates-claimed-teams", "no-upsert-skips-them"])
|
||||
def test_jwt_teams_array_claim_creates_and_joins_every_team_only_when_team_id_upsert_is_on(
|
||||
gateway: Gateway, tmp_path: Path, team_id_upsert: bool
|
||||
) -> None:
|
||||
private_key: Final = rsa.generate_private_key(public_exponent=65537, key_size=2048)
|
||||
public_jwk: Final = jwt.algorithms.RSAAlgorithm.to_jwk(private_key.public_key())
|
||||
|
||||
def respond(request: Request) -> Reply:
|
||||
assert request.method == "GET", request
|
||||
return _jwks_reply(public_jwk)
|
||||
|
||||
with wire_server(respond) as jwks, gateway.scenario() as scenario:
|
||||
config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
|
||||
config["general_settings"] = {
|
||||
**config["general_settings"],
|
||||
"enable_jwt_auth": True,
|
||||
"litellm_jwtauth": {
|
||||
"user_id_jwt_field": "sub",
|
||||
"user_id_upsert": True,
|
||||
"team_ids_jwt_field": "teams",
|
||||
"team_id_upsert": team_id_upsert,
|
||||
"sync_user_role_and_teams": True,
|
||||
},
|
||||
}
|
||||
path: Final = tmp_path / "jwt_team_claim_upsert.yaml"
|
||||
path.write_text(yaml.safe_dump(config))
|
||||
subject: Final = f"integration-jwt-{uuid.uuid4().hex}"
|
||||
claimed_teams: Final = sorted(f"jwt-team-{subject}-{suffix}" for suffix in ("a", "b"))
|
||||
token: Final = jwt.encode(
|
||||
{"sub": subject, "teams": claimed_teams, "iat": int(time.time()), "exp": int(time.time()) + 300},
|
||||
private_key,
|
||||
algorithm="RS256",
|
||||
headers={"kid": KEY_ID},
|
||||
)
|
||||
with owned_proxy(gateway, tmp_path, {"JWT_PUBLIC_KEY_URL": jwks.url}, config=path) as candidate:
|
||||
model: Final = scenario.model()
|
||||
scenario.cleanups.callback(scenario.delete_user, subject)
|
||||
scenario.cleanups.callback(gateway.request, "POST", "/team/delete", {"team_ids": claimed_teams})
|
||||
response: Final = candidate.request(
|
||||
"POST",
|
||||
"/v1/chat/completions",
|
||||
{"model": model, "messages": [{"role": "user", "content": "teams array claim"}]},
|
||||
key=token,
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
expected_teams: Final = claimed_teams if team_id_upsert else []
|
||||
user: Final = eventually(
|
||||
lambda: read_rows('SELECT teams FROM "LiteLLM_UserTable" WHERE user_id = %s', (subject,)),
|
||||
lambda rows: rows == [{"teams": expected_teams}],
|
||||
)
|
||||
teams: Final = read_rows(
|
||||
'SELECT team_id FROM "LiteLLM_TeamTable" WHERE team_id = ANY(%s) ORDER BY team_id', (claimed_teams,)
|
||||
)
|
||||
assert (user, [row["team_id"] for row in teams]) == ([{"teams": expected_teams}], expected_teams)
|
||||
|
|
@ -235,7 +235,7 @@ def test_sso_sign_in_upserts_the_team_id_jwt_field_team_only_when_team_id_upsert
|
|||
assert (len(teams), user) == (int(expect_member), [{"teams": [team_id] if expect_member else []}])
|
||||
|
||||
|
||||
def test_sso_sign_in_with_db_team_mappings_creates_the_team_id_jwt_field_team_and_only_joins_groups(
|
||||
def test_sso_sign_in_with_db_team_mappings_creates_the_team_id_jwt_field_team_and_the_group_teams(
|
||||
tmp_path: Path,
|
||||
) -> None:
|
||||
subject: Final = f"sso-mapped-user-{uuid.uuid4().hex[:12]}"
|
||||
|
|
@ -257,4 +257,7 @@ def test_sso_sign_in_with_db_team_mappings_creates_the_team_id_jwt_field_team_an
|
|||
'SELECT team_id FROM "LiteLLM_TeamTable" WHERE team_id IN (%s, %s)', (team_id, group_team_id)
|
||||
)
|
||||
user: Final = gateway.rows('SELECT teams FROM "LiteLLM_UserTable" WHERE user_id = %s', (subject,))
|
||||
assert (teams, user) == ([{"team_id": team_id}], [{"teams": [team_id]}])
|
||||
assert (sorted(row["team_id"] for row in teams), user) == (
|
||||
sorted([team_id, group_team_id]),
|
||||
[{"teams": sorted([team_id, group_team_id])}],
|
||||
)
|
||||
|
|
|
|||
|
|
@ -14,14 +14,14 @@ from fastapi import HTTPException
|
|||
from litellm.caching.dual_cache import DualCache
|
||||
from litellm.proxy._types import (
|
||||
DEFAULT_JWKS_STALE_TTL,
|
||||
JWTLiteLLMRoleMap,
|
||||
JWTIssuerConfig,
|
||||
LiteLLMRoutes,
|
||||
JWTLiteLLMRoleMap,
|
||||
LiteLLM_JWTAuth,
|
||||
LiteLLM_ModelTable,
|
||||
LiteLLM_TeamMembership,
|
||||
LiteLLM_TeamTable,
|
||||
LiteLLM_UserTable,
|
||||
LiteLLMRoutes,
|
||||
LitellmUserRoles,
|
||||
Member,
|
||||
ProxyErrorTypes,
|
||||
|
|
@ -8637,3 +8637,107 @@ async def test_decode_trusted_assertion_runs_the_configured_custom_validate_on_t
|
|||
else:
|
||||
assert isinstance(verified, AssertionRejection)
|
||||
assert "custom_validate" in verified.reason
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(("team_id_upsert", "expected_created"), [(True, ["team1", "team3"]), (False, [])])
|
||||
async def test_sync_user_role_and_teams_creates_the_claimed_array_teams_only_when_team_id_upsert_is_on(
|
||||
team_id_upsert: bool, expected_created: list[str]
|
||||
) -> None:
|
||||
jwt_handler = JWTHandler()
|
||||
jwt_handler.update_environment(
|
||||
prisma_client=None,
|
||||
user_api_key_cache=AsyncMock(),
|
||||
litellm_jwtauth=LiteLLM_JWTAuth(
|
||||
team_ids_jwt_field="my_id_teams", sync_user_role_and_teams=True, team_id_upsert=team_id_upsert
|
||||
),
|
||||
)
|
||||
token = {"my_id_teams": ["team3", "team1", "team2"]}
|
||||
user = LiteLLM_UserTable(user_id="u1", user_role=LitellmUserRoles.INTERNAL_USER.value, teams=["team2"])
|
||||
created: list[str] = [] # mutable-ok: records the teams the proxy would create
|
||||
|
||||
async def create_team(team_id: str, **kwargs: object) -> LiteLLM_TeamTable:
|
||||
created.append(team_id)
|
||||
return LiteLLM_TeamTable(team_id=team_id)
|
||||
|
||||
with (
|
||||
patch("litellm.proxy.auth.handle_jwt.get_team_object", create_team),
|
||||
patch(
|
||||
"litellm.proxy.management_endpoints.scim.scim_v2.patch_team_membership", new_callable=AsyncMock
|
||||
) as membership,
|
||||
):
|
||||
await JWTAuthManager.sync_user_role_and_teams(
|
||||
jwt_handler, token, user, AsyncMock(), user_api_key_cache=AsyncMock()
|
||||
)
|
||||
|
||||
assert created == expected_created
|
||||
assert sorted(membership.call_args.kwargs["teams_ids_to_add_user_to"]) == ["team1", "team3"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_sync_user_role_and_teams_still_joins_the_team_when_a_concurrent_request_created_it_first() -> None:
|
||||
jwt_handler = JWTHandler()
|
||||
jwt_handler.update_environment(
|
||||
prisma_client=None,
|
||||
user_api_key_cache=AsyncMock(),
|
||||
litellm_jwtauth=LiteLLM_JWTAuth(
|
||||
team_ids_jwt_field="my_id_teams", sync_user_role_and_teams=True, team_id_upsert=True
|
||||
),
|
||||
)
|
||||
user = LiteLLM_UserTable(user_id="u1", user_role=LitellmUserRoles.INTERNAL_USER.value, teams=[])
|
||||
|
||||
async def lost_the_create_race(team_id: str, **kwargs: object) -> LiteLLM_TeamTable:
|
||||
raise HTTPException(status_code=404, detail={"error": f"Team id = {team_id} already exists"})
|
||||
|
||||
with (
|
||||
patch("litellm.proxy.auth.handle_jwt.get_team_object", lost_the_create_race),
|
||||
patch(
|
||||
"litellm.proxy.management_endpoints.scim.scim_v2.patch_team_membership", new_callable=AsyncMock
|
||||
) as membership,
|
||||
):
|
||||
await JWTAuthManager.sync_user_role_and_teams(
|
||||
jwt_handler, {"my_id_teams": ["team1"]}, user, AsyncMock(), user_api_key_cache=AsyncMock()
|
||||
)
|
||||
|
||||
assert membership.call_args.kwargs["teams_ids_to_add_user_to"] == ["team1"]
|
||||
assert user.teams == ["team1"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(("team_id_upsert", "expected_team_id"), [(True, "team-1"), (False, None)])
|
||||
async def test_find_team_with_model_access_creates_the_claimed_team_only_when_team_id_upsert_is_on(
|
||||
monkeypatch: pytest.MonkeyPatch, team_id_upsert: bool, expected_team_id: str | None
|
||||
) -> None:
|
||||
import sys
|
||||
import types
|
||||
|
||||
from litellm.caching import DualCache
|
||||
from litellm.proxy.utils import ProxyLogging
|
||||
|
||||
proxy_server_module = types.ModuleType("proxy_server")
|
||||
proxy_server_module.llm_router = None
|
||||
monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", proxy_server_module)
|
||||
|
||||
async def team_table_with_upsert(team_id: str, **kwargs: object) -> LiteLLM_TeamTable:
|
||||
if not kwargs.get("team_id_upsert"):
|
||||
raise HTTPException(status_code=404, detail={"error": f"Team doesn't exist in db. Team={team_id}"})
|
||||
return LiteLLM_TeamTable(team_id=team_id)
|
||||
|
||||
monkeypatch.setattr("litellm.proxy.auth.handle_jwt.get_team_object", team_table_with_upsert)
|
||||
jwt_handler = JWTHandler()
|
||||
jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(team_id_upsert=team_id_upsert)
|
||||
user_api_key_cache = DualCache()
|
||||
|
||||
team_id, _ = await JWTAuthManager.find_team_with_model_access(
|
||||
team_ids={"team-1"},
|
||||
requested_model=None,
|
||||
route="/chat/completions",
|
||||
jwt_handler=jwt_handler,
|
||||
prisma_client=None,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
parent_otel_span=None,
|
||||
proxy_logging_obj=ProxyLogging(user_api_key_cache=user_api_key_cache),
|
||||
team_id_upsert=team_id_upsert,
|
||||
)
|
||||
|
||||
assert team_id == expected_team_id
|
||||
|
|
|
|||
|
|
@ -6,7 +6,7 @@ from collections.abc import Awaitable, Callable, Mapping, Sequence
|
|||
from contextlib import ExitStack, asynccontextmanager
|
||||
from types import SimpleNamespace
|
||||
from typing import Final
|
||||
from unittest.mock import ANY, AsyncMock, MagicMock, patch
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
|
@ -6683,8 +6683,7 @@ class TestAddMissingTeamMember:
|
|||
mock_add_member.assert_called_once_with(
|
||||
user_info=new_user_info,
|
||||
sso_teams=["entra-group-1", "entra-group-2"],
|
||||
team_id_to_upsert=None,
|
||||
upsert_team=ANY,
|
||||
upsert_team=None,
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -10967,48 +10966,55 @@ def _upsert_jwt_handler(**jwtauth: object) -> JWTHandler:
|
|||
@pytest.mark.parametrize(
|
||||
("jwtauth", "userinfo", "expected_team_ids", "expected_upsert"),
|
||||
[
|
||||
({"team_id_jwt_field": "dept", "team_id_upsert": True}, {"dept": "finance"}, ["finance"], "finance"),
|
||||
({"team_id_jwt_field": "dept", "team_id_upsert": False}, {"dept": "finance"}, ["finance"], None),
|
||||
({"team_id_jwt_field": "dept", "team_id_upsert": True}, {}, [], None),
|
||||
({"team_id_jwt_field": "dept", "team_id_upsert": True}, {"dept": ""}, [], None),
|
||||
({"team_id_jwt_field": "dept", "team_id_upsert": True}, {"dept": " "}, [" "], None),
|
||||
({"team_id_jwt_field": "dept", "team_id_upsert": True}, {"dept": ["a", "b"]}, ["a", "b"], None),
|
||||
({"team_id_jwt_field": "dept", "team_id_upsert": True}, {"dept": "finance"}, ["finance"], True),
|
||||
({"team_id_jwt_field": "dept", "team_id_upsert": False}, {"dept": "finance"}, ["finance"], False),
|
||||
({"team_id_jwt_field": "dept", "team_id_upsert": True}, {}, [], True),
|
||||
({"team_id_jwt_field": "dept", "team_id_upsert": True}, {"dept": ""}, [], True),
|
||||
({"team_id_jwt_field": "dept", "team_id_upsert": True}, {"dept": " "}, [], True),
|
||||
({"team_id_jwt_field": "dept", "team_id_upsert": True}, {"dept": ["a", "b"]}, ["a", "b"], True),
|
||||
(
|
||||
{"team_id_jwt_field": "dept", "team_id_upsert": True, "team_id_default": "everyone"},
|
||||
{},
|
||||
[],
|
||||
None,
|
||||
True,
|
||||
),
|
||||
(
|
||||
{"team_ids_jwt_field": "groups", "team_id_upsert": True},
|
||||
{"groups": ["g1", "g2"]},
|
||||
["g1", "g2"],
|
||||
None,
|
||||
True,
|
||||
),
|
||||
(
|
||||
{"team_ids_jwt_field": "groups", "team_id_upsert": False},
|
||||
{"groups": ["g1", "g2"]},
|
||||
["g1", "g2"],
|
||||
False,
|
||||
),
|
||||
(
|
||||
{"team_ids_jwt_field": "groups", "team_id_jwt_field": "dept", "team_id_upsert": True},
|
||||
{"groups": ["g1"], "dept": "finance"},
|
||||
["g1", "finance"],
|
||||
"finance",
|
||||
True,
|
||||
),
|
||||
],
|
||||
ids=[
|
||||
"upsert-on",
|
||||
"upsert-off",
|
||||
"claim-missing",
|
||||
"blank-claim-not-upserted",
|
||||
"whitespace-claim-not-upserted",
|
||||
"list-valued-claim-join-only",
|
||||
"default-team-not-upserted",
|
||||
"plural-claim-join-only",
|
||||
"blank-claim-dropped",
|
||||
"whitespace-claim-dropped",
|
||||
"list-valued-claim",
|
||||
"default-team-not-joined",
|
||||
"plural-claim-upsert-on",
|
||||
"plural-claim-upsert-off",
|
||||
"mixed",
|
||||
],
|
||||
)
|
||||
def test_generic_sso_login_upserts_only_the_team_id_jwt_field_claim(
|
||||
def test_generic_sso_login_marks_every_claimed_team_for_upsert_only_when_team_id_upsert_is_on(
|
||||
jwtauth: dict[str, object],
|
||||
userinfo: dict[str, object],
|
||||
expected_team_ids: list[str],
|
||||
expected_upsert: str | None,
|
||||
expected_upsert: bool,
|
||||
) -> None:
|
||||
from litellm.proxy.management_endpoints.ui_sso import generic_response_convertor
|
||||
|
||||
|
|
@ -11017,10 +11023,10 @@ def test_generic_sso_login_upserts_only_the_team_id_jwt_field_claim(
|
|||
jwt_handler=_upsert_jwt_handler(**jwtauth),
|
||||
)
|
||||
|
||||
assert (result.team_ids, result.team_id_to_upsert) == (expected_team_ids, expected_upsert)
|
||||
assert (result.team_ids, result.team_id_upsert) == (expected_team_ids, expected_upsert)
|
||||
|
||||
|
||||
def test_generic_sso_login_upserts_the_team_id_jwt_field_team_when_db_team_mappings_are_set() -> None:
|
||||
def test_generic_sso_login_keeps_the_team_id_jwt_field_team_when_db_team_mappings_are_set() -> None:
|
||||
from litellm.proxy.management_endpoints.ui_sso import generic_response_convertor
|
||||
|
||||
result: Final = generic_response_convertor(
|
||||
|
|
@ -11029,7 +11035,7 @@ def test_generic_sso_login_upserts_the_team_id_jwt_field_team_when_db_team_mappi
|
|||
team_mappings=TeamMappings(team_ids_jwt_field="groups"),
|
||||
)
|
||||
|
||||
assert (result.team_ids, result.team_id_to_upsert) == (["g1", "finance"], "finance")
|
||||
assert (result.team_ids, result.team_id_upsert) == (["g1", "finance"], True)
|
||||
|
||||
|
||||
def _recording_team_calls(events: list[tuple[str, str]]) -> tuple[AsyncMock, Callable[[str], Awaitable[object]]]:
|
||||
|
|
@ -11044,21 +11050,35 @@ def _recording_team_calls(events: list[tuple[str, str]]) -> tuple[AsyncMock, Cal
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_sso_login_creates_only_the_upsert_team_and_does_so_before_joining_it() -> None:
|
||||
async def test_sso_login_creates_every_missing_team_before_joining_it() -> None:
|
||||
from litellm.proxy.management_endpoints.ui_sso import add_missing_team_member
|
||||
|
||||
events: Final[list[tuple[str, str]]] = [] # mutable-ok: records call order across awaited tasks
|
||||
member_add, upsert_team = _recording_team_calls(events)
|
||||
user: Final = NewUserResponse(user_id="alice", key="sk-x", teams=None)
|
||||
user: Final = NewUserResponse(user_id="alice", key="sk-x", teams=["g0"])
|
||||
|
||||
with patch("litellm.proxy.management_endpoints.ui_sso.team_member_add", member_add):
|
||||
await add_missing_team_member(user, ["g1", "finance", "g2"], team_id_to_upsert="finance", upsert_team=upsert_team)
|
||||
await add_missing_team_member(user, ["g0", "g1", "finance", "g2"], upsert_team=upsert_team)
|
||||
|
||||
creates: Final = [team_id for kind, team_id in events if kind == "create"]
|
||||
joins: Final = [team_id for kind, team_id in events if kind == "join"]
|
||||
assert creates == ["finance"]
|
||||
assert sorted(creates) == ["finance", "g1", "g2"]
|
||||
assert sorted(joins) == ["finance", "g1", "g2"]
|
||||
assert events.index(("create", "finance")) < events.index(("join", "finance"))
|
||||
assert all(events.index(("create", team_id)) < events.index(("join", team_id)) for team_id in creates)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_sso_login_creates_no_team_when_team_id_upsert_is_off() -> None:
|
||||
from litellm.proxy.management_endpoints.ui_sso import add_missing_team_member
|
||||
|
||||
events: Final[list[tuple[str, str]]] = [] # mutable-ok: records call order across awaited tasks
|
||||
member_add, _ = _recording_team_calls(events)
|
||||
user: Final = NewUserResponse(user_id="alice", key="sk-x", teams=None)
|
||||
|
||||
with patch("litellm.proxy.management_endpoints.ui_sso.team_member_add", member_add):
|
||||
await add_missing_team_member(user, ["g1", "finance"])
|
||||
|
||||
assert events == [("join", "finance"), ("join", "g1")]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -11076,7 +11096,7 @@ async def test_sso_login_still_joins_the_team_when_a_concurrent_login_created_it
|
|||
user: Final = NewUserResponse(user_id="alice", key="sk-x", teams=None)
|
||||
|
||||
with patch("litellm.proxy.management_endpoints.ui_sso.team_member_add", member_add):
|
||||
await add_missing_team_member(user, ["finance"], team_id_to_upsert="finance", upsert_team=lost_the_create_race)
|
||||
await add_missing_team_member(user, ["finance"], upsert_team=lost_the_create_race)
|
||||
|
||||
assert joined == ["finance"]
|
||||
|
||||
|
|
@ -11127,7 +11147,7 @@ async def test_sso_login_warns_when_creating_the_team_fails_instead_of_calling_i
|
|||
AsyncMock(side_effect=TeamNotFoundError(team_id="finance")),
|
||||
),
|
||||
):
|
||||
await add_missing_team_member(user, ["finance"], team_id_to_upsert="finance", upsert_team=failed_create)
|
||||
await add_missing_team_member(user, ["finance"], upsert_team=failed_create)
|
||||
|
||||
warnings: Final = [r.getMessage() for r in caplog.records if r.levelno == logging.WARNING]
|
||||
assert len(warnings) == 1 and "finance" in warnings[0] and str(create_error.detail) in warnings[0]
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue