fix(proxy): retry deadlocks and requeue spend logs on any DB write error (#39883)

* fix(proxy): retry deadlocks and requeue spend logs on any DB write error

update_spend_logs dequeued the batch and only retried/requeued on transport
errors. A 40P01 deadlock surfaced as a plain prisma DataError and went through
poison-row isolation, which dropped every row it hit; every other DB error was
re-raised with the batch already gone from the queue.

Treat deadlocks as transient (retry, then requeue), keep them out of poison-row
isolation, and requeue the batch at the head of the queue on any other prisma
error so it lands once the DB is healthy.

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

* refactor(proxy): drop redundant docstrings and tighten test typing for spend-log requeue

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

* test(proxy): assert deadlock retries from mock call history instead of mutable lists

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

---------

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

litellm/proxy/db/exception_handler.py

@@ -199,6 +199,12 @@ class PrismaDBExceptionHandler:
             return True
         return False
 
+    @staticmethod
+    def is_prisma_error(e: Exception) -> bool:
+        import prisma
+
+        return isinstance(e, _exception_types(prisma.errors.PrismaError))
+
     @staticmethod
     def is_deadlock_error(e: Exception) -> bool:
         """True iff ``e`` is a Postgres deadlock (P2034 / 40P01) surfaced through prisma."""

litellm/proxy/utils.py

@@ -6386,10 +6386,17 @@ class ProxyUpdateSpend:
                         )
                     break
                 except Exception as e:
