fix(proxy): let /team/delete transactions wait for a pooled connection (#51099)

The team-model and team-key transactions in /team/delete used Prisma's 2s
default max_wait, so a busy writer pool failed the delete with P2028 while
every plain query in the same request waited up to the 60s pool timeout. Both
now get the 30s team-delete budget for max_wait, and the key transaction gets
it for timeout too.

litellm/proxy/management_endpoints/model_management_endpoints.py

@@ -295,7 +295,9 @@ class _TxModelTables(Protocol):
 
 @runtime_checkable
 class _TransactionFactory(Protocol):
-    def __call__(self, *, timeout: datetime.timedelta = ...) -> AbstractAsyncContextManager[_TxModelTables]: ...
+    def __call__(
+        self, *, timeout: datetime.timedelta = ..., max_wait: datetime.timedelta = ...
+    ) -> AbstractAsyncContextManager[_TxModelTables]: ...
 
 
 class _ModelTransactionClient(LiteLLMBaseModel):
@@ -2449,7 +2451,8 @@ async def delete_team_models(
     """
     deleted_model_ids: Final[list[str]] = []
     transaction_client: Final = _ModelTransactionClient.model_validate(prisma_client.db)
-    async with transaction_client.tx(timeout=datetime.timedelta(seconds=TEAM_DELETE_TRANSACTION_TIMEOUT_SECONDS)) as tx:
+    tx_budget: Final = datetime.timedelta(seconds=TEAM_DELETE_TRANSACTION_TIMEOUT_SECONDS)
+    async with transaction_client.tx(max_wait=tx_budget, timeout=tx_budget) as tx:
         for team_id in team_ids:
             rows = await _get_team_deployments(team_id, prisma_client, table=tx.litellm_proxymodeltable)
             model_ids = [row.model_id for row in rows]

litellm/proxy/utils.py

@@ -73,6 +73,7 @@ from litellm.constants import (
     SPEND_LOG_READ_STATEMENT_TIMEOUT_SECONDS,
     SPEND_LOG_WRITE_BATCH_MAX_BYTES,
     SPEND_LOG_WRITE_BATCH_MAX_ROWS,
+    TEAM_DELETE_TRANSACTION_TIMEOUT_SECONDS,
     VLLM_CHOICE_FIELDS,
 )
 from litellm.cost_calculator import BaseTokenUsageProcessor
@@ -5361,6 +5362,7 @@ _deprecated_key_cache: Final[LimitedSizeOrderedDict] = LimitedSizeOrderedDict(ma
 _DEPRECATED_KEY_CACHE_TTL_SECONDS: Final = 60
 _PRISMA_DEFAULT_TX_TIMEOUT: Final = timedelta(seconds=5)
 _PRISMA_DEFAULT_TX_MAX_WAIT: Final = timedelta(seconds=2)
+_TEAM_DELETE_TX_BUDGET: Final = timedelta(seconds=TEAM_DELETE_TRANSACTION_TIMEOUT_SECONDS)
 _SPEND_LOG_READ_STATEMENT_TIMEOUT: Final = timedelta(seconds=SPEND_LOG_READ_STATEMENT_TIMEOUT_SECONDS)
 _STATEMENT_TIMEOUT_TRANSACTION_GRACE: Final = timedelta(seconds=5)
 
@@ -6983,7 +6985,7 @@ class PrismaClient:
                 return {"deleted_teams": team_id_list}
             elif table_name == "key" and team_id_list is not None and isinstance(team_id_list, list):
                 # admin only endpoint -> `/team/delete`
-                async with self.tx() as transaction:
+                async with self.tx(max_wait=_TEAM_DELETE_TX_BUDGET, timeout=_TEAM_DELETE_TX_BUDGET) as transaction:
                     _ = await transaction.query_raw(_LOCK_KEYS_BY_TEAM_SQL, team_id_list)
                     _ = await transaction.litellm_verificationtoken.delete_many(where={"team_id": {"in": team_id_list}})
         except Exception as e:

tests/integration/management/test_team_delete_large_membership.py

@@ -21,6 +21,7 @@ import httpx
 import psycopg
 import pytest
 import yaml
+from psycopg.rows import TupleRow
 from pydantic import JsonValue
 from redis import Redis
 
@@ -55,7 +56,11 @@ FROM generate_series(1, %s::int) AS n
 """
 TAKE_TEAM_LOCK_SQL: Final = "SELECT pg_advisory_xact_lock(hashtext(%s))"
 LOCK_MODEL_ROW_SQL: Final = 'SELECT model_id FROM "LiteLLM_ProxyModelTable" WHERE model_id = %s FOR UPDATE'
+LOCK_TEAM_KEY_ROWS_SQL: Final = 'SELECT token FROM "LiteLLM_VerificationToken" WHERE team_id = %s FOR UPDATE'
+LOCK_BUDGET_ROW_SQL: Final = 'SELECT budget_id FROM "LiteLLM_BudgetTable" WHERE budget_id = %s FOR UPDATE'
 LOCK_HELD_PAST_DEFAULT_TRANSACTION_SECONDS: Final = 7
+POOL_HELD_PAST_DELETE_RETRIES_SECONDS: Final = 12
+SATURATED_POOL_LIMIT: Final = 2
 OLDEST_TRANSACTION_BLOCKED_BY_SQL: Final = """
 SELECT coalesce(max(extract(epoch FROM clock_timestamp() - xact_start)), 0)::float AS age
 FROM pg_stat_activity
@@ -146,6 +151,16 @@ def _waiters_on_lock_held_by(backend_pid: int) -> int:
     return waiting
 
 
+def _sessions_blocked_by(backend_pid: int) -> int:
+    rows: Final = read_rows(
+        "SELECT count(*)::int AS blocked FROM pg_stat_activity WHERE %s::int = ANY(pg_blocking_pids(pid))",
+        (str(backend_pid),),
+    )
+    blocked: Final = rows[0]["blocked"]
+    assert isinstance(blocked, int)
+    return blocked
+
+
 def _oldest_transaction_blocked_by(backend_pid: int) -> float:
     rows: Final = read_rows(OLDEST_TRANSACTION_BLOCKED_BY_SQL, (str(backend_pid),))
     age: Final = rows[0]["age"]
@@ -468,6 +483,83 @@ def test_delete_completes_after_waiting_on_the_team_lock_past_the_default_transa
         assert len(_tombstones(team)) == 1, _tombstones(team)
 
 
+def _park_on_team_key_rows(holder: psycopg.Connection[TupleRow], team_id: str) -> None:
+    _ = holder.execute(LOCK_TEAM_KEY_ROWS_SQL, (team_id,))
+
+
+def _park_on_deleted_key_records(holder: psycopg.Connection[TupleRow], _team_id: str) -> None:
+    _ = holder.execute('LOCK TABLE "LiteLLM_DeletedVerificationToken" IN SHARE MODE')
+
+
+def _requests_in_flight_including_this_one(proxy: Gateway) -> JsonValue:
+    return proxy.get("/health/backlog")["in_flight_requests"]
+
+
+@pytest.mark.timeout(240)
+@pytest.mark.parametrize(
+    "park",
+    [
+        pytest.param(_park_on_team_key_rows, id="next-transaction-is-the-team-model-delete"),
+        pytest.param(_park_on_deleted_key_records, id="next-transaction-is-the-team-key-delete"),
+    ],
+)
+def test_delete_waits_for_a_pooled_connection_when_every_one_is_busy(
+    gateway: Gateway,
+    tmp_path: Path,
+    park: Callable[[psycopg.Connection[TupleRow], str], None],
+) -> None:
+    database_url: Final = os.environ["DATABASE_URL"]
+    with (
+        gateway.scenario() as scenario,
+        owned_proxy_process(
+            gateway,
+            tmp_path,
+            {"DATABASE_URL": database_url},
+            config=_config_with_pool_limit(tmp_path, SATURATED_POOL_LIMIT),
+            remove_environment=("DATABASE_URL_READ_REPLICA",),
+        ) as owned,
+        ThreadPoolExecutor(max_workers=3) as pool,
+    ):
+        team: Final = string_value(
+            owned.gateway.post("/team/new", {"team_alias": f"integration-{uuid.uuid4().hex}"})["team_id"]
+        )
+        scenario.cleanups.callback(_remove_team_by_sql, team)
+        scenario.cleanups.callback(write_rows, 'DELETE FROM "LiteLLM_VerificationToken" WHERE team_id = %s', (team,))
+        keys: Final = _seeded_team_keys(team, 3)
+        budget: Final = f"integration-{uuid.uuid4().hex}"
+        owned.gateway.post("/budget/new", {"budget_id": budget, "max_budget": 100})
+        scenario.cleanups.callback(write_rows, 'DELETE FROM "LiteLLM_BudgetTable" WHERE budget_id = %s', (budget,))
+        update_budget: Final = {"budget_id": budget, "max_budget": 1}
+
+        with psycopg.connect(database_url) as budget_holder:
+            budget_holder.execute(LOCK_BUDGET_ROW_SQL, (budget,))
+            first_connection_taker: Final = _post_in_background(pool, owned.gateway, "/budget/update", update_budget)
+            eventually(lambda: _sessions_blocked_by(budget_holder.info.backend_pid), lambda count: count == 1)
+            with psycopg.connect(database_url) as parking:
+                park(parking, team)
+                pending_delete: Final = _post_in_background(pool, owned.gateway, "/team/delete", {"team_ids": [team]})
+                eventually(lambda: _sessions_blocked_by(parking.info.backend_pid), lambda count: count == 1)
+                second_connection_taker: Final = _post_in_background(
+                    pool, owned.gateway, "/budget/update", update_budget
+                )
+                eventually(lambda: _requests_in_flight_including_this_one(owned.gateway), lambda count: count == 4)
+            eventually(
+                pending_delete.done,
+                lambda done: done,
+                seconds=POOL_HELD_PAST_DELETE_RETRIES_SECONDS,
+                return_last_on_timeout=True,
+            )
+
+        response: Final = pending_delete.result(timeout=60)
+        assert response.status_code == 200, response.text
+        assert response.json() == {"deleted_teams": [team]}
+        takers: Final = (first_connection_taker, second_connection_taker)
+        assert [taker.result(timeout=60).status_code for taker in takers] == [200, 200]
+        assert _team_rows(team) == []
+        assert [_live_token(_hashed(key)) for key in keys] == [[], [], []]
+        assert len(_tombstones(team)) == 1, _tombstones(team)
+
+
 def test_deleting_two_teams_sharing_a_member_in_one_call_clears_both_from_the_member(gateway: Gateway) -> None:
     with gateway.scenario() as scenario:
         user: Final = scenario.user()