fix(bedrock): map real batch record counts and guard zero-count retire

litellm/llms/bedrock/batches/handler.py

@@ -1,3 +1,4 @@
+from collections.abc import Mapping
 from datetime import datetime
 from typing import TYPE_CHECKING, Any, Final, cast
 
@@ -68,6 +69,19 @@ def _predict_output_file_uri(output_prefix: str, input_uri: str, job_id: str | N
     return f"{output_prefix}{job_id}/{input_basename}.out"
 
 
+def _record_counts_from_response(response: Mapping[str, object]) -> BatchRequestCounts | None:
+    total_records: Final = response.get("totalRecordCount")
+    if not isinstance(total_records, int):
+        return None
+    success_records: Final = response.get("successRecordCount")
+    error_records: Final = response.get("errorRecordCount")
+    return BatchRequestCounts(
+        total=total_records,
+        completed=success_records if isinstance(success_records, int) else 0,
+        failed=error_records if isinstance(error_records, int) else 0,
+    )
+
+
 def _to_epoch(value: Any) -> int | None:
     if value is None:
         return None
@@ -271,11 +285,11 @@ class BedrockBatchesHandler:
                 ``aws_external_id``). Unknown keys are ignored.
 
         Returns:
-            ``LiteLLMBatch`` shaped like an OpenAI Batch resource. Note that
-            ``request_counts`` is always ``(0, 0, 0)`` because
-            ``GetModelInvocationJob`` does not surface per-record counts;
-            callers that need accurate counts should parse
-            ``manifest.json.out`` from the output S3 prefix.
+            ``LiteLLMBatch`` shaped like an OpenAI Batch resource.
+            ``request_counts`` maps ``GetModelInvocationJob``'s
+            ``totalRecordCount`` / ``successRecordCount`` / ``errorRecordCount``
+            when the provider reports them, and is ``None`` when it does not
+            (older botocore, or a status that omits counts).
         """
         try:
             import boto3
@@ -386,7 +400,7 @@ class BedrockBatchesHandler:
             failed_at=completed_at if openai_status == "failed" else None,
             cancelled_at=completed_at if openai_status == "cancelled" else None,
             expired_at=completed_at if openai_status == "expired" else None,
-            request_counts=BatchRequestCounts(total=0, completed=0, failed=0),
+            request_counts=_record_counts_from_response(response),
             metadata=openai_batch_metadata,
             completion_window="24h",
             endpoint="/v1/chat/completions",

litellm/proxy/openai_files_endpoints/common_utils.py

@@ -1351,15 +1351,16 @@ def _completed_batch_safe_to_retire(response: "LiteLLMBatch") -> bool:
     provider response briefly lags before the output id populates). Retiring in that
     window loses the spend record forever. Retire only once we can prove there is
     nothing left to recover: the output file has actually arrived, or the provider
-    reports no successful request lines. When counts are unknown, stay eligible so
-    the next poller pass revisits it. (#37713)
+    reported a positive total with zero successful request lines, proving it
+    enumerated the batch and none succeeded. A zero or unknown total means counts
+    are unreported, so stay eligible and let the next poller pass revisit it. (#37713)
     """
     if response.output_file_id is not None:
         return True
     request_counts = response.request_counts
     if request_counts is None:
         return False
