From 3d55f7f6abf60bd21ce02e70a1c6d79fd15e9e13 Mon Sep 17 00:00:00 2001 From: Ryan Crabbe Date: Fri, 6 Mar 2026 08:54:21 -0800 Subject: [PATCH 1/2] Fix batch list showing stale "validating" status after completion CheckBatchCost poller updated the status column but not the file_object JSON column. The list_batches endpoint reads status from file_object, so batches appeared stuck in "validating" even after Azure reported them as completed. Now update file_object alongside status in the per-job DB write. --- .../proxy/common_utils/check_batch_cost.py | 11 +++++++---- .../proxy/test_managed_files_access_check.py | 2 +- 2 files changed, 8 insertions(+), 5 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 4dcabb9c58..d4ab1d881d 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py @@ -239,8 +239,11 @@ class CheckBatchCost: # mark the job as complete completed_jobs.append(job) - if len(completed_jobs) > 0: - await self.prisma_client.db.litellm_managedobjecttable.update_many( - where={"id": {"in": [job.id for job in completed_jobs]}}, - data={"batch_processed": True, "status": "complete"}, + await self.prisma_client.db.litellm_managedobjecttable.update( + where={"id": job.id}, + data={ + "batch_processed": True, + "status": "complete", + "file_object": response.model_dump_json(), + }, ) diff --git a/tests/test_litellm/enterprise/proxy/test_managed_files_access_check.py b/tests/test_litellm/enterprise/proxy/test_managed_files_access_check.py index 2db5a2214c..7b7643aeec 100644 --- a/tests/test_litellm/enterprise/proxy/test_managed_files_access_check.py +++ b/tests/test_litellm/enterprise/proxy/test_managed_files_access_check.py @@ -134,7 +134,7 @@ async def test_check_batch_cost_should_call_afile_content_directly_with_credenti mock_prisma.db.litellm_managedobjecttable.find_many = AsyncMock( return_value=[mock_job] ) - mock_prisma.db.litellm_managedobjecttable.update_many = AsyncMock() + mock_prisma.db.litellm_managedobjecttable.update = AsyncMock() # Mock proxy_logging_obj — should NOT be called for file content mock_proxy_logging = MagicMock() From c4db53a98a65474ffb3a58c795211fb858d8e057 Mon Sep 17 00:00:00 2001 From: Ryan Crabbe Date: Fri, 6 Mar 2026 09:25:50 -0800 Subject: [PATCH 2/2] Address review feedback: remove dead code, add error handling, strengthen test assertions - Remove unused `completed_jobs` list (dead code after per-job update refactor) - Wrap DB update in try/except to prevent one failed update from aborting remaining jobs - Add test assertions verifying batch_processed, status, and file_object are written to DB --- .../proxy/common_utils/check_batch_cost.py | 25 ++++++++++--------- .../proxy/test_managed_files_access_check.py | 9 +++++++ 2 files changed, 22 insertions(+), 12 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 d4ab1d881d..10f7f98b71 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py @@ -78,8 +78,6 @@ class CheckBatchCost: "status": {"not_in": ["failed", "expired", "cancelled"]} } ) - completed_jobs = [] - for job in jobs: # get the model from the job unified_object_id = job.unified_object_id @@ -237,13 +235,16 @@ class CheckBatchCost: ) # mark the job as complete - completed_jobs.append(job) - - await self.prisma_client.db.litellm_managedobjecttable.update( - where={"id": job.id}, - data={ - "batch_processed": True, - "status": "complete", - "file_object": response.model_dump_json(), - }, - ) + try: + await self.prisma_client.db.litellm_managedobjecttable.update( + where={"id": job.id}, + data={ + "batch_processed": True, + "status": "complete", + "file_object": response.model_dump_json(), + }, + ) + except Exception as db_err: + verbose_proxy_logger.error( + f"CheckBatchCost: failed to mark job {job.id} complete in DB: {db_err}" + ) diff --git a/tests/test_litellm/enterprise/proxy/test_managed_files_access_check.py b/tests/test_litellm/enterprise/proxy/test_managed_files_access_check.py index 7b7643aeec..8cb642b7a4 100644 --- a/tests/test_litellm/enterprise/proxy/test_managed_files_access_check.py +++ b/tests/test_litellm/enterprise/proxy/test_managed_files_access_check.py @@ -198,3 +198,12 @@ async def test_check_batch_cost_should_call_afile_content_directly_with_credenti # managed_files_obj.afile_content should NOT have been called mock_managed_files_hook.afile_content.assert_not_called() + + # Verify the DB update writes batch_processed, status, and file_object + mock_prisma.db.litellm_managedobjecttable.update.assert_called_once() + update_call_kwargs = mock_prisma.db.litellm_managedobjecttable.update.call_args.kwargs + assert update_call_kwargs["data"]["batch_processed"] is True + assert update_call_kwargs["data"]["status"] == "complete" + assert "file_object" in update_call_kwargs["data"], ( + "file_object must be written to DB so list_batches reads updated status" + )