fix(proxy): drop only the daily spend batch whose failure cannot be re-sent, requeue the unsent ones

litellm/proxy/db/db_spend_update_writer.py

@@ -1332,15 +1332,7 @@ class DBSpendUpdateWriter:
                 daily_spend_transactions=cast(dict[str, _DailySpendTransactionT], transactions),
             )
         except Exception as e:  # noqa: BLE001  # whatever failed here, the other tables must still flush
-            if not _daily_spend_commit_failure_is_requeue_safe(e):
-                spend_log_error(
-                    "Spend tracking - dropped %d daily %s spend rows: the failed commit may have applied "
-                    "or the database refused the data, so re-sending it is not safe. Error: %s",
-                    len(transactions),
-                    entity_type,
-                    str(e),
-                    exc=e,
-                )
+            if not transactions:
                 return
             spend_log_error(
                 "Spend tracking - failed to commit daily %s spend updates. "
@@ -2050,13 +2042,25 @@ class DBSpendUpdateWriter:
                             sql, params = build_bulk_upsert(table=table, batch=merged_batch)
                             await prisma_client.db.execute_raw(sql, *params)
                         except Exception as batch_error:
-                            # Log detailed error information for debugging batch upsert failures
-                            # This helps diagnose issues like unique constraint violations
+                            if _daily_spend_commit_failure_is_requeue_safe(batch_error):
+                                spend_log_error(
+                                    "Daily %s spend batch upsert failed. Table: %s, Rows: %d, Error: %s",
+                                    entity_type,
+                                    table.name,
+                                    len(transactions_to_process),
+                                    str(batch_error),
+                                    exc=batch_error,
+                                )
+                                raise
+                            for key in transactions_to_process:
+                                daily_spend_transactions.pop(key, None)
                             spend_log_error(
-                                "Daily %s spend batch upsert failed. Table: %s, Rows: %d, Error: %s",
+                                "Spend tracking - dropped %d daily %s spend rows: the failed statement may have "
+                                "applied or the database refused the data, so re-sending it is not safe. "
+                                "Table: %s, Error: %s",
+                                len(transactions_to_process),
                                 entity_type,
                                 table.name,
-                                len(transactions_to_process),
                                 str(batch_error),
                                 exc=batch_error,
                             )

tests/test_litellm/proxy/db/test_db_spend_update_writer.py

@@ -1624,6 +1624,33 @@ async def test_update_daily_spend_keeps_failed_transactions_for_retry():
     assert daily_spend_transactions == expected
 
 
+@pytest.mark.asyncio
+async def test_update_daily_spend_drops_the_batch_whose_failure_cannot_be_resent():
+    """A reply lost after the statement was sent may already have applied, so the batch is
+    taken out of the caller's dict before the error propagates: whichever requeue the caller
+    runs afterwards, the Redis restore included, cannot send it a second time."""
+
+    def lose_the_reply() -> int:
+        raise httpx.ReadTimeout("no reply")
+
+    prisma_client = _RecordingPrisma(execute_raw=lose_the_reply)
+    daily_spend_transactions = {"user-key": _daily_txn(user_id="user-1")}
+    proxy_logging_obj = MagicMock()
+    proxy_logging_obj.failure_handler = AsyncMock()
+
+    with pytest.raises(httpx.ReadTimeout):
+        await DBSpendUpdateWriter._update_daily_spend(
+            n_retry_times=0,
+            prisma_client=prisma_client,
+            proxy_logging_obj=proxy_logging_obj,
+            daily_spend_transactions=daily_spend_transactions,
+            entity_type="user",
+            entity_id_field="user_id",
+        )
+
+    assert daily_spend_transactions == {}
+
+
 @pytest.mark.asyncio
 async def test_commit_key_spend_updates_includes_last_active():
     """
@@ -2875,6 +2902,33 @@ async def test_failed_daily_spend_commit_is_requeued_only_when_the_rows_are_prov
     assert db_writer.daily_spend_update_queue.update_queue.empty()
 
 
+@pytest.mark.asyncio
+async def test_failed_daily_spend_commit_drops_only_the_batch_that_was_sent():
+    """A tick holding more than one batch of 100 rows sends them one statement at a time, and
+    a reply lost on one statement says nothing about the batches after it: only the batch that
+    was on the wire is dropped, the ones never sent go back on the queue and land next tick."""
+    db_writer = DBSpendUpdateWriter()
+    await db_writer.daily_spend_update_queue.add_update(
+        {f"user-{i:03d}": _daily_txn(user_id=f"user-{i:03d}") for i in range(150)}
+    )
+    db = _DailySpendFakeDB(failing_table="LiteLLM_DailyUserSpend", failure=httpx.ReadTimeout("no reply"))
+    db_writer._flush_tool_discovery_queue = AsyncMock()
+    proxy_logging_obj = MagicMock()
+    proxy_logging_obj.failure_handler = AsyncMock()
+
+    await db_writer._commit_spend_updates_to_db_without_redis_buffer(
+        prisma_client=_WindowSpendFakePrisma(db), n_retry_times=0, proxy_logging_obj=proxy_logging_obj
+    )
+    db.failing_table = None
+    await db_writer._commit_spend_updates_to_db_without_redis_buffer(
+        prisma_client=_WindowSpendFakePrisma(db), n_retry_times=0, proxy_logging_obj=proxy_logging_obj
+    )
+
+    (upsert,) = _daily_upserts(db, "LiteLLM_DailyUserSpend")
+    assert _row_values(upsert, "user_id") == [f"user-{i:03d}" for i in range(100, 150)]
+    assert db_writer.daily_spend_update_queue.update_queue.empty()
+
+
 @pytest.mark.asyncio
 async def test_failed_daily_spend_commit_requeues_the_rows_and_flushes_the_other_tables():
     """With the Redis buffer off, a daily batch that failed to commit was discarded along