From 685c3ec6a29cedff43ce6d7de117f14aeb9e928e Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Wed, 14 Oct 2026 15:09:16 -0700 Subject: [PATCH] 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 | 76 ++++++++++++++-- tests/unit/proxy/db/test_prisma_client.py | 104 +++++++++++++++++++++- 2 files changed, 171 insertions(+), 9 deletions(-) diff --git a/litellm/proxy/db/prisma_client.py b/litellm/proxy/db/prisma_client.py index 593f7c0503f..0cbce6339fd 100644 --- a/litellm/proxy/db/prisma_client.py +++ b/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: diff --git a/tests/unit/proxy/db/test_prisma_client.py b/tests/unit/proxy/db/test_prisma_client.py index 7b8d000a8d5..99b500ec315 100644 --- a/tests/unit/proxy/db/test_prisma_client.py +++ b/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