From 24f0be80219cd3403ec28cddbed9be946fad1cb0 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 5 Sep 2026 22:47:41 -0700 Subject: [PATCH] fix(spend): leave a batch uncharged when the database refuses the takeover The takeover of a $0 row an older proxy left behind used to charge the batch when the update could not reach the database. That leaves the row still reading $0, so every later retrieve finds the same row and charges the batch again, which is the repeat charging this PR exists to stop. The retrieve that does take the row over is the one that charges, and a batch nobody retrieves again after that failure is never charged, the same as one whose proxy died inside the write window. --- litellm/proxy/db/db_spend_update_writer.py | 8 +++++--- .../proxy/db/test_db_spend_update_writer.py | 12 ++++++++---- 2 files changed, 13 insertions(+), 7 deletions(-) diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index ae3bc5663eb..ee7802a45d3 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -412,9 +412,11 @@ class DBSpendUpdateWriter: "spend": 0.0, }, ) - except Exception as e: # noqa: BLE001 # prisma raises its own hierarchy; a row it cannot take over charges the batch - verbose_proxy_logger.warning("Could not take over spend row %s for a batch's cost: %s", request_id, e) - return True + except Exception as e: # noqa: BLE001 # prisma raises its own hierarchy; the next retrieve takes the row over + verbose_proxy_logger.warning( + "Could not take over spend row %s, leaving this batch's cost to the next retrieve: %s", request_id, e + ) + return False if taken_over == 0: verbose_proxy_logger.debug("Cost tracking skipped: spend row %s already charged this batch", request_id) return False diff --git a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py index 33b7af06e1a..4efb94b60aa 100644 --- a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py +++ b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py @@ -3093,17 +3093,21 @@ async def test_update_database_charges_a_batch_whose_row_a_pre_upgrade_poll_left @pytest.mark.asyncio -async def test_update_database_charges_a_batch_whose_zero_row_it_could_not_take_over(): - """A DB that refuses the takeover must not swallow the batch's cost.""" +async def test_update_database_leaves_a_batch_whose_zero_row_it_could_not_take_over_to_the_next_retrieve(): + """ + A DB that refuses the takeover leaves the row reading $0, so charging here would charge + the batch again on every later retrieve. The retrieve that does take the row over is the + one that charges. + """ db_writer = DBSpendUpdateWriter() db_writer._batch_database_updates = AsyncMock() existing = SimpleNamespace(call_type="aretrieve_batch", status="success", spend=0.0) prisma = _spend_logs_prisma(0, existing) prisma.db.litellm_spendlogs.update_many = AsyncMock(side_effect=RuntimeError("db unreachable")) - assert await _update_database_with(db_writer, prisma, _batch_cost_payload()) is True + assert await _update_database_with(db_writer, prisma, _batch_cost_payload()) is False - assert db_writer._batch_database_updates.await_count == 1 + assert db_writer._batch_database_updates.await_count == 0 @pytest.mark.asyncio