From 37da5b6f4d8ebf6e33379f09e36fcccc9bf11c4f Mon Sep 17 00:00:00 2001 From: kerry Date: Sat, 19 Sep 2026 00:10:17 +0000 Subject: [PATCH 1/3] fix(bedrock): sign batch retrieve and cancel with deployment credentials when AWS_BEARER_TOKEN_BEDROCK is set Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/llms/bedrock/batches/handler.py | 10 ++++ .../llms/bedrock/batches/test_handler.py | 50 +++++++++++++++++++ 2 files changed, 60 insertions(+) diff --git a/litellm/llms/bedrock/batches/handler.py b/litellm/llms/bedrock/batches/handler.py index 6239973eb7c..fd4c3dc1659 100644 --- a/litellm/llms/bedrock/batches/handler.py +++ b/litellm/llms/bedrock/batches/handler.py @@ -10,6 +10,8 @@ from litellm.types.llms.bedrock import AwsAuthParams, AwsSessionTag from litellm.types.utils import LiteLLMBatch if TYPE_CHECKING: + from botocore.config import Config + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj # AWS Bedrock model-invocation-job statuses → OpenAI Batch statuses. @@ -31,6 +33,12 @@ _BEDROCK_MIJ_STATUS_TO_OPENAI: Final = { _CANCEL_IDEMPOTENT_STATUSES: Final = frozenset({"cancelling", "cancelled", "completed", "failed", "expired"}) +def _sigv4_config() -> "Config": + from botocore.config import Config + + return Config(signature_version="v4") + + def _extract_region_from_bedrock_arn(arn: str) -> str | None: """ARN shape: ``arn:aws:bedrock:::/``""" try: @@ -150,6 +158,7 @@ class BedrockBatchesHandler: aws_access_key_id=creds.access_key, aws_secret_access_key=creds.secret_key, aws_session_token=creds.token, + config=_sigv4_config(), ) def job_status() -> "LiteLLMBatch": @@ -309,6 +318,7 @@ class BedrockBatchesHandler: aws_access_key_id=creds.access_key, aws_secret_access_key=creds.secret_key, aws_session_token=creds.token, + config=_sigv4_config(), ) if logging_obj is not None: diff --git a/tests/test_litellm/llms/bedrock/batches/test_handler.py b/tests/test_litellm/llms/bedrock/batches/test_handler.py index 03daafcad72..056378f97c9 100644 --- a/tests/test_litellm/llms/bedrock/batches/test_handler.py +++ b/tests/test_litellm/llms/bedrock/batches/test_handler.py @@ -570,3 +570,53 @@ def test_cancel_batch_stops_and_polls_the_job_with_the_tagged_session(monkeypatc fake_bedrock.stop_model_invocation_job.assert_called_once_with(jobIdentifier=JOB_ARN) assert batch.status == "cancelled" assert [kwargs["aws_access_key_id"] for kwargs in bedrock_client_kwargs] == ["ASIABATCHCANCELTAGGED"] * 2 + + +def _sigv4_capture_send(sent_headers: list[dict[str, str]], body: dict): + import json + + from botocore.awsrequest import AWSResponse + + def send(_self, request): + sent_headers.append({k: v.decode() if isinstance(v, bytes) else v for k, v in request.headers.items()}) + raw = MagicMock() + raw.stream.return_value = iter([json.dumps(body, default=str).encode()]) + return AWSResponse(request.url, 200, {"content-type": "application/json"}, raw) + + return send + + +def test_retrieve_signs_with_deployment_credentials_when_env_bearer_token_is_set(monkeypatch): + """A proxy-wide AWS_BEARER_TOKEN_BEDROCK must not override the deployment's own SigV4 credentials.""" + monkeypatch.setenv("AWS_BEARER_TOKEN_BEDROCK", "env-bearer-token") + sent_headers: list[dict[str, str]] = [] + + with patch("botocore.httpsession.URLLib3Session.send", _sigv4_capture_send(sent_headers, _fake_boto3_response())): + batch = BedrockBatchesHandler._handle_model_invocation_job_status( + batch_id=JOB_ARN, + aws_access_key_id="AKIADEPLOYMENTKEY", + aws_secret_access_key="deployment-secret", + ) + + assert batch.status == "completed" + assert len(sent_headers) == 1 + assert sent_headers[0]["Authorization"].startswith("AWS4-HMAC-SHA256 Credential=AKIADEPLOYMENTKEY/") + + +def test_cancel_signs_with_deployment_credentials_when_env_bearer_token_is_set(monkeypatch): + monkeypatch.setenv("AWS_BEARER_TOKEN_BEDROCK", "env-bearer-token") + sent_headers: list[dict[str, str]] = [] + + with patch( + "botocore.httpsession.URLLib3Session.send", + _sigv4_capture_send(sent_headers, _fake_boto3_response(status="Stopped")), + ): + batch = BedrockBatchesHandler.cancel_batch( + batch_id=JOB_ARN, + aws_access_key_id="AKIADEPLOYMENTKEY", + aws_secret_access_key="deployment-secret", + ) + + assert batch.status == "cancelled" + assert len(sent_headers) == 2 + assert all(h["Authorization"].startswith("AWS4-HMAC-SHA256 Credential=AKIADEPLOYMENTKEY/") for h in sent_headers) From 6b082d3a018b9956425babb47ce8cb2aee260a6f Mon Sep 17 00:00:00 2001 From: kerry Date: Sat, 19 Sep 2026 00:25:01 +0000 Subject: [PATCH 2/3] test(bedrock): type the SigV4 request recorder and drop caller-owned mutation Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../llms/bedrock/batches/test_handler.py | 51 +++++++++++-------- 1 file changed, 31 insertions(+), 20 deletions(-) diff --git a/tests/test_litellm/llms/bedrock/batches/test_handler.py b/tests/test_litellm/llms/bedrock/batches/test_handler.py index 056378f97c9..d328e09056b 100644 --- a/tests/test_litellm/llms/bedrock/batches/test_handler.py +++ b/tests/test_litellm/llms/bedrock/batches/test_handler.py @@ -8,10 +8,14 @@ the tests don't hit AWS. from __future__ import annotations +import json +from collections.abc import Iterator, Mapping from datetime import datetime, timedelta, timezone +from typing import Final from unittest.mock import MagicMock, patch import pytest +from botocore.awsrequest import AWSPreparedRequest, AWSResponse from litellm.llms.bedrock.batches.handler import ( # noqa: E402 @@ -572,26 +576,36 @@ def test_cancel_batch_stops_and_polls_the_job_with_the_tagged_session(monkeypatc assert [kwargs["aws_access_key_id"] for kwargs in bedrock_client_kwargs] == ["ASIABATCHCANCELTAGGED"] * 2 -def _sigv4_capture_send(sent_headers: list[dict[str, str]], body: dict): - import json +class _JsonBody: + def __init__(self, payload: bytes) -> None: + self._payload: Final = payload - from botocore.awsrequest import AWSResponse + def stream(self) -> Iterator[bytes]: + return iter((self._payload,)) - def send(_self, request): - sent_headers.append({k: v.decode() if isinstance(v, bytes) else v for k, v in request.headers.items()}) - raw = MagicMock() - raw.stream.return_value = iter([json.dumps(body, default=str).encode()]) - return AWSResponse(request.url, 200, {"content-type": "application/json"}, raw) - return send +class _AuthorizationRecorder: + """Stands in for botocore's HTTP session and records the Authorization header of every request it receives.""" + + def __init__(self, body: Mapping[str, object]) -> None: + self._payload: Final = json.dumps(body, default=str).encode() + self.authorization_headers: tuple[str, ...] = () + + def send(self, request: AWSPreparedRequest) -> AWSResponse: + raw_authorization: Final = request.headers["Authorization"] + authorization: Final = ( + raw_authorization.decode() if isinstance(raw_authorization, bytes) else str(raw_authorization) + ) + self.authorization_headers = (*self.authorization_headers, authorization) + return AWSResponse(request.url, 200, {"content-type": "application/json"}, _JsonBody(self._payload)) def test_retrieve_signs_with_deployment_credentials_when_env_bearer_token_is_set(monkeypatch): """A proxy-wide AWS_BEARER_TOKEN_BEDROCK must not override the deployment's own SigV4 credentials.""" monkeypatch.setenv("AWS_BEARER_TOKEN_BEDROCK", "env-bearer-token") - sent_headers: list[dict[str, str]] = [] + recorder: Final = _AuthorizationRecorder(_fake_boto3_response()) - with patch("botocore.httpsession.URLLib3Session.send", _sigv4_capture_send(sent_headers, _fake_boto3_response())): + with patch("botocore.httpsession.URLLib3Session.send", recorder.send): batch = BedrockBatchesHandler._handle_model_invocation_job_status( batch_id=JOB_ARN, aws_access_key_id="AKIADEPLOYMENTKEY", @@ -599,18 +613,15 @@ def test_retrieve_signs_with_deployment_credentials_when_env_bearer_token_is_set ) assert batch.status == "completed" - assert len(sent_headers) == 1 - assert sent_headers[0]["Authorization"].startswith("AWS4-HMAC-SHA256 Credential=AKIADEPLOYMENTKEY/") + assert len(recorder.authorization_headers) == 1 + assert recorder.authorization_headers[0].startswith("AWS4-HMAC-SHA256 Credential=AKIADEPLOYMENTKEY/") def test_cancel_signs_with_deployment_credentials_when_env_bearer_token_is_set(monkeypatch): monkeypatch.setenv("AWS_BEARER_TOKEN_BEDROCK", "env-bearer-token") - sent_headers: list[dict[str, str]] = [] + recorder: Final = _AuthorizationRecorder(_fake_boto3_response(status="Stopped")) - with patch( - "botocore.httpsession.URLLib3Session.send", - _sigv4_capture_send(sent_headers, _fake_boto3_response(status="Stopped")), - ): + with patch("botocore.httpsession.URLLib3Session.send", recorder.send): batch = BedrockBatchesHandler.cancel_batch( batch_id=JOB_ARN, aws_access_key_id="AKIADEPLOYMENTKEY", @@ -618,5 +629,5 @@ def test_cancel_signs_with_deployment_credentials_when_env_bearer_token_is_set(m ) assert batch.status == "cancelled" - assert len(sent_headers) == 2 - assert all(h["Authorization"].startswith("AWS4-HMAC-SHA256 Credential=AKIADEPLOYMENTKEY/") for h in sent_headers) + assert len(recorder.authorization_headers) == 2 + assert all(h.startswith("AWS4-HMAC-SHA256 Credential=AKIADEPLOYMENTKEY/") for h in recorder.authorization_headers) From 99659e9e7e50b01c86e88b2e570179ff0f943acb Mon Sep 17 00:00:00 2001 From: kerry Date: Sat, 19 Sep 2026 00:43:51 +0000 Subject: [PATCH 3/3] test(bedrock): drop redundant recorder docstring Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/test_litellm/llms/bedrock/batches/test_handler.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/tests/test_litellm/llms/bedrock/batches/test_handler.py b/tests/test_litellm/llms/bedrock/batches/test_handler.py index d328e09056b..e69098a460d 100644 --- a/tests/test_litellm/llms/bedrock/batches/test_handler.py +++ b/tests/test_litellm/llms/bedrock/batches/test_handler.py @@ -585,8 +585,6 @@ class _JsonBody: class _AuthorizationRecorder: - """Stands in for botocore's HTTP session and records the Authorization header of every request it receives.""" - def __init__(self, body: Mapping[str, object]) -> None: self._payload: Final = json.dumps(body, default=str).encode() self.authorization_headers: tuple[str, ...] = ()