feat(batch-job): bedrock batch model invocation job retrieval (#26834)

* feat(bedrock): support retrieve for model-invocation-job batch ARNs `bedrock.retrieve_batch` previously only handled `:async-invoke/` ARNs (Twelve Labs Marengo embeddings). The `:model-invocation-job/` ARNs returned by `CreateModelInvocationJob` (the bulk batch inference API behind `bedrock.create_batch`) fell through and returned a misleading data-plane error, leaving created jobs unretrievable through the LiteLLM batches API. The two ARN families live on different AWS service endpoints (`bedrock-runtime` data plane vs `bedrock` control plane), so they need distinct handlers. This adds: * `BedrockBatchesHandler._handle_model_invocation_job_status` — calls the control plane via boto3 (`bedrock:GetModelInvocationJob`), reusing `BaseAWSLLM.get_credentials` for credential resolution so model_list / env / role-assumption configs continue to apply. The response is reshaped into a `LiteLLMBatch` with the same status mapping `transform_create_batch_response` already uses. * Output-file-URI prediction. Bedrock surfaces the user-supplied `s3OutputDataConfig.s3Uri` *prefix* in `GetModelInvocationJob`, but results actually land at `<prefix>/<job-id>/<basename(input)>.out`. We compute that single-file URI client-side and surface it as `output_file_id`, so OpenAI-style `client.files.content(...)` works without an extra `ListObjectsV2` round-trip. The bare prefix stays in metadata for callers that want the manifest. * Dispatch in `litellm/batches/main.py` for the new ARN family, alongside the existing async-invoke branch. * Unit tests covering ARN parsing, output-URI prediction (incl. edge cases), the full status mapping, region resolution precedence, and failure-message propagation. Note: `request_counts` is intentionally `(0, 0, 0)` — `GetModelInvocationJob` does not report per-record counts; getting accurate numbers requires parsing `manifest.json.out` from the output S3 prefix, which is left to callers. Made-with: Cursor * 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 * test(bedrock): cover retrieve_batch dispatch for both ARN families Codecov flagged 8 uncovered lines on `litellm/batches/main.py` after this PR refactored the Bedrock dispatch into a single guard with two sub-branches (`async-invoke` + `model-invocation-job`). Existing tests exercised the handlers directly but not the dispatch in `main.py`. Adds `tests/test_litellm/batches/test_retrieve_batch_bedrock_dispatch.py` with 6 mocked tests that exercise `litellm.retrieve_batch` end-to-end for the dispatch logic: - async-invoke ARN routes to `_handle_async_invoke_status` - async-invoke ARN with no region falls back to "us-east-1" (preserves prior behavior on this branch) - model-invocation-job ARN routes to the new `_handle_model_invocation_job_status` handler - model-invocation-job ARN with no region forwards None (so the new handler can sniff region from the ARN itself, rather than getting silently routed to us-east-1) - unrelated bedrock ARN family falls through to the generic provider-config retrieve path (neither special handler invoked) - non-bedrock batch ids skip the bedrock dispatch entirely Both handlers are mocked at the import site so the tests don't hit AWS — the focus here is purely the new dispatch logic in main.py. Co-authored-by: Cursor <cursoragent@cursor.com> * test(bedrock): move retrieve_batch dispatch test to tests/test_litellm/ The dispatch test landed under `tests/test_litellm/batches/`, a new directory that no upstream `test-unit-*.yml` workflow's `test-path` allow-list includes. As a result, the test was never executed in CI and codecov reported `litellm/batches/main.py` patch coverage at 11.11% (8 lines uncovered) — the lines belonging to this PR's dispatch refactor itself. Move the file up one level so it matches the `tests/test_litellm/test_*.py` glob that `test-unit-misc.yml` already runs, and adjust `sys.path.insert` for the new depth. The companion handler tests under `tests/test_litellm/llms/bedrock/batches/test_handler.py` are unaffected — they're picked up by the `llms` directory in `test-unit-llm-providers.yml`. Made-with: Cursor --------- Co-authored-by: Cursor <cursoragent@cursor.com>
2026-08-28 05:25:59 +00:00 · 2026-05-11 13:22:26 -07:00 · 2026-05-11 13:22:26 -07:00 · 0751886680
commit 0751886680
parent 5833d3eadd
4 changed files with 769 additions and 17 deletions
--- a/litellm/batches/main.py
+++ b/litellm/batches/main.py
@ -617,24 +617,35 @@ def retrieve_batch(
        _is_async = kwargs.pop("aretrieve_batch", False) is True
        client = kwargs.get("client", None)

-        # Check if this is an async invoke ARN (different from regular batch ARN)
-        # Async invoke ARNs have format: arn:aws(-[^:]+)?:bedrock:[a-z0-9-]{1,20}:[0-9]{12}:async-invoke/[a-z0-9]{12}
-        if (
-            batch_id.startswith("arn:aws")
-            and ":bedrock:" in batch_id
-            and ":async-invoke/" in batch_id
-        ):
-            # Handle async invoke status check
-            # Remove aws_region_name from kwargs to avoid duplicate parameter
-            async_kwargs = kwargs.copy()
-            async_kwargs.pop("aws_region_name", None)
+        # Bedrock has two distinct ARN families that need different APIs:
+        #   * async-invoke ARNs       (Twelve Labs Marengo embeddings)        -> bedrock-runtime data plane
+        #   * model-invocation-job ARNs (CreateModelInvocationJob batch)      -> bedrock control plane
+        # They live on different AWS service endpoints and can't share a handler.
+        # ARN shapes:
+        #   arn:aws(-[^:]+)?:bedrock:<region>:<account>:async-invoke/<id>
+        #   arn:aws(-[^:]+)?:bedrock:<region>:<account>:model-invocation-job/<id>
+        if batch_id.startswith("arn:aws") and ":bedrock:" in batch_id:
+            if ":async-invoke/" in batch_id:
+                # Remove aws_region_name from kwargs to avoid duplicate parameter
+                async_kwargs = kwargs.copy()
+                async_kwargs.pop("aws_region_name", None)

-            return BedrockBatchesHandler._handle_async_invoke_status(
-                batch_id=batch_id,
-                aws_region_name=kwargs.get("aws_region_name", "us-east-1"),
-                logging_obj=litellm_logging_obj,
-                **async_kwargs,
-            )
+                return BedrockBatchesHandler._handle_async_invoke_status(
+                    batch_id=batch_id,
+                    aws_region_name=kwargs.get("aws_region_name", "us-east-1"),
+                    logging_obj=litellm_logging_obj,
+                    **async_kwargs,
+                )
+            if ":model-invocation-job/" in batch_id:
+                mij_kwargs = kwargs.copy()
+                mij_kwargs.pop("aws_region_name", None)
+
+                return BedrockBatchesHandler._handle_model_invocation_job_status(
+                    batch_id=batch_id,
+                    aws_region_name=kwargs.get("aws_region_name"),
+                    logging_obj=litellm_logging_obj,
+                    **mij_kwargs,
+                )

        # Try to use provider config first (for providers like bedrock)
        model: Optional[str] = kwargs.get("model", None)
--- a/litellm/llms/bedrock/batches/handler.py
+++ b/litellm/llms/bedrock/batches/handler.py
@ -1,8 +1,79 @@
+from datetime import datetime
+from typing import Any, Optional, cast
+
 from openai.types.batch import BatchRequestCounts
 from openai.types.batch import Metadata as OpenAIBatchMetadata

 from litellm.types.utils import LiteLLMBatch

+# AWS Bedrock model-invocation-job statuses → OpenAI Batch statuses.
+# Mirrors the mapping used by `BedrockBatchesConfig.transform_create_batch_response`
+# so create / retrieve return consistent statuses.
+_BEDROCK_MIJ_STATUS_TO_OPENAI = {
+    "Submitted": "validating",
+    "Validating": "validating",
+    "Scheduled": "validating",
+    "InProgress": "in_progress",
+    "Stopping": "cancelling",
+    "Stopped": "cancelled",
+    "Completed": "completed",
+    "PartiallyCompleted": "completed",
+    "Failed": "failed",
+    "Expired": "expired",
+}
+
+
+def _extract_region_from_bedrock_arn(arn: str) -> Optional[str]:
+    """ARN shape: ``arn:aws:bedrock:<region>:<account>:<type>/<id>``"""
+    try:
+        parts = arn.split(":")
+        if len(parts) >= 4 and parts[2] == "bedrock":
+            return parts[3] or None
+    except Exception:
+        pass
+    return None
+
+
+def _extract_job_id_from_arn(arn: str) -> Optional[str]:
+    """``arn:aws:bedrock:<region>:<acct>:model-invocation-job/<job-id>`` -> ``<job-id>``."""
+    if ":model-invocation-job/" not in arn:
+        return None
+    return arn.rsplit("/", 1)[-1] or None
+
+
+def _predict_output_file_uri(
+    output_prefix: str, input_uri: str, job_id: Optional[str]
+) -> Optional[str]:
+    """
+    Compute the deterministic per-job result file URI Bedrock writes to.
+
+    Bedrock lays results out as::
+
+        <output_prefix>/<job-id>/<basename(input_uri)>.out
+
+    We compute it client-side so OpenAI-style ``client.files.content(output_file_id)``
+    works without an extra S3 ``ListObjectsV2`` round-trip. Returns ``None`` if we
+    don't have enough info; callers should fall back to the bare prefix.
+    """
+    if not output_prefix or not input_uri or not job_id:
+        return None
+    if not output_prefix.endswith("/"):
+        output_prefix = output_prefix + "/"
+    input_basename = input_uri.rsplit("/", 1)[-1]
+    if not input_basename:
+        return None
+    return f"{output_prefix}{job_id}/{input_basename}.out"
+
+
+def _to_epoch(value: Any) -> Optional[int]:
+    if value is None:
+        return None
+    if isinstance(value, (int, float)):
+        return int(value)
+    if isinstance(value, datetime):
+        return int(value.timestamp())
+    return None
+

 class BedrockBatchesHandler:
    """
@ -97,3 +168,173 @@ class BedrockBatchesHandler:
        with concurrent.futures.ThreadPoolExecutor() as executor:
            future = executor.submit(run_in_thread)
            return future.result()
+
+    @staticmethod
+    def _handle_model_invocation_job_status(
+        batch_id: str,
+        aws_region_name: Optional[str] = None,
+        logging_obj=None,
+        **kwargs,
+    ) -> "LiteLLMBatch":
+        """
+        Handle ``GetModelInvocationJob`` status check for AWS Bedrock bulk batch
+        inference jobs (the ARN type returned by ``CreateModelInvocationJob``).
+
+        ``CreateModelInvocationJob`` lives on the Bedrock **control plane**
+        (``bedrock.<region>.amazonaws.com``), distinct from the data-plane
+        ``bedrock-runtime`` endpoint that serves Twelve Labs async-invoke ARNs.
+        The two ARN families therefore can't share a handler — see
+        ``litellm/batches/main.py`` for the dispatch.
+
+        Args:
+            batch_id: A ``arn:aws:bedrock:<region>:<acct>:model-invocation-job/<id>``
+                ARN (or just the trailing job id; both are accepted by
+                ``GetModelInvocationJob``).
+            aws_region_name: Region for the boto3 ``bedrock`` client. If omitted,
+                we fall back to parsing the region out of ``batch_id`` itself.
+            logging_obj: Optional litellm logging object.
+            **kwargs: Optional AWS credential overrides
+                (``aws_access_key_id``, ``aws_secret_access_key``,
+                ``aws_session_token``, ``aws_profile_name``,
+                ``aws_role_name``, ``aws_session_name``,
+                ``aws_web_identity_token``, ``aws_sts_endpoint``,
+                ``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.
+        """
+        try:
+            import boto3
+        except ImportError as exc:
+            raise ImportError(
+                "Missing boto3 to call bedrock. Run 'pip install boto3'."
+            ) from exc
+
+        # Resolve region: explicit > parsed-from-ARN > us-east-1 (boto3 default).
+        region = (
+            aws_region_name or _extract_region_from_bedrock_arn(batch_id) or "us-east-1"
+        )
+
+        # Resolve credentials through the same path the rest of the bedrock
+        # provider uses, so model_list / env / role-assumption configs are
+        # honored. We instantiate BedrockBatchesConfig (which extends
+        # BaseAWSLLM) lazily to avoid a circular import at module load.
+        from litellm.llms.bedrock.batches.transformation import BedrockBatchesConfig
+
+        creds = BedrockBatchesConfig().get_credentials(
+            aws_access_key_id=kwargs.get("aws_access_key_id"),
+            aws_secret_access_key=kwargs.get("aws_secret_access_key"),
+            aws_session_token=kwargs.get("aws_session_token"),
+            aws_region_name=region,
+            aws_session_name=kwargs.get("aws_session_name"),
+            aws_profile_name=kwargs.get("aws_profile_name"),
+            aws_role_name=kwargs.get("aws_role_name"),
+            aws_web_identity_token=kwargs.get("aws_web_identity_token"),
+            aws_sts_endpoint=kwargs.get("aws_sts_endpoint"),
+            aws_external_id=kwargs.get("aws_external_id"),
+        )
+
+        client = boto3.client(
+            "bedrock",
+            region_name=region,
+            aws_access_key_id=creds.access_key,
+            aws_secret_access_key=creds.secret_key,
+            aws_session_token=creds.token,
+        )
+
+        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="",
+                additional_args={
+                    "complete_input_dict": {"jobIdentifier": batch_id},
+                    "api_base": (
+                        f"https://bedrock.{region}.amazonaws.com/"
+                        f"model-invocation-job/{url_path_id}"
+                    ),
+                },
+            )
+
+        response = client.get_model_invocation_job(jobIdentifier=batch_id)
+
+        if logging_obj is not None:
+            logging_obj.post_call(
+                input=batch_id,
+                api_key="",
+                original_response=response,
+                additional_args={"complete_input_dict": {"jobIdentifier": batch_id}},
+            )
+
+        bedrock_status = str(response.get("status", ""))
+        openai_status = cast(
+            Any,
+            _BEDROCK_MIJ_STATUS_TO_OPENAI.get(bedrock_status, "in_progress"),
+        )
+
+        input_uri = (
+            response.get("inputDataConfig", {})
+            .get("s3InputDataConfig", {})
+            .get("s3Uri", "")
+        )
+        output_prefix = (
+            response.get("outputDataConfig", {})
+            .get("s3OutputDataConfig", {})
+            .get("s3Uri", "")
+        )
+
+        # 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. 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)
+
+        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,
+            "job_name": response.get("jobName", ""),
+            "failure_message": response.get("message") or "",
+            "input_s3_uri": input_uri,
+            "output_s3_uri": output_prefix,
+            "output_file_uri": output_file_uri or "",
+        }
+
+        return LiteLLMBatch(
+            id=job_arn,
+            object="batch",
+            status=openai_status,
+            created_at=_to_epoch(response.get("submitTime")) or 0,
+            in_progress_at=_to_epoch(response.get("lastModifiedTime")),
+            completed_at=completed_at if openai_status == "completed" else None,
+            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),
+            metadata=openai_batch_metadata,
+            completion_window="24h",
+            endpoint="/v1/chat/completions",
+            input_file_id=input_uri,
+            output_file_id=output_file_uri if openai_status == "completed" else None,
+        )
--- a/tests/test_litellm/llms/bedrock/batches/test_handler.py
+++ b/tests/test_litellm/llms/bedrock/batches/test_handler.py
@ -0,0 +1,338 @@
+"""Unit tests for ``BedrockBatchesHandler._handle_model_invocation_job_status``.
+
+These cover the upstream support for retrieving Bedrock bulk batch jobs
+(``arn:aws:bedrock:<region>:<acct>:model-invocation-job/<id>``) — the ARN
+type returned by ``CreateModelInvocationJob``. We mock the boto3 client so
+the tests don't hit AWS.
+"""
+
+from __future__ import annotations
+
+import os
+import sys
+from datetime import datetime, timezone
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+sys.path.insert(0, os.path.abspath("../../../../.."))
+
+from litellm.llms.bedrock.batches.handler import (  # noqa: E402
+    BedrockBatchesHandler,
+    _extract_job_id_from_arn,
+    _extract_region_from_bedrock_arn,
+    _predict_output_file_uri,
+    _to_epoch,
+)
+
+JOB_ID = "abc1234567"
+JOB_ARN = f"arn:aws:bedrock:us-west-2:123456789012:model-invocation-job/{JOB_ID}"
+INPUT_URI = "s3://my-bucket/inputs/qwen3-235b-a22b-2507-batch.jsonl"
+OUTPUT_PREFIX = "s3://my-bucket/litellm-batch-outputs/litellm-bedrock-files-qwen-uuid/"
+SUBMIT_TIME = datetime(2026, 4, 28, 12, 0, 0, tzinfo=timezone.utc)
+END_TIME = datetime(2026, 4, 28, 12, 30, 0, tzinfo=timezone.utc)
+
+
+def _fake_boto3_response(status: str = "Completed", end_time=END_TIME):
+    return {
+        "jobArn": JOB_ARN,
+        "jobName": "litellm-bedrock-files-qwen-uuid",
+        "modelId": "bedrock/qwen.qwen3-235b-a22b-2507-v1:0",
+        "status": status,
+        "submitTime": SUBMIT_TIME,
+        "lastModifiedTime": end_time,
+        "endTime": end_time,
+        "inputDataConfig": {"s3InputDataConfig": {"s3Uri": INPUT_URI}},
+        "outputDataConfig": {"s3OutputDataConfig": {"s3Uri": OUTPUT_PREFIX}},
+    }
+
+
+@pytest.fixture
+def patched_boto3():
+    """Yield a stub bedrock client whose `get_model_invocation_job` is a MagicMock."""
+    fake_client = MagicMock()
+    fake_client.get_model_invocation_job.return_value = _fake_boto3_response()
+    with (
+        patch("boto3.client", return_value=fake_client) as boto_client_factory,
+        patch(
+            "litellm.llms.bedrock.batches.transformation.BedrockBatchesConfig.get_credentials",
+            return_value=MagicMock(access_key="AKIA", secret_key="SECRET", token=None),
+        ),
+    ):
+        yield fake_client, boto_client_factory
+
+
+def test_extract_region_from_arn():
+    assert _extract_region_from_bedrock_arn(JOB_ARN) == "us-west-2"
+    assert _extract_region_from_bedrock_arn("arn:aws:bedrock::123:foo/bar") is None
+    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 (
+        _extract_job_id_from_arn("arn:aws:bedrock:us-west-2:1:async-invoke/x") is None
+    )
+
+
+def test_predict_output_file_uri_happy_path():
+    expected = f"{OUTPUT_PREFIX}{JOB_ID}/qwen3-235b-a22b-2507-batch.jsonl.out"
+    assert _predict_output_file_uri(OUTPUT_PREFIX, INPUT_URI, JOB_ID) == expected
+
+
+def test_predict_output_file_uri_adds_trailing_slash():
+    prefix_no_slash = OUTPUT_PREFIX.rstrip("/")
+    expected = f"{OUTPUT_PREFIX}{JOB_ID}/qwen3-235b-a22b-2507-batch.jsonl.out"
+    assert _predict_output_file_uri(prefix_no_slash, INPUT_URI, JOB_ID) == expected
+
+
+@pytest.mark.parametrize(
+    "missing_arg",
+    [
+        ("", INPUT_URI, JOB_ID),
+        (OUTPUT_PREFIX, "", JOB_ID),
+        (OUTPUT_PREFIX, INPUT_URI, None),
+    ],
+)
+def test_predict_output_file_uri_returns_none_when_missing_input(missing_arg):
+    assert _predict_output_file_uri(*missing_arg) is None
+
+
+def test_handle_model_invocation_job_status_completed(patched_boto3):
+    fake_client, boto_client_factory = patched_boto3
+
+    batch = BedrockBatchesHandler._handle_model_invocation_job_status(batch_id=JOB_ARN)
+
+    fake_client.get_model_invocation_job.assert_called_once_with(jobIdentifier=JOB_ARN)
+
+    # Region should be sniffed from the ARN.
+    _, kwargs = boto_client_factory.call_args
+    assert kwargs["region_name"] == "us-west-2"
+
+    assert batch.id == JOB_ARN
+    assert batch.status == "completed"
+    assert batch.input_file_id == INPUT_URI
+    expected_out = f"{OUTPUT_PREFIX}{JOB_ID}/qwen3-235b-a22b-2507-batch.jsonl.out"
+    assert batch.output_file_id == expected_out
+    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.metadata["job_arn"] == JOB_ARN
+    assert batch.metadata["output_file_uri"] == expected_out
+    assert batch.metadata["output_s3_uri"] == OUTPUT_PREFIX
+
+
+@pytest.mark.parametrize(
+    "bedrock_status,openai_status",
+    [
+        ("Submitted", "validating"),
+        ("Validating", "validating"),
+        ("Scheduled", "validating"),
+        ("InProgress", "in_progress"),
+        ("Stopping", "cancelling"),
+        ("Stopped", "cancelled"),
+        ("Completed", "completed"),
+        ("PartiallyCompleted", "completed"),
+        ("Failed", "failed"),
+        ("Expired", "expired"),
+        # Unknown/unmapped Bedrock status falls back to "in_progress" so we
+        # don't 500 on a future AWS-side enum addition.
+        ("MyBrandNewStatus", "in_progress"),
+    ],
+)
+def test_status_mapping(patched_boto3, bedrock_status, openai_status):
+    fake_client, _ = patched_boto3
+    fake_client.get_model_invocation_job.return_value = _fake_boto3_response(
+        status=bedrock_status
+    )
+
+    batch = BedrockBatchesHandler._handle_model_invocation_job_status(batch_id=JOB_ARN)
+
+    assert batch.status == openai_status
+    # output_file_id is only populated for terminal-completed jobs, so callers
+    # don't accidentally try to download a non-existent file mid-run.
+    if openai_status == "completed":
+        assert batch.output_file_id is not None
+    else:
+        assert batch.output_file_id is None
+
+
+def test_explicit_region_overrides_arn(patched_boto3):
+    _, boto_client_factory = patched_boto3
+    BedrockBatchesHandler._handle_model_invocation_job_status(
+        batch_id=JOB_ARN, aws_region_name="eu-central-1"
+    )
+    _, kwargs = boto_client_factory.call_args
+    assert kwargs["region_name"] == "eu-central-1"
+
+
+def test_failure_message_propagates(patched_boto3):
+    fake_client, _ = patched_boto3
+    failed_response = _fake_boto3_response(status="Failed")
+    failed_response["message"] = "Input file failed validation"
+    fake_client.get_model_invocation_job.return_value = failed_response
+
+    batch = BedrockBatchesHandler._handle_model_invocation_job_status(batch_id=JOB_ARN)
+
+    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}"
+    )
--- a/tests/test_litellm/test_retrieve_batch_bedrock_dispatch.py
+++ b/tests/test_litellm/test_retrieve_batch_bedrock_dispatch.py
@ -0,0 +1,162 @@
+"""Cover the Bedrock-ARN dispatch in ``litellm.batches.main.retrieve_batch``.
+
+The dispatch picks one of two Bedrock handlers depending on the ARN
+family in ``batch_id``:
+
+* ``:async-invoke/<id>``        -> ``_handle_async_invoke_status`` (data plane)
+* ``:model-invocation-job/<id>`` -> ``_handle_model_invocation_job_status``
+                                    (control plane, added in this PR)
+
+Anything else falls through to the generic ``provider_config`` retrieve
+flow. We mock the two handlers so the tests don't hit AWS — the focus
+here is purely the dispatch logic that lives in ``main.py``.
+"""
+
+from __future__ import annotations
+
+import os
+import sys
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+sys.path.insert(0, os.path.abspath("../.."))
+
+import litellm  # noqa: E402
+
+ASYNC_INVOKE_ARN = "arn:aws:bedrock:us-west-2:123456789012:async-invoke/abc123def456"
+MIJ_ARN = "arn:aws:bedrock:us-west-2:123456789012:model-invocation-job/abc1234567"
+
+
+@pytest.fixture
+def mock_handlers():
+    """Patch both Bedrock retrieve handlers and yield the mocks.
+
+    We patch at the import site (litellm.batches.main) rather than the
+    definition site so the ``BedrockBatchesHandler`` reference inside
+    ``retrieve_batch`` resolves to our mocks.
+    """
+    fake_batch = MagicMock(name="LiteLLMBatch")
+    with (
+        patch(
+            "litellm.batches.main.BedrockBatchesHandler._handle_async_invoke_status",
+            return_value=fake_batch,
+        ) as async_invoke,
+        patch(
+            "litellm.batches.main.BedrockBatchesHandler._handle_model_invocation_job_status",
+            return_value=fake_batch,
+        ) as mij,
+    ):
+        yield async_invoke, mij, fake_batch
+
+
+def test_async_invoke_arn_routes_to_async_invoke_handler(mock_handlers):
+    """``:async-invoke/`` ARNs go to the data-plane handler."""
+    async_invoke, mij, fake_batch = mock_handlers
+
+    result = litellm.retrieve_batch(
+        batch_id=ASYNC_INVOKE_ARN,
+        custom_llm_provider="bedrock",
+        aws_region_name="us-west-2",
+    )
+
+    assert result is fake_batch
+    async_invoke.assert_called_once()
+    mij.assert_not_called()
+    call_kwargs = async_invoke.call_args.kwargs
+    assert call_kwargs["batch_id"] == ASYNC_INVOKE_ARN
+    assert call_kwargs["aws_region_name"] == "us-west-2"
+    # Region must be stripped from the forwarded kwargs to avoid TypeError
+    # (it's already an explicit positional/keyword arg).
+    assert "aws_region_name" not in {
+        k
+        for k in call_kwargs
+        if k not in {"batch_id", "aws_region_name", "logging_obj"}
+    }
+
+
+def test_async_invoke_arn_falls_back_to_default_region_when_unset(mock_handlers):
+    """If no ``aws_region_name`` is passed, the data-plane handler defaults
+    to ``us-east-1`` (preserving prior behavior on this branch)."""
+    async_invoke, _mij, _ = mock_handlers
+
+    litellm.retrieve_batch(
+        batch_id=ASYNC_INVOKE_ARN,
+        custom_llm_provider="bedrock",
+    )
+
+    async_invoke.assert_called_once()
+    assert async_invoke.call_args.kwargs["aws_region_name"] == "us-east-1"
+
+
+def test_model_invocation_job_arn_routes_to_mij_handler(mock_handlers):
+    """``:model-invocation-job/`` ARNs go to the new control-plane handler."""
+    _async_invoke, mij, fake_batch = mock_handlers
+
+    result = litellm.retrieve_batch(
+        batch_id=MIJ_ARN,
+        custom_llm_provider="bedrock",
+        aws_region_name="us-west-2",
+    )
+
+    assert result is fake_batch
+    mij.assert_called_once()
+    _async_invoke.assert_not_called()
+    call_kwargs = mij.call_args.kwargs
+    assert call_kwargs["batch_id"] == MIJ_ARN
+    assert call_kwargs["aws_region_name"] == "us-west-2"
+
+
+def test_model_invocation_job_arn_with_no_region_passes_none(mock_handlers):
+    """The MIJ handler is responsible for sniffing region from the ARN
+    when none is explicitly provided. Dispatch must forward ``None``
+    rather than substituting a default — otherwise per-region jobs in
+    other AWS regions would silently route to ``us-east-1``."""
+    _async_invoke, mij, _ = mock_handlers
+
+    litellm.retrieve_batch(
+        batch_id=MIJ_ARN,
+        custom_llm_provider="bedrock",
+    )
+
+    mij.assert_called_once()
+    assert mij.call_args.kwargs["aws_region_name"] is None
+
+
+def test_unrelated_bedrock_arn_falls_through_to_provider_config(mock_handlers):
+    """Bedrock ARNs that aren't async-invoke or model-invocation-job
+    must NOT hit either special handler — they should fall through to
+    the existing generic provider_config path. We don't fully exercise
+    that path here (it requires a real provider config); we just assert
+    neither special handler is invoked."""
+    async_invoke, mij, _ = mock_handlers
+
+    # Use a plausible-but-unsupported Bedrock ARN family.
+    unrelated_arn = "arn:aws:bedrock:us-west-2:123456789012:provisioned-model/xyz"
+
+    with pytest.raises(Exception):
+        # Will raise because no provider_config exists for this path —
+        # that's fine, we just need to assert neither bedrock handler ran
+        # before the failure.
+        litellm.retrieve_batch(
+            batch_id=unrelated_arn,
+            custom_llm_provider="bedrock",
+        )
+
+    async_invoke.assert_not_called()
+    mij.assert_not_called()
+
+
+def test_non_bedrock_id_skips_bedrock_dispatch_entirely(mock_handlers):
+    """Plain (non-ARN) batch ids must not even enter the Bedrock dispatch
+    block — they belong to other providers' retrieve flows."""
+    async_invoke, mij, _ = mock_handlers
+
+    with pytest.raises(Exception):
+        litellm.retrieve_batch(
+            batch_id="batch_abc123",
+            custom_llm_provider="openai",
+        )
+
+    async_invoke.assert_not_called()
+    mij.assert_not_called()