diff --git a/litellm/proxy/management_endpoints/scim/scim_v2.py b/litellm/proxy/management_endpoints/scim/scim_v2.py index 19e6856e9f3..5957f290a3d 100644 --- a/litellm/proxy/management_endpoints/scim/scim_v2.py +++ b/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: diff --git a/litellm/proxy/management_helpers/utils.py b/litellm/proxy/management_helpers/utils.py index 72246a570c9..1c4eada67fb 100644 --- a/litellm/proxy/management_helpers/utils.py +++ b/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( diff --git a/tests/unit/proxy/management_endpoints/scim/test_scim_v2_endpoints.py b/tests/unit/proxy/management_endpoints/scim/test_scim_v2_endpoints.py index 11fa0d2bf51..2ed2fa538e7 100644 --- a/tests/unit/proxy/management_endpoints/scim/test_scim_v2_endpoints.py +++ b/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