-                    if not PrismaDBExceptionHandler.is_database_transport_error(e):
+                    if not _is_transient_spend_log_write_error(e):
+                        if PrismaDBExceptionHandler.is_prisma_error(e):
+                            await enqueue_spend_logs(prisma_client, logs_to_process, at_head=True)
+                            verbose_proxy_logger.warning(
+                                "Spend tracking - DB error writing spend logs, requeued %d rows for the next flush. error=%s",
+                                len(logs_to_process),
+                                str(e),
+                            )
                         raise
                     verbose_proxy_logger.warning(
-                        "Spend tracking - DB connection error writing spend logs, retry %d/%d. logs_count=%d, error=%s",
+                        "Spend tracking - transient DB error writing spend logs, retry %d/%d. logs_count=%d, error=%s",
                         i + 1,
                         n_retry_times,
                         len(logs_to_process),
@@ -6732,6 +6739,10 @@ async def _monitor_spend_logs_queue(
 MAX_SPEND_LOG_ISOLATION_FAILURES_PER_BATCH: Final = 256
 
 
+def _is_transient_spend_log_write_error(e: Exception) -> bool:
+    return PrismaDBExceptionHandler.is_database_transport_error(e) or PrismaDBExceptionHandler.is_deadlock_error(e)
+
+
 async def _create_spend_logs_with_poison_isolation(
     repo: SpendLogsRepository,
     rows: Sequence[Mapping[str, object]],
@@ -6767,6 +6778,8 @@ async def _create_spend_logs_with_poison_isolation(
             raise
         if PrismaDBExceptionHandler.is_database_service_unavailable_error(e):
             raise
+        if PrismaDBExceptionHandler.is_deadlock_error(e):
+            raise
         budget_left: Final = max(failure_budget - 1, 0)
         if len(rows) == 1:
             request_id: Final = rows[0].get("request_id")

tests/test_litellm/proxy/utils/prisma_and_spend/test_proxy_update_spend.py

@@ -473,6 +473,110 @@ async def test_update_spend_logs_retries_and_requeues_batch_on_db_outage(
     assert [row["request_id"] for row in mock_prisma_client.spend_log_transactions] == ["a", "b", "c"]
 
 
+def _deadlock_error() -> Exception:
+    return _data_error(
+        'Error occurred during query execution: ConnectorError(ConnectorError { user_facing_error: None, '
+        'kind: QueryError(PostgresError { code: "40P01", message: "deadlock detected", severity: "ERROR" }) })'
+    )
+
+
+@pytest.mark.asyncio
+async def test_update_spend_logs_retries_deadlock_and_keeps_every_row(
+    mock_prisma_client: Any, make_spend_log_row: Any, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    """A 40P01 deadlock aborts the whole insert, so the same rows succeed on replay.
+    Before the fix the deadlock surfaced as a plain ``DataError`` and went through
+    poison-row isolation, which bisected the batch and dropped every row the
+    deadlock happened to hit as if Postgres had rejected it.
+    """
+
+    async def _fake_sleep(_: float) -> None:
+        return None
+
+    monkeypatch.setattr(utils_mod.asyncio, "sleep", _fake_sleep)
+    create_many = AsyncMock(side_effect=[_deadlock_error(), _deadlock_error(), None])
+    mock_prisma_client.db.litellm_spendlogs.create_many = create_many
+    proxy_logging = MagicMock()
+    proxy_logging.failure_handler = AsyncMock()
+    mock_prisma_client.spend_log_transactions = []
+
+    await ProxyUpdateSpend.update_spend_logs(
+        n_retry_times=2,
+        prisma_client=mock_prisma_client,
+        db_writer_client=None,
+        proxy_logging_obj=proxy_logging,
+        logs_to_process=[make_spend_log_row(request_id="a"), make_spend_log_row(request_id="b")],
+    )
+
+    attempts = tuple(
+        tuple(row["request_id"] for row in call.kwargs["data"]) for call in create_many.await_args_list
+    )
+    assert attempts == (("a", "b"), ("a", "b"), ("a", "b"))
+    assert mock_prisma_client.spend_log_transactions == []
+
+
+@pytest.mark.asyncio
+async def test_update_spend_logs_requeues_batch_once_deadlock_retries_exhaust(
+    mock_prisma_client: Any, make_spend_log_row: Any, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    """If every retry deadlocks, the batch goes back to the head of the queue for
+    the next flush instead of being dropped.
+    """
+
+    async def _fake_sleep(_: float) -> None:
+        return None
+
+    monkeypatch.setattr(utils_mod.asyncio, "sleep", _fake_sleep)
+    mock_prisma_client.db.litellm_spendlogs.create_many = AsyncMock(side_effect=_deadlock_error())
+    proxy_logging = MagicMock()
+    proxy_logging.failure_handler = AsyncMock()
+    mock_prisma_client.spend_log_transactions = [make_spend_log_row(request_id="c")]
+
+    with pytest.raises(type(_deadlock_error())):
+        await ProxyUpdateSpend.update_spend_logs(
+            n_retry_times=1,
+            prisma_client=mock_prisma_client,
+            db_writer_client=None,
+            proxy_logging_obj=proxy_logging,
+            logs_to_process=[make_spend_log_row(request_id="a"), make_spend_log_row(request_id="b")],
+        )
+
+    assert mock_prisma_client.db.litellm_spendlogs.create_many.await_count == 2
+    assert [row["request_id"] for row in mock_prisma_client.spend_log_transactions] == ["a", "b", "c"]
+
+
+@pytest.mark.asyncio
+async def test_update_spend_logs_requeues_batch_on_non_transport_db_error(
+    mock_prisma_client: Any, make_spend_log_row: Any
+) -> None:
+    """A DB error that is neither transport nor deadlock (here P2021, the table is
+    gone mid-migration) is not retried in place, but the dequeued batch must not
+    be lost either: it goes back to the head of the queue so it lands once the
+    DB is healthy again.
+    """
+    from prisma.errors import TableNotFoundError
+
+    err = TableNotFoundError(
+        {"user_facing_error": {"error_code": "P2021", "message": "The table does not exist", "meta": {}}}
+    )
+    mock_prisma_client.db.litellm_spendlogs.create_many = AsyncMock(side_effect=err)
+    proxy_logging = MagicMock()
+    proxy_logging.failure_handler = AsyncMock()
+    mock_prisma_client.spend_log_transactions = [make_spend_log_row(request_id="c")]
+
+    with pytest.raises(TableNotFoundError):
+        await ProxyUpdateSpend.update_spend_logs(
+            n_retry_times=2,
+            prisma_client=mock_prisma_client,
+            db_writer_client=None,
+            proxy_logging_obj=proxy_logging,
+            logs_to_process=[make_spend_log_row(request_id="a"), make_spend_log_row(request_id="b")],
+        )
+
+    assert mock_prisma_client.db.litellm_spendlogs.create_many.await_count == 1
+    assert [row["request_id"] for row in mock_prisma_client.spend_log_transactions] == ["a", "b", "c"]
+
+
 @pytest.mark.asyncio
 async def test_requeue_after_outage_drops_oldest_logs_past_the_byte_budget(
     mock_prisma_client: Any, make_spend_log_row: Any
@@ -549,8 +653,9 @@ async def test_flush_returns_the_bytes_it_took_off_the_queue(mock_prisma_client:
 async def test_update_spend_logs_does_not_requeue_non_transport_failures(
     mock_prisma_client: Any, make_spend_log_row: Any
 ) -> None:
-    """Only transport failures are worth replaying. A rejection the DB will keep
-    rejecting must not be requeued, or it would wedge the queue forever.
+    """Only DB failures are worth replaying. A row the proxy itself cannot
+    serialize would fail the same way on every flush, so requeueing it would
+    wedge the head of the queue forever.
     """
     mock_prisma_client.db.litellm_spendlogs.create_many = AsyncMock(side_effect=ValueError("bad payload"))
     proxy_logging = MagicMock()