fix(scim): re-check the roster under the team lock before restoring a member's team

litellm/proxy/management_endpoints/scim/scim_v2.py

@@ -74,7 +74,7 @@ from litellm.proxy.management_endpoints.team_endpoints import (
     team_member_add,
     team_member_delete,
 )
-from litellm.proxy.management_helpers.access_group_team_sync import USER_ROWS_LOCK_SQL
+from litellm.proxy.management_helpers.access_group_team_sync import TEAM_ADVISORY_LOCK_SQL, USER_ROWS_LOCK_SQL
 from litellm.proxy.management_helpers.bulk_user_creation import (
     BULK_NEW_USER_CONCURRENCY,
     _bounded,  # pyright: ignore[reportPrivateUsage]  # the bounded gather /user/bulk_new runs its row writes through
@@ -94,7 +94,6 @@ from litellm.proxy.management_helpers.team_roster_sync import (
     team_tx_db,
     user_tx_db,
 )
-from litellm.proxy.management_helpers.utils import append_team_id_if_absent
 from litellm.proxy.utils import (  # noqa: F401  # legacy module exports
     PrismaClient,
     _premium_user_check,  # pyright: ignore[reportPrivateUsage,reportUnusedImport]  # backwards-compatible package export
@@ -2699,7 +2698,19 @@ async def _add_user_to_team(user_id: str, team_id: str) -> None:
     except ProxyException as e:
         if e.type != ProxyErrorTypes.team_member_already_in_team:
             raise
-        await append_team_id_if_absent(await _get_prisma_client_or_raise_exception(), user_id, team_id)
+        await _restore_team_on_roster_member(user_id, team_id)
+
+
+async def _restore_team_on_roster_member(user_id: str, team_id: str) -> None:
+    prisma_client: Final = await _get_prisma_client_or_raise_exception()
+    async with prisma_client.tx() as tx:
+        await tx.query_raw(TEAM_ADVISORY_LOCK_SQL, team_id)
+        roster: Final = await TeamRepository(prisma_client).get_members_with_roles_locked(tx, team_id)
+        if roster is not None and any(member.user_id == user_id for member in roster):
+            await user_tx_db(tx).update_many(
+                where={"user_id": user_id, "NOT": {"teams": {"has": team_id}}},
+                data={"teams": {"push": [team_id]}},
+            )
 
 
 async def _remove_user_from_team(user_id: str, team_id: str) -> None:

litellm/proxy/management_helpers/utils.py

@@ -400,7 +400,7 @@ async def _resolve_member_budget_id(
     return response.budget_id
 
 
-async def append_team_id_if_absent(
+async def _append_team_id_if_absent(
     prisma_client: PrismaClient, user_id: str, team_id: str, tx: MemberWriteTx | None = None
 ) -> None:
     """Append team_id to a user's teams array, only if it is not already present.
@@ -453,7 +453,7 @@ async def add_new_member(
             new_member.user_id,
             {"teams": [team_id], **new_user_defaults},
         )
-        await append_team_id_if_absent(prisma_client, new_member.user_id, team_id, tx)
+        await _append_team_id_if_absent(prisma_client, new_member.user_id, team_id, tx)
         if _returned_user is not None:
             returned_user = LiteLLM_UserTable.model_validate(_returned_user.model_dump())
     elif new_member.user_email is not None:
@@ -471,7 +471,7 @@ async def add_new_member(
                 returned_user = LiteLLM_UserTable.model_validate(_returned_user.model_dump())
         elif len(existing_user_row) == 1:
             user_info: Final = existing_user_row[0]
-            await append_team_id_if_absent(prisma_client, user_info.user_id, team_id, tx)
+            await _append_team_id_if_absent(prisma_client, user_info.user_id, team_id, tx)
             returned_user = LiteLLM_UserTable.model_validate(user_info.model_dump())
         elif len(existing_user_row) > 1:
             raise HTTPException(

tests/unit/proxy/management_endpoints/scim/test_scim_v2_endpoints.py

@@ -9,6 +9,7 @@ from types import MappingProxyType
 from typing import Final, assert_never
 from unittest.mock import ANY, AsyncMock, MagicMock, call
 
+from litellm.proxy.management_helpers.access_group_team_sync import TEAM_ADVISORY_LOCK_SQL
 import pytest
 from fastapi import FastAPI, HTTPException
 from httpx import ASGITransport, AsyncClient
@@ -1107,8 +1108,15 @@ async def test_handle_existing_user_by_email_roster_add_failure_blocks_teams_wri
 
 
 @pytest.mark.asyncio
-async def test_handle_existing_user_by_email_roster_add_already_member_repairs_user_teams(mocker):
-    """A user already on the roster whose teams list lacks the team gets it appended, not skipped"""
+@pytest.mark.parametrize(
+    ("roster", "repaired"),
+    [([{"user_id": "uid", "role": "user"}], True), ([], False)],
+    ids=["still-on-roster", "removed-before-the-lock"],
+)
+async def test_handle_existing_user_by_email_roster_add_already_member_repairs_user_teams(
+    mocker, roster: list[dict[str, str]], repaired: bool
+):
+    """An already-on-roster user missing the team gets it back, but only if the roster still has them under the lock"""
     existing_user = mocker.MagicMock()
     existing_user.user_id = "uid"
     existing_user.user_email = "member@example.com"
@@ -1121,7 +1129,13 @@ async def test_handle_existing_user_by_email_roster_add_already_member_repairs_u
     mock_prisma_client.db.litellm_usertable = mocker.MagicMock()
     mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(return_value=existing_user)
     mock_prisma_client.db.litellm_usertable.update = AsyncMock(return_value={})
-    mock_prisma_client.db.litellm_usertable.update_many = AsyncMock(return_value=1)
+    tx = mocker.MagicMock()
+    tx.query_raw = AsyncMock(side_effect=[[{"locked": None}], [{"members_with_roles": roster}]])
+    tx.litellm_usertable.update_many = AsyncMock(return_value=1)
+    tx_context = mocker.MagicMock()
+    tx_context.__aenter__ = AsyncMock(return_value=tx)
+    tx_context.__aexit__ = AsyncMock(return_value=False)
+    mock_prisma_client.tx = mocker.MagicMock(return_value=tx_context)
     mocker.patch(  # test-quality-ok: the roster helpers resolve the proxy prisma client at module level
         "litellm.proxy.management_endpoints.scim.scim_v2._get_prisma_client_or_raise_exception",
         AsyncMock(return_value=mock_prisma_client),
@@ -1157,10 +1171,14 @@ async def test_handle_existing_user_by_email_roster_add_already_member_repairs_u
     )
 
     mock_team_member_add.assert_awaited_once()
-    mock_prisma_client.db.litellm_usertable.update_many.assert_awaited_once_with(
-        where={"user_id": "uid", "NOT": {"teams": {"has": "team-x"}}},
-        data={"teams": {"push": ["team-x"]}},
-    )
+    assert tx.query_raw.await_args_list[0].args == (TEAM_ADVISORY_LOCK_SQL, "team-x")
+    if repaired:
+        tx.litellm_usertable.update_many.assert_awaited_once_with(
+            where={"user_id": "uid", "NOT": {"teams": {"has": "team-x"}}},
+            data={"teams": {"push": ["team-x"]}},
+        )
+    else:
+        tx.litellm_usertable.update_many.assert_not_awaited()
     update_calls = mock_prisma_client.db.litellm_usertable.update.call_args_list
     assert len(update_calls) == 1