fix(proxy): give /team/delete transactions a 30s timeout (#46979)

Both interactive transactions in /team/delete used Prisma's 5s default.
A lock wait or a short database stall past 5s expired the transaction
after the keys were already deleted, so the call returned 503 and left
the team row, memberships and team models behind

litellm/constants.py

@@ -1859,6 +1859,7 @@ SPEND_COUNTER_RESEED_LOCKS_MAX_SIZE: Final = get_env_int("SPEND_COUNTER_RESEED_L
 PROXY_DB_LOOKUP_MAX_CONCURRENCY: Final = max(1, get_env_int("PROXY_DB_LOOKUP_MAX_CONCURRENCY", 25))
 PROXY_DB_LOOKUP_DEADLINE_SECONDS: Final = max(0.1, get_env_float("PROXY_DB_LOOKUP_DEADLINE_SECONDS", 10.0))
 PROXY_DB_LOOKUP_STALL_WINDOW_SECONDS: Final = max(0.0, get_env_float("PROXY_DB_LOOKUP_STALL_WINDOW_SECONDS", 30.0))
+TEAM_DELETE_TRANSACTION_TIMEOUT_SECONDS: Final = 30
 DEFAULT_CRON_JOB_LOCK_TTL_SECONDS: Final = get_env_int("DEFAULT_CRON_JOB_LOCK_TTL_SECONDS", 60)  # 1 minute
 PROXY_BUDGET_RESCHEDULER_MIN_TIME: Final = get_env_int("PROXY_BUDGET_RESCHEDULER_MIN_TIME", 597)
 RESET_BUDGET_JOB_BATCH_SIZE: Final = max(1, get_env_int("RESET_BUDGET_JOB_BATCH_SIZE", 500))

litellm/proxy/management_endpoints/model_management_endpoints.py

@@ -28,7 +28,7 @@ from pydantic import BaseModel, ConfigDict, Field, TypeAdapter, ValidationError,
 import litellm
 from litellm._logging import verbose_proxy_logger
 from litellm._uuid import uuid
-from litellm.constants import LITELLM_PROXY_ADMIN_NAME
+from litellm.constants import LITELLM_PROXY_ADMIN_NAME, TEAM_DELETE_TRANSACTION_TIMEOUT_SECONDS
 from litellm.litellm_core_utils.credential_accessor import CredentialAccessor
 from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap
 from litellm.litellm_core_utils.ptu_pricing import (
@@ -1881,7 +1881,7 @@ 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() as tx:
+    async with transaction_client.tx(timeout=datetime.timedelta(seconds=TEAM_DELETE_TRANSACTION_TIMEOUT_SECONDS)) 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/management_endpoints/team_endpoints.py

@@ -17,7 +17,7 @@ import traceback
 from collections.abc import Iterable, Mapping, Sequence
 from collections.abc import Set as AbstractSet
 from dataclasses import dataclass
-from datetime import datetime, timezone
+from datetime import datetime, timedelta, timezone
 from itertools import chain
 from types import MappingProxyType
 from typing import (
@@ -42,6 +42,7 @@ from litellm._internal_context import with_service_target
 from litellm._logging import verbose_proxy_logger
 from litellm._uuid import uuid
 from litellm.caching.redis_batch import active_request_redis_batches
+from litellm.constants import TEAM_DELETE_TRANSACTION_TIMEOUT_SECONDS
 from litellm.integrations.prometheus import PrometheusLogger
 from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
 from litellm.proxy._types import (
@@ -4674,7 +4675,7 @@ async def delete_team(
     # or is still waiting on the lock, in which case its own re-read happens after this commits
     # and sees the row gone before it writes anything.
     delete_filter: Final[_TeamIdInFilter] = {"team_id": {"in": data.team_ids}}
-    async with prisma_client.tx() as tx:
+    async with prisma_client.tx(timeout=timedelta(seconds=TEAM_DELETE_TRANSACTION_TIMEOUT_SECONDS)) as tx:
         for team_id in sorted(data.team_ids):
             await tx.query_raw(TEAM_ADVISORY_LOCK_SQL, team_id)
         await tx.litellm_teamtable.delete_many(where=delete_filter)

tests/integration/management/test_team_delete_large_membership.py

@@ -53,6 +53,13 @@ SELECT encode(sha256(convert_to(%s || '-' || n::text, 'UTF8')), 'hex'), %s, '{}'
 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_HELD_PAST_DEFAULT_TRANSACTION_SECONDS: Final = 7
+OLDEST_TRANSACTION_BLOCKED_BY_SQL: Final = """
+SELECT coalesce(max(extract(epoch FROM clock_timestamp() - xact_start)), 0)::float AS age
+FROM pg_stat_activity
+WHERE %s::int = ANY(pg_blocking_pids(pid))
+"""
 # Sessions blocked on the advisory lock the given backend holds, and nothing else on the shared rig.
 WAITERS_ON_HELD_LOCK_SQL: Final = """
 SELECT count(*)::int AS waiting
@@ -131,6 +138,13 @@ def _waiters_on_lock_held_by(backend_pid: int) -> int:
     return waiting
 
 
+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"]
+    assert isinstance(age, float)
+    return age
+
+
 def _remove_team_by_sql(team_id: str) -> None:
     """Cleanup for a team the test expects to have deleted itself. Whatever a failed delete left behind
     (row, memberships, `teams` references) goes by SQL so the shared rig stays clean without sending
@@ -161,6 +175,38 @@ def _create_team(scenario: Scenario, **fields: JsonValue) -> str:
     return team_id
 
 
+def _create_team_model(scenario: Scenario, team_id: str) -> str:
+    created: Final = scenario.gateway.post(
+        "/model/new",
+        {
+            "model_name": f"integration-{uuid.uuid4().hex}",
+            "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "sk-integration"},
+            "model_info": {"team_id": team_id},
+        },
+    )
+    model_id: Final = string_value(object_value(created["model_info"])["id"])
+    scenario.cleanups.callback(write_rows, 'DELETE FROM "LiteLLM_ProxyModelTable" WHERE model_id = %s', (model_id,))
+    return model_id
+
+
+def _team_model_rows(model_id: str) -> list[dict[str, JsonValue]]:
+    return read_rows('SELECT model_id FROM "LiteLLM_ProxyModelTable" WHERE model_id = %s', (model_id,))
+
+
+def _delete_while_locked(gateway: Gateway, team_id: str, lock_sql: str, lock_target: str) -> httpx.Response:
+    with ThreadPoolExecutor(max_workers=1) as pool:
+        with psycopg.connect(os.environ["DATABASE_URL"]) as holder:
+            holder.execute(lock_sql, (lock_target,))
+            pending: Final = _post_in_background(pool, gateway, "/team/delete", {"team_ids": [team_id]})
+            eventually(
+                lambda: _oldest_transaction_blocked_by(holder.info.backend_pid),
+                lambda age: age > LOCK_HELD_PAST_DEFAULT_TRANSACTION_SECONDS,
+                seconds=30,
+            )
+            assert not pending.done(), "delete returned while the lock was still held"
+        return pending.result(timeout=60)
+
+
 def _generate_key(scenario: Scenario, **fields: JsonValue) -> str:
     """A key the team delete is expected to remove; cleanup only deletes it if it is still live."""
     key: Final = string_value(scenario.gateway.post("/key/generate", fields)["key"])
@@ -392,6 +438,36 @@ def test_delete_waits_for_the_team_advisory_lock_and_completes_after_release(
         assert len(_tombstones(team)) == 1, _tombstones(team)
 
 
+def test_delete_completes_after_waiting_on_a_locked_team_model_row_past_the_default_transaction_timeout(
+    gateway: Gateway,
+) -> None:
+    with gateway.scenario() as scenario:
+        team: Final = _create_team(scenario)
+        model_id: Final = _create_team_model(scenario, team)
+        response: Final = _delete_while_locked(gateway, team, LOCK_MODEL_ROW_SQL, model_id)
+        assert response.status_code == 200, response.text
+        assert response.json() == {"deleted_teams": [team]}
+        assert _team_rows(team) == []
+        assert _team_model_rows(model_id) == []
+        assert len(_tombstones(team)) == 1, _tombstones(team)
+
+
+def test_delete_completes_after_waiting_on_the_team_lock_past_the_default_transaction_timeout(
+    gateway: Gateway,
+) -> None:
+    with gateway.scenario() as scenario:
+        user: Final = scenario.user()
+        team: Final = _create_team(scenario)
+        _bulk_member_add(gateway, team, [user])
+        response: Final = _delete_while_locked(gateway, team, TAKE_TEAM_LOCK_SQL, team)
+        assert response.status_code == 200, response.text
+        assert response.json() == {"deleted_teams": [team]}
+        assert _team_rows(team) == []
+        assert _membership_user_ids(team) == []
+        assert _user_teams(user) == []
+        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()