fix(bedrock): address PR feedback on model-invocation-job retrieve

Addresses Greptile P2 findings on #26834:

1. Use the bare job id (not the full ARN) when constructing the
   `api_base` URL for `pre_call` logging. Passing the full ARN double-
   counts the `model-invocation-job/` segment and embeds colons in the
   path, producing misleading log lines.

2. Drop the `or output_prefix` fallback when `_predict_output_file_uri`
   returns None. A bare prefix is not a downloadable object and surfacing
   it as `output_file_id` re-creates the very NoSuchKey bug this handler
   exists to fix. The bare prefix is still preserved in
   `metadata["output_s3_uri"]` for callers that want to do their own S3
   listing or read `manifest.json.out`.

   `metadata["output_file_uri"]` uses "" rather than None to satisfy the
   OpenAI Batch metadata schema (`dict[str, str]`); callers should branch
   on the typed `output_file_id` field instead.

Also expands test coverage on the new code path:
- new "stay None" regression test for the prediction-fail case
- pre_call/post_call logging hook assertions (incl. the bare-id URL)
- explicit cancelled_at / expired_at coverage
- _to_epoch type-handling matrix and the boto3 ImportError branch
- defensive _extract_region_from_bedrock_arn exception path
- empty-basename case for _predict_output_file_uri

Patch coverage on the changed lines is now 100% (the only remaining
uncovered lines in the file belong to the pre-existing
`_handle_async_invoke_status` method, which this PR does not touch).

Made-with: Cursor

litellm/llms/bedrock/batches/handler.py

@@ -247,6 +247,11 @@ class BedrockBatchesHandler:
         )
 
         if logging_obj is not None:
+            # Use the bare job id in the logged URL so we don't double up the
+            # `model-invocation-job/` segment when `batch_id` is a full ARN.
+            # `GetModelInvocationJob` accepts either form, but only the bare id
+            # produces a sensible-looking URL in logs.
+            url_path_id = _extract_job_id_from_arn(batch_id) or batch_id
             logging_obj.pre_call(
                 input=batch_id,
                 api_key="",
@@ -254,7 +259,7 @@ class BedrockBatchesHandler:
                     "complete_input_dict": {"jobIdentifier": batch_id},
                     "api_base": (
                         f"https://bedrock.{region}.amazonaws.com/"
-                        f"model-invocation-job/{batch_id}"
+                        f"model-invocation-job/{url_path_id}"
                     ),
                 },
             )
@@ -289,16 +294,23 @@ class BedrockBatchesHandler:
         # Bedrock returns the output *prefix* the user supplied at job creation.
         # Actual results land at <prefix>/<job-id>/<basename(input)>.out — we
         # surface that single-file URI as `output_file_id` so the OpenAI-style
-        # download flow works without an extra S3 listing call. The bare
-        # prefix is preserved in metadata for callers that want the manifest.
+        # download flow works without an extra S3 listing call. We deliberately
+        # do NOT fall back to the bare prefix when prediction fails: a prefix
+        # is not a downloadable object, so handing it back as `output_file_id`
+        # would reproduce the very NoSuchKey bug this handler exists to fix.
+        # The bare prefix is preserved in metadata for callers that want the
+        # `manifest.json.out` or want to do their own listing.
         job_arn = response.get("jobArn", batch_id)
         job_id = _extract_job_id_from_arn(job_arn)
-        output_file_uri = (
-            _predict_output_file_uri(output_prefix, input_uri, job_id) or output_prefix
-        )
+        output_file_uri = _predict_output_file_uri(output_prefix, input_uri, job_id)
 
         completed_at = _to_epoch(response.get("endTime"))
 
