scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-08-28 05:25:59 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 35

fix(jwt): scope dual-claim membership sync to fallback_to_db_teams

Browse source 
The membership sync read both plural and singular JWT team claims via
get_all_jwt_team_ids unconditionally, which silently changed reconciliation
for every deployment using sync_user_role_and_teams, not just those opting
into fallback_to_db_teams: a singular-only IdP token that previously stripped
all DB teams would now be recognized. Gate the dual-claim read on
fallback_to_db_teams so flag-off deployments keep the upstream plural-only
behavior, honoring the PR's contract that existing deployments are unchanged.
Regression: test_sync_user_role_and_teams_singular_claim_only_recognized_under_flag.
This commit is contained in:
mateo-berri 2026-06-26 02:02:23 +00:00
parent 35221235aa
commit 1f038f8800
No known key found for this signature in database
2 changed files with 51 additions and 5 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

15
litellm/proxy/auth/handle_jwt.py
View file

@ -1835,11 +1835,16 @@ class JWTAuthManager:
ttl=DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL,
)
# Sync team memberships; include both plural and singular claim shapes so
# IdPs that populate only the singular field (e.g. Okta/Auth0) are not
# treated as claimless, which would otherwise leave stale DB memberships
# eligible for selection by _resolve_db_team_fallback on later requests.
jwt_team_ids = set(jwt_handler.get_all_jwt_team_ids(jwt_valid_token))
# Sync team memberships. With fallback_to_db_teams on, read both plural and
# singular claim shapes so a singular-only IdP token (e.g. Okta/Auth0) is
# 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 = 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)
)
existing_teams = set(user_object.teams or [])
teams_to_add = jwt_team_ids - existing_teams
preserve_db_teams_without_claims = (

41
tests/test_litellm/proxy/auth/test_handle_jwt.py
View file

@ -5392,3 +5392,44 @@ async def test_auth_builder_header_cannot_override_rbac_team_under_db_fallback()
)
assert exc_info.value.status_code == 403
@pytest.mark.asyncio
async def test_sync_user_role_and_teams_singular_claim_only_recognized_under_flag():
"""Reading the singular team claim during sync is scoped to fallback_to_db_teams.
With the flag off, sync keeps the upstream plural-only reconciliation, so a
singular-only token is treated as claimless and existing DB teams are removed
exactly as before this PR; the new dual-claim behavior must not silently change
membership reconciliation for deployments that never opted in."""
jwt_handler = JWTHandler()
jwt_handler.update_environment(
prisma_client=None,
user_api_key_cache=AsyncMock(),
litellm_jwtauth=LiteLLM_JWTAuth(
team_id_jwt_field="primary_team",
sync_user_role_and_teams=True,
fallback_to_db_teams=False,
),
)
token = {"sub": "u_flag_off", "primary_team": "team_primary"}
user = LiteLLM_UserTable(
user_id="u_flag_off",
user_role=LitellmUserRoles.INTERNAL_USER.value,
teams=["team_existing"],
)
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()
)
mock_patch.assert_awaited_once()
assert set(mock_patch.call_args.kwargs["teams_ids_to_remove_user_from"]) == {
"team_existing"
}
assert mock_patch.call_args.kwargs["teams_ids_to_add_user_to"] == []
assert user.teams == []