From 1f038f880079a76700a495c2f9fd0d163fb227e5 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Fri, 26 Jun 2026 02:02:23 +0000
Subject: [PATCH] fix(jwt): scope dual-claim membership sync to
 fallback_to_db_teams

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.
---
 litellm/proxy/auth/handle_jwt.py              | 15 ++++---
 .../proxy/auth/test_handle_jwt.py             | 41 +++++++++++++++++++
 2 files changed, 51 insertions(+), 5 deletions(-)

diff --git a/litellm/proxy/auth/handle_jwt.py b/litellm/proxy/auth/handle_jwt.py
index e448c676939..bfccace9f85 100644
--- a/litellm/proxy/auth/handle_jwt.py
+++ b/litellm/proxy/auth/handle_jwt.py
@@ -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 = (
diff --git a/tests/test_litellm/proxy/auth/test_handle_jwt.py b/tests/test_litellm/proxy/auth/test_handle_jwt.py
index 45e61ea825f..3573850a0f9 100644
--- a/tests/test_litellm/proxy/auth/test_handle_jwt.py
+++ b/tests/test_litellm/proxy/auth/test_handle_jwt.py
@@ -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 == []