fix(proxy): retry spend-log writes on any error before dropping the batch

update_spend_logs pops the batch off spend_log_transactions before writing and
deliberately never puts it back, so an attempt not made here is billing data
gone for good. The inner retry only caught DB_CONNECTION_ERROR_TYPES, which is
httpx transport errors, while the failures concurrency actually produces
(deadlock detected, serialization failure, statement and lock timeouts, pool
exhaustion) arrive as prisma errors and got zero attempts before the whole
batch, up to MAX_LOGS_PER_INTERVAL rows, was discarded

Retrying any error is safe here: the write is idempotent (create_many with
skip_duplicates), and a row Postgres rejects on its data never reaches this
handler because _create_spend_logs_with_poison_isolation drops it and returns,
so a poisoned batch cannot spin. Retries stay bounded by n_retry_times and the
batch is still surfaced once they are spent

basedpyright-code-budget.json

@@ -117,7 +117,7 @@
     "limit": 118
   },
   "reportUnnecessaryComparison": {
-    "limit": 699
+    "limit": 698
   },
   "reportUnnecessaryContains": {
     "limit": 5

litellm/proxy/utils.py

@@ -25,7 +25,6 @@ from litellm.constants import (
     SPEND_LOG_WRITE_BATCH_MAX_BYTES,
 )
 from litellm.proxy._types import (
-    DB_CONNECTION_ERROR_TYPES,
     DB_RETRY_SAFE_ERROR_TYPES,
     CommonProxyErrors,
     ProxyErrorTypes,
@@ -5561,18 +5560,27 @@ class ProxyUpdateSpend:
                             "%s logs processed. Remaining in queue: %s", len(logs_to_process), remaining_count
                         )
                     break
-                except DB_CONNECTION_ERROR_TYPES as e:
-                    if i is None:
-                        i = 0
+                except Exception as e:
+                    # The batch was popped off the queue before this write, and the
+                    # handler below deliberately does not put it back, so whatever is
+                    # not retried here is billing data lost for good. Retry on any
+                    # error rather than only DB_CONNECTION_ERROR_TYPES (transport):
+                    # the errors concurrency actually produces - deadlock detected,
+                    # serialization failure, statement/lock timeout, pool exhaustion -
+                    # reach us as prisma errors, and got zero attempts. Retrying them
+                    # is safe because the write is idempotent (create_many with
+                    # skip_duplicates), and a row Postgres rejects on its data never
+                    # gets here: _create_spend_logs_with_poison_isolation drops it and
+                    # returns, so a poisoned batch cannot spin.
+                    if i >= n_retry_times:
+                        raise
                     verbose_proxy_logger.warning(
-                        "Spend tracking - DB connection error writing spend logs, retry %d/%d. logs_count=%d, error=%s",
+                        "Spend tracking - error writing spend logs, retry %d/%d. logs_count=%d, error=%s",
                         i + 1,
                         n_retry_times,
                         len(logs_to_process),
                         str(e),
                     )
-                    if i >= n_retry_times:
-                        raise
                     await asyncio.sleep(2**i)
         except Exception as e:
             # Logs already removed from queue at start - don't put them back

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

@@ -290,6 +290,90 @@ async def test_update_spend_logs_failure_raises_after_retries(
         )
 
 
+@pytest.mark.asyncio
+async def test_update_spend_logs_retries_a_postgres_level_error_and_persists_the_batch(
+    mock_prisma_client: Any,
+    make_spend_log_row: Any,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """A deadlock / serialization failure must not cost the batch.
+
+    The batch is popped off ``spend_log_transactions`` before the write and the
+    failure handler deliberately does not put it back, so an attempt not made
+    here is billing data lost for good. Only transport errors
+    (``DB_CONNECTION_ERROR_TYPES``) used to be retried, while the errors
+    concurrency actually produces - deadlock detected, serialization failure,
+    lock timeout, pool exhaustion - arrive as prisma errors and got zero
+    attempts. Retrying is safe because the write is idempotent.
+    """
+    from prisma.errors import TransactionError
+
+    async def _fake_sleep(_: float) -> None:
+        return None
+
+    monkeypatch.setattr(utils_mod.asyncio, "sleep", _fake_sleep)
+
+    attempts: List[List[str]] = []
+    written: List[str] = []
+
+    async def _create_many(*, data: Any, skip_duplicates: bool) -> None:
+        ids = [row["request_id"] for row in data]
+        attempts.append(ids)
+        if len(attempts) == 1:
+            raise TransactionError("deadlock detected")
+        written.extend(ids)
+
+    mock_prisma_client.db.litellm_spendlogs.create_many = AsyncMock(side_effect=_create_many)
+    proxy_logging = MagicMock()
+    proxy_logging.failure_handler = AsyncMock()
+
+    await ProxyUpdateSpend.update_spend_logs(
+        n_retry_times=3,
+        prisma_client=mock_prisma_client,
+        db_writer_client=None,
+        proxy_logging_obj=proxy_logging,
+        logs_to_process=[make_spend_log_row(request_id="r1"), make_spend_log_row(request_id="r2")],
+    )
+
+    assert written == ["r1", "r2"], f"the batch was dropped instead of retried; attempts={attempts}"
+    assert len(attempts) == 2
+
+
+@pytest.mark.asyncio
+async def test_update_spend_logs_stops_retrying_a_postgres_level_error_at_the_limit(
+    mock_prisma_client: Any,
+    make_spend_log_row: Any,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """Retries stay bounded by ``n_retry_times``: a persistently failing write must
+    still surface rather than spin the flush job forever."""
+    from prisma.errors import TransactionError
+
+    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=TransactionError("deadlock detected")
+    )
+    proxy_logging = MagicMock()
+    proxy_logging.failure_handler = AsyncMock()
+
+    with pytest.raises(TransactionError):
+        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="r1")],
+        )
+
+    assert mock_prisma_client.db.litellm_spendlogs.create_many.await_count == 3, (
+        "expected the initial write plus n_retry_times retries"
+    )
+
+
 def _data_error(message: str) -> Any:
     from prisma.errors import DataError