Merge pull request #22982 from BerriAI/litellm_fix_batch_stale_status

Fix batch list showing stale "validating" status after completion

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,10 +235,16 @@ 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"},
-                )
+                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}"
+                    )

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()
@@ -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"
+        )