From 7361b097befab8aa0f513c1902f9c4fe013bfcd3 Mon Sep 17 00:00:00 2001 From: sailikhithk Date: Mon, 17 Aug 2026 18:53:38 -0500 Subject: [PATCH] fix(batches): bill cancelled/failed batches stamped terminal by a client poll A client polling GET /v1/batches/{id} can stamp the provider's terminal status (cancelled/failed) before the batch cost job runs. The pickup query excluded cancelled/failed from the not_in list, so the job never revisited that row and the spend for the completed requests was permanently lost. Remove cancelled/failed from the exclusion so the cost job picks them up like complete/completed. The staleness sweep is extended to also retire old failed/cancelled rows with batch_processed=False, serving as the upgrade guard so pre-existing terminal rows do not spike on deploy. Fixes #37217 --- .../proxy/common_utils/check_batch_cost.py | 21 ++++-- .../proxy_unit_tests/test_check_batch_cost.py | 73 +++++++++++++++++-- 2 files changed, 80 insertions(+), 14 deletions(-) diff --git a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py index a8e46349917..63279856875 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py @@ -193,18 +193,22 @@ class CheckBatchCost: # A row already in a terminal status is never rewritten by the sweep above, so # without this it keeps a poll-page slot forever and starves newer batches. + # failed/cancelled are included because a client poll can stamp them before the + # cost job bills the completed portion (issue #37217); rows older than the cutoff + # were never billed and never will be, so retire them instead of billing a spike + # on deploy. retired: Final = await self.prisma_client.db.litellm_managedobjecttable.update_many( where={ "file_purpose": "batch", "batch_processed": False, - "status": {"in": ["complete", "completed"]}, + "status": {"in": ["complete", "completed", "failed", "cancelled"]}, "created_at": {"lt": cutoff}, }, data={"batch_processed": True}, ) if retired > 0: verbose_proxy_logger.warning( - f"CheckBatchCost: gave up on {retired} completed managed objects older than " + f"CheckBatchCost: gave up on {retired} terminal managed objects older than " f"{MANAGED_OBJECT_STALENESS_CUTOFF_DAYS} days that were never costed" ) @@ -792,19 +796,20 @@ class CheckBatchCost: # every subsequent poll cycle. if self._has_batch_processed_column: try: - # Include "complete"/"completed" batches: the retrieve_batch - # endpoint may transition a batch to "complete" before - # CheckBatchCost runs. The batch_processed=False filter - # already prevents reprocessing finished batches. + # Include "complete"/"completed"/"failed"/"cancelled" batches: a + # client polling GET /v1/batches/{id} can stamp any of these terminal + # statuses before CheckBatchCost runs (issue #37217). A cancelled or + # failed batch may still carry an output_file_id with completed work that + # must be billed. The batch_processed=False filter prevents + # reprocessing, and the staleness sweep retires old terminal rows before + # this query runs so they cannot spike or starve newer batches. jobs = await self.prisma_client.db.litellm_managedobjecttable.find_many( where={ "file_purpose": "batch", "batch_processed": False, "status": { "not_in": [ - "failed", "expired", - "cancelled", "stale_expired", ] }, diff --git a/tests/proxy_unit_tests/test_check_batch_cost.py b/tests/proxy_unit_tests/test_check_batch_cost.py index 1dbbbfc43a0..fa2f6e2c860 100644 --- a/tests/proxy_unit_tests/test_check_batch_cost.py +++ b/tests/proxy_unit_tests/test_check_batch_cost.py @@ -174,13 +174,17 @@ class TestCheckBatchCost: assert find_call[1]["order"] == {"created_at": "asc"} not_in = find_call[1]["where"]["status"]["not_in"] assert "stale_expired" in not_in - # "complete"/"completed" are intentionally NOT excluded from the - # primary query — the batch_processed=False filter is sufficient. - # This allows CheckBatchCost to pick up batches that were - # transitioned to "complete" by the retrieve_batch endpoint - # before CheckBatchCost had a chance to process them. + assert "expired" in not_in + # "complete"/"completed"/"failed"/"cancelled" are intentionally NOT excluded + # from the primary query. A client polling GET /v1/batches/{id} can stamp any + # of these terminal statuses before CheckBatchCost runs (issue #37217), and a + # cancelled/failed batch may still carry an output_file_id with completed work + # that must be billed. The batch_processed=False filter prevents reprocessing, + # and the staleness sweep retires old terminal rows before this query runs. assert "complete" not in not_in assert "completed" not in not_in + assert "failed" not in not_in + assert "cancelled" not in not_in assert find_call[1]["where"]["batch_processed"] is False assert check_batch_cost_instance.batch_processed_support_confirmed is True @@ -1268,6 +1272,63 @@ class TestCheckBatchCost: update_data["status"] == terminal_status ), f"billed {terminal_status} batch must keep its real terminal status in the DB" + @pytest.mark.asyncio + @pytest.mark.parametrize("terminal_status", ["cancelled", "failed"]) + async def test_client_poll_terminal_status_still_picked_up_by_query( + self, + check_batch_cost_instance, + mock_prisma_client, + terminal_status, + ): + """Regression for issue #37217: a client polling GET /v1/batches/{id} stamps + the provider's terminal status (cancelled/failed) before the cost job runs. + The pickup query must still select that row so the completed portion is billed. + Before the fix, cancelled/failed were in the not_in list and the row was never + queried again, permanently losing the spend for the completed requests. + """ + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( + return_value=0 + ) + mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock( + return_value=[] + ) + + await check_batch_cost_instance.check_batch_cost() + + find_call = mock_prisma_client.db.litellm_managedobjecttable.find_many.call_args + not_in = find_call[1]["where"]["status"]["not_in"] + assert terminal_status not in not_in, ( + f"{terminal_status} must not be excluded from the pickup query, " + "otherwise a client poll that stamps it first permanently loses the spend" + ) + + @pytest.mark.asyncio + async def test_stale_sweep_retires_old_failed_and_cancelled_rows( + self, check_batch_cost_instance, mock_prisma_client + ): + """The staleness sweep must also retire old failed/cancelled rows with + batch_processed=False so they don't spike or starve newer batches on deploy + when the pickup query starts including them (issue #37217 upgrade guard). + """ + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( + return_value=0 + ) + mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock( + return_value=[] + ) + + await check_batch_cost_instance.check_batch_cost() + + calls = ( + mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args_list + ) + retired_call = calls[1] + retired_statuses = retired_call[1]["where"]["status"]["in"] + assert "failed" in retired_statuses + assert "cancelled" in retired_statuses + assert retired_call[1]["where"]["batch_processed"] is False + assert retired_call[1]["data"] == {"batch_processed": True} + @pytest.mark.asyncio async def test_terminal_batch_with_missing_output_file_is_retired_unbilled( self, check_batch_cost_instance, mock_prisma_client, mock_llm_router @@ -2502,7 +2563,7 @@ class TestPollPageStarvation: where = calls[1][1]["where"] assert where["file_purpose"] == "batch" assert where["batch_processed"] is False - assert where["status"] == {"in": ["complete", "completed"]} + assert where["status"] == {"in": ["complete", "completed", "failed", "cancelled"]} assert "created_at" in where assert calls[1][1]["data"] == {"batch_processed": True}