-    return request_counts.completed == 0
+    return request_counts.total > 0 and request_counts.completed == 0
 
 
 async def update_batch_in_database(

tests/test_litellm/llms/bedrock/batches/test_handler.py

@@ -150,14 +150,38 @@ def test_handle_model_invocation_job_status_completed(patched_boto3):
     assert batch.completed_at == int(END_TIME.timestamp())
     assert batch.failed_at is None
     assert batch.cancelled_at is None
-    # Per-record counts aren't reported by GetModelInvocationJob, so we leave
-    # them zeroed; consumers should parse manifest.json.out for accurate counts.
-    assert batch.request_counts.total == 0
+    assert batch.request_counts is None
     assert batch.metadata["job_arn"] == JOB_ARN
     assert batch.metadata["output_file_uri"] == expected_out
     assert batch.metadata["output_s3_uri"] == OUTPUT_PREFIX
 
 
+@pytest.mark.parametrize("success_count,error_count", [(100, 0), (86, 14)])
+def test_completed_job_maps_provider_record_counts(patched_boto3, success_count, error_count):
+    fake_client, _ = patched_boto3
+    counted_response = _fake_boto3_response()
+    counted_response.update(totalRecordCount=100, successRecordCount=success_count, errorRecordCount=error_count)
+    fake_client.get_model_invocation_job.return_value = counted_response
+
+    batch = BedrockBatchesHandler._handle_model_invocation_job_status(batch_id=JOB_ARN)
+
+    assert batch.request_counts is not None
+    assert (batch.request_counts.total, batch.request_counts.completed, batch.request_counts.failed) == (
+        100,
+        success_count,
+        error_count,
+    )
+
+
+def test_missing_record_counts_leave_request_counts_none(patched_boto3):
+    fake_client, _ = patched_boto3
+    fake_client.get_model_invocation_job.return_value = _fake_boto3_response()
+
+    batch = BedrockBatchesHandler._handle_model_invocation_job_status(batch_id=JOB_ARN)
+
+    assert batch.request_counts is None
+
+
 @pytest.mark.parametrize(
     "bedrock_status,openai_status",
     [

tests/test_litellm/proxy/openai_files_endpoint/test_files_common_utils.py

@@ -430,13 +430,15 @@ def test_add_internal_model_credentials_survives_a_failing_deployment_lookup():
     assert data == {"batch_id": "unified-batch-id"}
 
 
+from openai.types.batch import BatchRequestCounts
+
 from litellm.proxy.openai_files_endpoints.common_utils import (
     _completed_batch_safe_to_retire,
 )
 
 
 def _completed_batch_for_retire(
-    output_file_id: str | None, completed: int | None = None
+    output_file_id: str | None, counts: BatchRequestCounts | None = None
 ) -> LiteLLMBatch:
     kwargs = dict(
         id="batch-1",
@@ -449,26 +451,30 @@ def _completed_batch_for_retire(
         output_file_id=output_file_id,
         error_file_id=None,
     )
-    if completed is not None:
-        kwargs["request_counts"] = {"total": completed, "completed": completed, "failed": 0}
+    if counts is not None:
+        kwargs["request_counts"] = counts
     return LiteLLMBatch(**kwargs)
 
 
 class TestCompletedBatchSafeToRetire:
     """A completed batch is only safe to retire from cost recovery once its output
-    file has arrived or the provider proves no successful lines (#37713)."""
+    file has arrived or the provider proves it enumerated a positive total of
+    request lines and none succeeded (#37713, LIT-6360)."""
 
     def test_output_file_present_is_safe(self):
         assert _completed_batch_safe_to_retire(_completed_batch_for_retire("file-out")) is True
 
-    def test_no_output_and_no_successful_lines_is_safe(self):
-        # Every request line errored -> nothing left to recover.
-        assert _completed_batch_safe_to_retire(_completed_batch_for_retire(None, completed=0)) is True
+    def test_no_output_and_synthesized_zero_counts_is_not_safe(self):
+        counts = BatchRequestCounts(total=0, completed=0, failed=0)
+        assert _completed_batch_safe_to_retire(_completed_batch_for_retire(None, counts)) is False
 
     def test_no_output_but_successful_lines_is_not_safe(self):
-        # The bug: output_file_id is lagging; retiring here loses the spend record.
-        assert _completed_batch_safe_to_retire(_completed_batch_for_retire(None, completed=5)) is False
+        counts = BatchRequestCounts(total=100, completed=100, failed=0)
+        assert _completed_batch_safe_to_retire(_completed_batch_for_retire(None, counts)) is False
+
+    def test_no_output_and_all_lines_failed_is_safe(self):
+        counts = BatchRequestCounts(total=100, completed=0, failed=100)
+        assert _completed_batch_safe_to_retire(_completed_batch_for_retire(None, counts)) is True
 
     def test_no_output_and_unknown_counts_is_not_safe(self):
-        # Counts unknown -> stay eligible so the next poller pass revisits it.
         assert _completed_batch_safe_to_retire(_completed_batch_for_retire(None)) is False