From eef5bd1ad227b3398d35bb6eee3c85ed1239012d Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 8 Jun 2026 19:36:58 +0000 Subject: [PATCH] test(batches): add failing regression for completed-batch cost tracking gap Adds a regression test demonstrating that a user GET that transitions a batch to complete marks batch_processed=True on a plain status sync. Because CheckBatchCost only picks up batches with batch_processed=False, any batch whose asynchronous cost callback failed is permanently skipped and its cost is never recorded, causing billing discrepancies. The test encodes the expected behavior (a status sync must not mark a batch cost-tracked) and currently fails against the existing code, pinpointing the bug to fix. It will pass once batch_processed is set only from the cost-recording path. --- ..._batch_update_db_managed_output_file_id.py | 51 +++++++++++++++++++ 1 file changed, 51 insertions(+) diff --git a/tests/test_litellm/enterprise/proxy/test_batch_update_db_managed_output_file_id.py b/tests/test_litellm/enterprise/proxy/test_batch_update_db_managed_output_file_id.py index d8669960674..aa48b37e418 100644 --- a/tests/test_litellm/enterprise/proxy/test_batch_update_db_managed_output_file_id.py +++ b/tests/test_litellm/enterprise/proxy/test_batch_update_db_managed_output_file_id.py @@ -52,6 +52,57 @@ def _build_prisma_mock(): return mock +@pytest.mark.asyncio +async def test_status_sync_to_complete_does_not_mark_batch_cost_tracked(): + """ + Regression for the batch cost-tracking gap. + + When a user GET /v1/batches/{id} transitions a batch from validating -> + completed, the endpoint calls update_batch_in_database to sync the status. The + cost callback for that retrieve is enqueued asynchronously and can fail. The + batch_processed flag means "cost has been recorded" and is what CheckBatchCost + uses (batch_processed=False) to find batches whose cost still needs tracking. + + If update_batch_in_database marks batch_processed=True just for syncing status, + a batch whose async cost callback failed is permanently hidden from CheckBatchCost + and its cost is never recorded. So a plain status sync to complete must NOT set + batch_processed=True; that flag belongs to the cost-recording path. This test + fails while the bug is present and passes once the flag is no longer set here. + """ + batch_id = "batch_complete_status_sync" + unified_batch_id = ( + "litellm_proxy;model_id:m;llm_batch_id:batch_complete_status_sync" + ) + response = _build_batch_response( + batch_id=batch_id, + status="completed", + output_file_id=None, + hidden_params={"model_id": "m", "model_name": "openai/gpt-4o"}, + ) + + mock_prisma = _build_prisma_mock() + + await update_batch_in_database( + batch_id=batch_id, + unified_batch_id=unified_batch_id, + response=response, + managed_files_obj=_build_managed_files_mock(), + prisma_client=mock_prisma, + verbose_proxy_logger=MagicMock(), + user_api_key_dict=UserAPIKeyAuth(user_id="user-abc"), + ) + + update_data = mock_prisma.db.litellm_managedobjecttable.update.call_args.kwargs[ + "data" + ] + assert update_data["status"] == "complete" + assert update_data.get("batch_processed") is not True, ( + "A status sync to complete must not mark batch_processed=True; doing so hides " + "batches whose async cost callback failed from CheckBatchCost, so their cost " + "is never recorded." + ) + + @pytest.mark.asyncio async def test_update_batch_in_database_stores_unified_output_file_id(): raw_output_file_id = "file-rawoutput789"