fix(proxy): finish prisma engine requests when their caller is cancelled (#46503)

Cancelling a task mid-request closed the query engine socket, the engine
handed the still-busy pool connection back, the next transaction start
queued its BEGIN behind the orphaned statement, and when that waiter was
cancelled too the engine panicked on its dropped reply channel
(interactive_transactions/actors.rs) and exited under the shutdown spend
flush, losing daily user and team spend for requests that had returned 200.

Every engine request now runs in its own task behind asyncio.shield so only
the caller is cancelled, and a transaction that opens after its caller left
is rolled back.

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

litellm/proxy/db/prisma_client.py

@@ -11,9 +11,9 @@ import subprocess
 import time
 import urllib
 import urllib.parse
-from collections.abc import Callable
+from collections.abc import Awaitable, Callable, Coroutine
 from datetime import datetime, timedelta
-from typing import TYPE_CHECKING, Any, Final, Protocol
+from typing import TYPE_CHECKING, Any, Final, Protocol, TypeVar
 
 from litellm._logging import verbose_proxy_logger
 from litellm.proxy.db.db_span import db_span
@@ -96,10 +96,20 @@ class _PrismaDrainTracker:
         await self._drained.wait()
 
 
+_EngineResponseT: Final = TypeVar("_EngineResponseT")
+
+
 class _TrackedPrismaEngine:
+    """Engine wrapper that counts in-flight work and never lets a cancelled caller drop an engine request.
+
+    A dropped request leaves the engine holding a busy pool connection and can crash it once that statement
+    finishes, so every request runs to completion in its own task and only the caller is cancelled.
+    """
+
     def __init__(self, engine: _PrismaEngine, tracker: _PrismaDrainTracker) -> None:
         self._engine = engine
         self.tracker = tracker
+        self._detached: frozenset[asyncio.Task[object]] = frozenset()
 
     @property
     def process(self) -> _PrismaProcess:
@@ -109,7 +119,60 @@ class _TrackedPrismaEngine:
         return getattr(self._engine, name)
 
     async def query(self, content: str, *, tx_id: str | None) -> object:
+        return await self._finish_even_if_cancelled(self._tracked_query(content, tx_id))
+
+    async def start_transaction(self, *, content: str) -> str:
+        return await self._finish_even_if_cancelled(
+            self._tracked_start_transaction(content), abandoned=self.rollback_transaction
+        )
+
+    async def commit_transaction(self, tx_id: str) -> None:
+        await self._finish_even_if_cancelled(self._tracked_commit_transaction(tx_id))
+
+    async def rollback_transaction(self, tx_id: str) -> None:
+        await self._finish_even_if_cancelled(self._tracked_rollback_transaction(tx_id))
+
+    async def _finish_even_if_cancelled(
+        self,
+        request: Coroutine[object, object, _EngineResponseT],
+        abandoned: Callable[[_EngineResponseT], Awaitable[None]] | None = None,
+    ) -> _EngineResponseT:
         self.tracker.begin_operation()
+        task: Final = asyncio.ensure_future(request)
+        try:
+            return await asyncio.shield(task)
+        except asyncio.CancelledError:
+            self._detached = self._detached.union((task,))
+            task.add_done_callback(lambda done: self._detached_request_finished(done, abandoned))
+            raise
+
+    def _detached_request_finished(
+        self,
+        done: "asyncio.Task[_EngineResponseT]",
+        abandoned: Callable[[_EngineResponseT], Awaitable[None]] | None,
+    ) -> None:
+        self._detached = self._detached.difference((done,))
+        if done.cancelled():
+            return
+        failure: Final = done.exception()
+        if failure is not None:
+            verbose_proxy_logger.debug("Prisma request abandoned by a cancelled caller failed: %s", failure)
+            return
+        if abandoned is None:
+            return
+        settle: Final = asyncio.ensure_future(abandoned(done.result()))
+        self._detached = self._detached.union((settle,))
+        settle.add_done_callback(self._abandoned_transaction_settled)
+
+    def _abandoned_transaction_settled(self, settle: "asyncio.Task[object]") -> None:
+        self._detached = self._detached.difference((settle,))
+        if settle.cancelled() or settle.exception() is None:
+            return
+        verbose_proxy_logger.debug(
+            "Rolling back a transaction abandoned by a cancelled caller failed: %s", settle.exception()
+        )
+
+    async def _tracked_query(self, content: str, tx_id: str | None) -> object:
         try:
             if db_io_claimed():
                 record_db_io()
@@ -121,9 +184,8 @@ class _TrackedPrismaEngine:
         finally:
             self.tracker.end_operation()
 
-    async def start_transaction(self, *, content: str) -> str:
+    async def _tracked_start_transaction(self, content: str) -> str:
         record_db_io()
-        self.tracker.begin_operation()
         try:
             transaction_id: Final = await self._engine.start_transaction(content=content)
         except (Exception, asyncio.CancelledError):
@@ -132,18 +194,16 @@ class _TrackedPrismaEngine:
         self.tracker.transaction_started(transaction_id)
         return transaction_id
 
-    async def commit_transaction(self, tx_id: str) -> None:
+    async def _tracked_commit_transaction(self, tx_id: str) -> None:
         record_db_io()
-        self.tracker.begin_operation()
         try:
             await self._engine.commit_transaction(tx_id)
         finally:
             self.tracker.end_operation()
             self.tracker.transaction_finished(tx_id)
 
-    async def rollback_transaction(self, tx_id: str) -> None:
+    async def _tracked_rollback_transaction(self, tx_id: str) -> None:
         record_db_io()
-        self.tracker.begin_operation()
         try:
             await self._engine.rollback_transaction(tx_id)
         finally:

tests/unit/proxy/db/test_prisma_client.py

@@ -1,8 +1,10 @@
+import asyncio
 import json
 import os
 import signal
 import sys
 import urllib.parse
+from typing import Any, Final, cast
 from unittest.mock import AsyncMock, MagicMock, Mock, patch
 
 import pytest
@@ -10,7 +12,13 @@ from fastapi.testclient import TestClient
 
 
 
-from litellm.proxy.db.prisma_client import PrismaManager, PrismaWrapper, should_update_prisma_schema
+from litellm.proxy.db.prisma_client import (
+    PrismaManager,
+    PrismaWrapper,
+    _PrismaDrainTracker,
+    _TrackedPrismaEngine,
+    should_update_prisma_schema,
+)
 
 
 @pytest.fixture(autouse=True)
@@ -470,3 +478,97 @@ def test_without_proxy_extras_the_index_build_reports_failure_instead_of_raising
     monkeypatch.setitem(sys.modules, "litellm_proxy_extras.utils", None)
 
     assert run() is outcome
+
+
+class _SlowEngine:
+    """Engine double whose requests only answer once the test lets them, recording what the engine saw."""
+
+    def __init__(self) -> None:
+        self.gates: dict[str, asyncio.Event] = {}  # mutable-ok: one gate per request the test lets through
+        self.queries_finished = 0
+        self.requests_cancelled = 0
+        self.rolled_back: list[str] = []  # mutable-ok: test records
+        self.process = object()
+
+    def gate(self, request: str) -> asyncio.Event:
+        return self.gates.setdefault(request, asyncio.Event())
+
+    async def _answer_once_released(self, request: str) -> None:
+        try:
+            await self.gate(request).wait()
+        except asyncio.CancelledError:
+            self.requests_cancelled += 1
+            raise
+
+    async def query(self, content: str, *, tx_id: str | None) -> object:
+        await self._answer_once_released(content)
+        self.queries_finished += 1
+        return {"content": content, "tx_id": tx_id}
+
+    async def start_transaction(self, *, content: str) -> str:
+        await self._answer_once_released("start")
+        return "tx-opened-after-caller-left"
+
+    async def commit_transaction(self, tx_id: str) -> None:
+        raise AssertionError(f"unexpected commit of {tx_id}")
+
+    async def rollback_transaction(self, tx_id: str) -> None:
+        self.rolled_back.append(tx_id)
+
+
+def _tracked(engine: _SlowEngine) -> tuple[_TrackedPrismaEngine, _PrismaDrainTracker]:
+    tracker: Final = _PrismaDrainTracker()
+    return _TrackedPrismaEngine(cast(Any, engine), tracker), tracker
+
+
+@pytest.mark.asyncio
+async def test_cancelling_a_query_caller_lets_the_engine_finish_the_statement() -> None:
+    engine: Final = _SlowEngine()
+    tracked, tracker = _tracked(engine)
+    caller: Final = asyncio.create_task(tracked.query("x", tx_id=None))
+    await asyncio.sleep(0)
+    caller.cancel()
+    with pytest.raises(asyncio.CancelledError):
+        await caller
+    drained: Final = asyncio.create_task(tracker.wait_until_drained())
+    await asyncio.sleep(0)
+    assert engine.requests_cancelled == 0
+    assert not drained.done()
+    engine.gate("x").set()
+    await asyncio.wait_for(drained, timeout=1)
+    assert engine.queries_finished == 1
+
+
+@pytest.mark.asyncio
+async def test_transaction_opened_after_its_caller_was_cancelled_is_rolled_back() -> None:
+    engine: Final = _SlowEngine()
+    tracked, tracker = _tracked(engine)
+    caller: Final = asyncio.create_task(tracked.start_transaction(content="{}"))
+    await asyncio.sleep(0)
+    caller.cancel()
+    with pytest.raises(asyncio.CancelledError):
+        await caller
+    engine.gate("start").set()
+    await asyncio.wait_for(tracker.wait_until_drained(), timeout=1)
+    assert engine.requests_cancelled == 0
+    assert engine.rolled_back == ["tx-opened-after-caller-left"]
+
+
+@pytest.mark.asyncio
+async def test_a_query_issued_while_the_last_one_finishes_keeps_the_engine_drain_open() -> None:
+    engine: Final = _SlowEngine()
+    tracked, tracker = _tracked(engine)
+    first: Final = asyncio.create_task(tracked.query("first", tx_id=None))
+    await asyncio.sleep(0)
+    await asyncio.sleep(0)
+    drained: Final = asyncio.create_task(tracker.wait_until_drained())
+    await asyncio.sleep(0)
+    second: Final = asyncio.create_task(tracked.query("second", tx_id=None))
+    engine.gate("first").set()
+    await first
+    await asyncio.sleep(0)
+    assert not drained.done()
+    engine.gate("second").set()
+    await second
+    await asyncio.wait_for(drained, timeout=1)
+    assert engine.queries_finished == 2