From 1644ae516fa914501b192c69216610802bae0561 Mon Sep 17 00:00:00 2001 From: milan Date: Wed, 19 Aug 2026 21:55:32 +0000 Subject: [PATCH] fix(ui_sso): remove stale team memberships on SSO login when sync_user_role_and_teams is set Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../management_endpoints/scim/scim_v2.py | 4 +- litellm/proxy/management_endpoints/ui_sso.py | 48 ++++++ .../proxy/management_endpoints/test_ui_sso.py | 147 ++++++++++++++++++ 3 files changed, 197 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/management_endpoints/scim/scim_v2.py b/litellm/proxy/management_endpoints/scim/scim_v2.py index e50e9bf0537..9853c988858 100644 --- a/litellm/proxy/management_endpoints/scim/scim_v2.py +++ b/litellm/proxy/management_endpoints/scim/scim_v2.py @@ -1854,8 +1854,8 @@ def _is_user_not_in_team_error(exc: HTTPException) -> bool: async def patch_team_membership( user_id: str, - teams_ids_to_add_user_to: list[str], - teams_ids_to_remove_user_from: list[str], + teams_ids_to_add_user_to: Sequence[str], + teams_ids_to_remove_user_from: Sequence[str], raise_on_error: bool = False, ) -> bool: """ diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index 46af5dd80e1..ec3e04b99ec 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -1910,6 +1910,47 @@ async def _sync_user_role_from_jwt_role_map( ) +async def _sync_sso_team_memberships( + jwt_handler: JWTHandler | None, + sso_team_ids: Sequence[str], + user_info: LiteLLM_UserTable | NewUserResponse | None, + user_api_key_cache: UserApiKeyCache, +) -> None: + """ + Remove stale team memberships during SSO login. + + Mirrors JWTAuthManager.sync_user_role_and_teams so UI/SSO logins reconcile + team membership the same way API/JWT auth does: teams the IdP no longer + reports for the user are removed. Additions are handled separately by + add_user_to_teams_from_sso_response. + """ + if jwt_handler is None or user_info is None or user_info.user_id is None: + return + if not jwt_handler.litellm_jwtauth.sync_user_role_and_teams: + return + + sso_teams: Final = frozenset(sso_team_ids) + existing_teams: Final = frozenset(user_info.teams or ()) + preserve_db_teams_without_claims: Final = jwt_handler.litellm_jwtauth.fallback_to_db_teams and not sso_teams + teams_to_remove: Final = frozenset() if preserve_db_teams_without_claims else existing_teams - sso_teams + if not teams_to_remove: + return + + from litellm.proxy.management_endpoints.scim.scim_v2 import patch_team_membership + + await patch_team_membership( + user_id=user_info.user_id, + teams_ids_to_add_user_to=(), + teams_ids_to_remove_user_from=sorted(teams_to_remove), + ) + user_info.teams = sorted(sso_teams) + await user_api_key_cache.async_set_cache( + key=user_info.user_id, + value=user_info, + model_type=LiteLLM_UserTable, + ) + + def apply_user_info_values_to_sso_user_defined_values( user_info: LiteLLM_UserTable | NewUserResponse | None, user_defined_values: SSOUserDefinedValues | None, @@ -3532,6 +3573,13 @@ class SSOAuthenticationHandler: user_defined_values=user_defined_values, ) + await _sync_sso_team_memberships( + jwt_handler=jwt_handler, + sso_team_ids=getattr(result, "team_ids", ()) or (), + user_info=user_info, + user_api_key_cache=user_api_key_cache, + ) + user_defined_values = apply_user_info_values_to_sso_user_defined_values( user_info=user_info, user_defined_values=user_defined_values ) diff --git a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py index b83b862d6b8..df6a5f7c0f1 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py +++ b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py @@ -25,6 +25,7 @@ from litellm.proxy.management_endpoints.ui_sso import ( MicrosoftSSOHandler, SSOAuthenticationHandler, _setup_team_mappings, + _sync_sso_team_memberships, _sync_user_role_from_jwt_role_map, normalize_email, process_sso_jwt_access_token, @@ -7163,6 +7164,152 @@ class TestSyncUserRoleFromJwtRoleMap: prisma.db.litellm_usertable.update.assert_not_called() +class TestSyncSsoTeamMemberships: + """Tests for _sync_sso_team_memberships (SSO login team removal sync).""" + + @staticmethod + def _make_jwt_handler(sync_user_role_and_teams=True, fallback_to_db_teams=False): + from litellm.caching.caching import DualCache + from litellm.proxy._types import LiteLLM_JWTAuth + + handler = JWTHandler() + handler.update_environment( + prisma_client=None, + user_api_key_cache=DualCache(), + litellm_jwtauth=LiteLLM_JWTAuth( + sync_user_role_and_teams=sync_user_role_and_teams, + fallback_to_db_teams=fallback_to_db_teams, + ), + ) + return handler + + @pytest.mark.asyncio + async def test_removes_stale_teams_when_sync_enabled(self): + from litellm.caching.caching import DualCache + + user = LiteLLM_UserTable(user_id="testuser@example.com", teams=["team-a", "team-b"]) + + with patch( + "litellm.proxy.management_endpoints.scim.scim_v2.patch_team_membership", + new_callable=AsyncMock, + ) as mock_patch: + await _sync_sso_team_memberships( + jwt_handler=self._make_jwt_handler(), + sso_team_ids=["team-a"], + user_info=user, + user_api_key_cache=DualCache(), + ) + + mock_patch.assert_called_once_with( + user_id="testuser@example.com", + teams_ids_to_add_user_to=(), + teams_ids_to_remove_user_from=["team-b"], + ) + assert user.teams == ["team-a"] + + @pytest.mark.asyncio + async def test_no_removal_when_sync_disabled(self): + from litellm.caching.caching import DualCache + + user = LiteLLM_UserTable(user_id="testuser@example.com", teams=["team-a", "team-b"]) + + with patch( + "litellm.proxy.management_endpoints.scim.scim_v2.patch_team_membership", + new_callable=AsyncMock, + ) as mock_patch: + await _sync_sso_team_memberships( + jwt_handler=self._make_jwt_handler(sync_user_role_and_teams=False), + sso_team_ids=["team-a"], + user_info=user, + user_api_key_cache=DualCache(), + ) + + mock_patch.assert_not_called() + assert user.teams == ["team-a", "team-b"] + + @pytest.mark.asyncio + async def test_preserves_db_teams_when_no_sso_teams_and_fallback_enabled(self): + from litellm.caching.caching import DualCache + + user = LiteLLM_UserTable(user_id="testuser@example.com", teams=["team-a"]) + + with patch( + "litellm.proxy.management_endpoints.scim.scim_v2.patch_team_membership", + new_callable=AsyncMock, + ) as mock_patch: + await _sync_sso_team_memberships( + jwt_handler=self._make_jwt_handler(fallback_to_db_teams=True), + sso_team_ids=[], + user_info=user, + user_api_key_cache=DualCache(), + ) + + mock_patch.assert_not_called() + assert user.teams == ["team-a"] + + @pytest.mark.asyncio + async def test_removes_all_teams_when_no_sso_teams_and_no_fallback(self): + from litellm.caching.caching import DualCache + + user = LiteLLM_UserTable(user_id="testuser@example.com", teams=["team-a"]) + + with patch( + "litellm.proxy.management_endpoints.scim.scim_v2.patch_team_membership", + new_callable=AsyncMock, + ) as mock_patch: + await _sync_sso_team_memberships( + jwt_handler=self._make_jwt_handler(), + sso_team_ids=[], + user_info=user, + user_api_key_cache=DualCache(), + ) + + mock_patch.assert_called_once_with( + user_id="testuser@example.com", + teams_ids_to_add_user_to=(), + teams_ids_to_remove_user_from=["team-a"], + ) + assert user.teams == [] + + @pytest.mark.asyncio + async def test_noop_when_teams_match(self): + from litellm.caching.caching import DualCache + + user = LiteLLM_UserTable(user_id="testuser@example.com", teams=["team-a"]) + + with patch( + "litellm.proxy.management_endpoints.scim.scim_v2.patch_team_membership", + new_callable=AsyncMock, + ) as mock_patch: + await _sync_sso_team_memberships( + jwt_handler=self._make_jwt_handler(), + sso_team_ids=["team-a"], + user_info=user, + user_api_key_cache=DualCache(), + ) + + mock_patch.assert_not_called() + + @pytest.mark.asyncio + async def test_noop_when_jwt_handler_none(self): + from litellm.caching.caching import DualCache + + user = LiteLLM_UserTable(user_id="testuser@example.com", teams=["team-a"]) + + with patch( + "litellm.proxy.management_endpoints.scim.scim_v2.patch_team_membership", + new_callable=AsyncMock, + ) as mock_patch: + await _sync_sso_team_memberships( + jwt_handler=None, + sso_team_ids=[], + user_info=user, + user_api_key_cache=DualCache(), + ) + + mock_patch.assert_not_called() + + # ── VERIA-34 regression: PKCE state-to-session-cookie binding ───────────────