fix(team): keep member writes on the lock holder's connection after merge

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

litellm/proxy/management_helpers/utils.py

@@ -134,11 +134,11 @@ async def _find_users_by_email(
 
 
 async def _create_user_row(
-    prisma_client: PrismaClient, tx: MemberWriteTx | None, user_data: Mapping[str, object]
+    prisma_client: PrismaClient, tx: MemberWriteTx | None, user_data: dict[str, object]
 ) -> _PrismaUserRecord | None:
     if tx is not None:
-        return await tx.litellm_usertable.create(data=jsonify_object(dict(user_data)))
-    return await prisma_client.insert_data(data=dict(user_data), table_name="user")
+        return await tx.litellm_usertable.create(data=jsonify_object(user_data))
+    return await prisma_client.insert_data(data=user_data, table_name="user")
 
 
 def get_new_internal_user_defaults(user_id: str, user_email: str | None = None) -> dict[str, object]:

tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py

@@ -4,7 +4,7 @@ from contextlib import asynccontextmanager
 from datetime import datetime, timezone
 from types import SimpleNamespace
 from typing import Optional, cast
-from unittest.mock import AsyncMock, MagicMock, call, patch
+from unittest.mock import AsyncMock, MagicMock, PropertyMock, call, patch
 
 import pytest
 from fastapi import HTTPException
@@ -1954,11 +1954,28 @@ async def test_add_team_members_runs_member_writes_on_the_lock_holding_transacti
         _add_team_members_to_team,
     )
 
+    added_user = MagicMock()
+    added_user.user_id = "bob"
+    added_user.model_dump.return_value = {"user_id": "bob", "teams": ["team-pool"]}
+    created_budget = MagicMock()
+    created_budget.budget_id = "budget-pool"
+    membership = MagicMock()
+    membership.model_dump.return_value = {
+        "team_id": "team-pool",
+        "user_id": "bob",
+        "budget_id": "budget-pool",
+        "litellm_budget_table": None,
+    }
+
     tx = MagicMock()
     tx.query_raw = AsyncMock(return_value=[{"members_with_roles": []}])
     tx.litellm_teamtable.update = AsyncMock(
         return_value=LiteLLM_TeamTable(team_id="team-pool", members_with_roles=[])
     )
+    tx.litellm_usertable.upsert = AsyncMock(return_value=added_user)
+    tx.litellm_usertable.update_many = AsyncMock()
+    tx.litellm_budgettable.create = AsyncMock(return_value=created_budget)
+    tx.litellm_teammembership.create = AsyncMock(return_value=membership)
 
     tx_cm = MagicMock()
     tx_cm.__aenter__ = AsyncMock(return_value=tx)
@@ -1966,27 +1983,25 @@ async def test_add_team_members_runs_member_writes_on_the_lock_holding_transacti
 
     prisma_client = MagicMock()
     prisma_client.tx = MagicMock(return_value=tx_cm)
-
-    process_team_members = AsyncMock(return_value=([], []))
-    with patch(
-        "litellm.proxy.management_endpoints.team_endpoints._process_team_members",
-        new=process_team_members,
-    ):
-        await _add_team_members_to_team(
-            data=TeamMemberAddRequest(
-                team_id="team-pool",
-                member=Member(user_id="bob", role="user"),
-            ),
-            complete_team_data=LiteLLM_TeamTable(team_id="team-pool", members_with_roles=[]),
-            prisma_client=cast(object, prisma_client),
-            user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN),
-            litellm_proxy_admin_name="admin",
-        )
-
-    assert process_team_members.call_args.kwargs["tx"] is tx, (
-        "member writes must reuse the lock holder's connection, not check out another one"
+    type(prisma_client).db = PropertyMock(
+        side_effect=AssertionError("member writes must not reach for a second pooled connection")
     )
 
+    _, updated_users, updated_team_memberships = await _add_team_members_to_team(
+        data=TeamMemberAddRequest(
+            team_id="team-pool",
+            member=Member(user_id="bob", role="user"),
+            max_budget_in_team=50.0,
+        ),
+        complete_team_data=LiteLLM_TeamTable(team_id="team-pool", members_with_roles=[]),
+        prisma_client=cast(object, prisma_client),
+        user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN),
+        litellm_proxy_admin_name="admin",
+    )
+
+    assert [user.user_id for user in updated_users] == ["bob"]
+    assert [tm.budget_id for tm in updated_team_memberships] == ["budget-pool"]
+
 
 @pytest.mark.asyncio
 async def test_add_team_members_writes_nothing_when_the_team_is_deleted_mid_request():