+        # Note: metadata uses "" (not None) for unknown URIs to satisfy the
+        # OpenAI Batch metadata schema, which is `dict[str, str]`. The
+        # `output_file_id` field on the LiteLLMBatch itself does carry None
+        # correctly (see below), so callers should branch on that, not on
+        # `metadata["output_file_uri"]`.
         openai_batch_metadata: OpenAIBatchMetadata = {
             "model_arn": response.get("modelId", ""),
             "job_arn": job_arn,
@@ -306,7 +318,7 @@ class BedrockBatchesHandler:
             "failure_message": response.get("message") or "",
             "input_s3_uri": input_uri,
             "output_s3_uri": output_prefix,
-            "output_file_uri": output_file_uri,
+            "output_file_uri": output_file_uri or "",
         }
 
         return LiteLLMBatch(

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

@@ -22,6 +22,7 @@ from litellm.llms.bedrock.batches.handler import (  # noqa: E402
     _extract_job_id_from_arn,
     _extract_region_from_bedrock_arn,
     _predict_output_file_uri,
+    _to_epoch,
 )
 
 JOB_ID = "abc1234567"
@@ -67,6 +68,42 @@ def test_extract_region_from_arn():
     assert _extract_region_from_bedrock_arn("not-an-arn") is None
 
 
+def test_extract_region_swallows_unexpected_split_errors():
+    """Defensive `except Exception` branch — anything that isn't a plain str
+    should fall through to ``None`` rather than blow up."""
+
+    class WeirdArn:
+        def split(self, _sep):
+            raise RuntimeError("boom")
+
+    assert _extract_region_from_bedrock_arn(WeirdArn()) is None  # type: ignore[arg-type]
+
+
+def test_predict_output_file_uri_returns_none_for_directory_input_uri():
+    """Input URI ending in `/` has an empty basename — we must bail rather
+    than emit ``<prefix>/<job-id>/.out``."""
+    assert (
+        _predict_output_file_uri(OUTPUT_PREFIX, "s3://bucket/inputs/", JOB_ID) is None
+    )
+
+
+_DT = datetime(2026, 4, 28, 12, 0, 0, tzinfo=timezone.utc)
+
+
+@pytest.mark.parametrize(
+    "value,expected",
+    [
+        (None, None),
+        (1730000000, 1730000000),
+        (1730000000.5, 1730000000),
+        (_DT, int(_DT.timestamp())),
+        ("2026-04-28T12:00:00Z", None),  # strings aren't supported -> None
+    ],
+)
+def test_to_epoch_handles_supported_types(value, expected):
+    assert _to_epoch(value) == expected
+
+
 def test_extract_job_id_from_arn():
     assert _extract_job_id_from_arn(JOB_ARN) == JOB_ID
     assert (
@@ -179,3 +216,123 @@ def test_failure_message_propagates(patched_boto3):
     assert batch.status == "failed"
     assert batch.failed_at == int(END_TIME.timestamp())
     assert batch.metadata["failure_message"] == "Input file failed validation"
+
+
+def test_completed_with_unpredictable_output_uri_stays_none(patched_boto3):
+    """
+    Regression guard for the original NoSuchKey bug: if Bedrock's response is
+    missing pieces we need to compute the per-job output file path (here, the
+    input s3Uri), `output_file_id` must stay `None` rather than fall back to
+    the bare prefix. Falling back to the prefix is what produced the original
+    NoSuchKey error this PR fixes.
+    """
+    fake_client, _ = patched_boto3
+    incomplete_response = _fake_boto3_response(status="Completed")
+    incomplete_response["inputDataConfig"] = {"s3InputDataConfig": {"s3Uri": ""}}
+    fake_client.get_model_invocation_job.return_value = incomplete_response
+
+    batch = BedrockBatchesHandler._handle_model_invocation_job_status(batch_id=JOB_ARN)
+
+    assert batch.status == "completed"
+    # output_file_id MUST be None (not the bare prefix) — that's the whole
+    # point of this regression test. Callers branch on this field.
+    assert batch.output_file_id is None
+    # The metadata field uses "" because OpenAI Batch metadata is dict[str, str];
+    # callers should branch on `output_file_id` (above) instead.
+    assert batch.metadata["output_file_uri"] == ""
+    # The bare prefix is still preserved in metadata so callers can list it.
+    assert batch.metadata["output_s3_uri"] == OUTPUT_PREFIX
+
+
+def test_cancelled_status_sets_cancelled_at(patched_boto3):
+    fake_client, _ = patched_boto3
+    fake_client.get_model_invocation_job.return_value = _fake_boto3_response(
+        status="Stopped"
+    )
+
+    batch = BedrockBatchesHandler._handle_model_invocation_job_status(batch_id=JOB_ARN)
+
+    assert batch.status == "cancelled"
+    assert batch.cancelled_at == int(END_TIME.timestamp())
+    assert batch.completed_at is None
+    assert batch.failed_at is None
+    assert batch.expired_at is None
+
+
+def test_expired_status_sets_expired_at(patched_boto3):
+    fake_client, _ = patched_boto3
+    fake_client.get_model_invocation_job.return_value = _fake_boto3_response(
+        status="Expired"
+    )
+
+    batch = BedrockBatchesHandler._handle_model_invocation_job_status(batch_id=JOB_ARN)
+
+    assert batch.status == "expired"
+    assert batch.expired_at == int(END_TIME.timestamp())
+    assert batch.completed_at is None
+    assert batch.failed_at is None
+    assert batch.cancelled_at is None
+
+
+def test_logging_obj_pre_and_post_call_invoked(patched_boto3):
+    """`pre_call` / `post_call` get called with sensible payloads when a
+    `logging_obj` is supplied."""
+    _, _ = patched_boto3
+    logging_obj = MagicMock()
+
+    BedrockBatchesHandler._handle_model_invocation_job_status(
+        batch_id=JOB_ARN, logging_obj=logging_obj
+    )
+
+    logging_obj.pre_call.assert_called_once()
+    logging_obj.post_call.assert_called_once()
+
+    pre_kwargs = logging_obj.pre_call.call_args.kwargs
+    assert pre_kwargs["input"] == JOB_ARN
+    assert pre_kwargs["additional_args"]["complete_input_dict"] == {
+        "jobIdentifier": JOB_ARN
+    }
+    # Logged URL must use the bare job id, not the full ARN, so it doesn't
+    # double the `model-invocation-job/` segment or embed colons in the path.
+    assert pre_kwargs["additional_args"]["api_base"] == (
+        f"https://bedrock.us-west-2.amazonaws.com/model-invocation-job/{JOB_ID}"
+    )
+
+    post_kwargs = logging_obj.post_call.call_args.kwargs
+    assert post_kwargs["input"] == JOB_ARN
+    assert post_kwargs["original_response"]["jobArn"] == JOB_ARN
+
+
+def test_missing_boto3_raises_helpful_import_error():
+    """If boto3 isn't installed we should raise a clear, actionable
+    ImportError rather than letting a NameError escape."""
+    real_import = (
+        __builtins__["__import__"]
+        if isinstance(__builtins__, dict)
+        else __builtins__.__import__
+    )
+
+    def fake_import(name, *args, **kwargs):
+        if name == "boto3":
+            raise ImportError("No module named 'boto3'")
+        return real_import(name, *args, **kwargs)
+
+    with patch("builtins.__import__", side_effect=fake_import):
+        with pytest.raises(ImportError, match="pip install boto3"):
+            BedrockBatchesHandler._handle_model_invocation_job_status(batch_id=JOB_ARN)
+
+
+def test_logging_url_uses_bare_id_when_only_id_passed(patched_boto3):
+    """If the caller passes just the trailing job id (also valid for
+    `GetModelInvocationJob`), the logged URL should use it as-is."""
+    _, _ = patched_boto3
+    logging_obj = MagicMock()
+
+    BedrockBatchesHandler._handle_model_invocation_job_status(
+        batch_id=JOB_ID, aws_region_name="us-west-2", logging_obj=logging_obj
+    )
+
+    pre_kwargs = logging_obj.pre_call.call_args.kwargs
+    assert pre_kwargs["additional_args"]["api_base"] == (
+        f"https://bedrock.us-west-2.amazonaws.com/model-invocation-job/{JOB_ID}"
+    )