Merge 1644ae516f into c2c2a623c0

litellm/proxy/management_endpoints/scim/scim_v2.py

@@ -2277,8 +2277,8 @@ async def _remove_user_from_team(user_id: str, team_id: str) -> None:
 
 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:
     """

litellm/proxy/management_endpoints/ui_sso.py

@@ -1936,6 +1936,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,
@@ -3582,6 +3623,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
         )

tests/test_litellm/proxy/management_endpoints/test_ui_sso.py

@@ -20,6 +20,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,
@@ -7602,6 +7603,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 ───────────────