From 2aa005fed2911f874edda62af4ce9ec1740eebf4 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 2 Sep 2026 14:28:40 -0700 Subject: [PATCH 1/3] fix(bedrock): skip the SigV4 credential chain when a bearer token is configured A deployment authenticating with api_key or AWS_BEARER_TOKEN_BEDROCK still ran boto3's credential chain before every call, so an unloadable default profile (a login_session profile without botocore[crt]) made Converse, embeddings, image generation, image edit, and the Bedrock guardrail hook fail with MissingDependencyException even though the bearer token alone signs the request. The chain now runs only when no bearer token is configured --- litellm/llms/bedrock/base_aws_llm.py | 52 +++++++++--------- litellm/llms/bedrock/chat/converse_handler.py | 28 +++++----- litellm/llms/bedrock/embed/embedding.py | 48 +++++++++-------- litellm/llms/bedrock/image_edit/handler.py | 6 ++- .../bedrock/image_generation/image_handler.py | 6 ++- .../guardrail_hooks/bedrock_guardrails.py | 53 +++++++++---------- .../secret_managers/aws_secret_manager_v2.py | 7 ++- .../chat/test_bedrock_converse_handler.py | 24 +++++++++ .../bedrock/embed/test_bedrock_embedding.py | 26 +++++++++ .../image/test_bedrock_image_bearer_token.py | 21 ++++++++ .../test_amazon_nova_canvas_image_edit.py | 21 ++++++++ .../test_bedrock_guardrails.py | 22 ++++++++ .../test_bedrock_invoke_guardrail_checks.py | 28 ++++++++++ 13 files changed, 250 insertions(+), 92 deletions(-) diff --git a/litellm/llms/bedrock/base_aws_llm.py b/litellm/llms/bedrock/base_aws_llm.py index 1e634ced29b..1f00bf7792e 100644 --- a/litellm/llms/bedrock/base_aws_llm.py +++ b/litellm/llms/bedrock/base_aws_llm.py @@ -49,11 +49,16 @@ SIGV4_COMPUTED_HEADERS: Final = frozenset({"authorization", "x-amz-date", "x-amz class Boto3CredentialsInfo(BaseModel): - credentials: Credentials + credentials: Credentials | None aws_region_name: str aws_bedrock_runtime_endpoint: str | None +def bedrock_bearer_token(api_key: str | None) -> str | None: + token: Final = api_key if api_key is not None else get_secret_str("AWS_BEARER_TOKEN_BEDROCK") + return token or None + + class _WebIdentityTokenClaims(BaseModel): aud: str | list[str] | None = None iss: str | None = None @@ -1388,7 +1393,7 @@ class BaseAWSLLM: return f"https://bedrock-runtime.{aws_region_name}.{dns_suffix}" def _get_boto_credentials_from_optional_params( - self, optional_params: dict, model: str | None = None + self, optional_params: dict, model: str | None = None, bearer_token: str | None = None ) -> Boto3CredentialsInfo: """ Get boto3 credentials from optional params @@ -1420,17 +1425,21 @@ class BaseAWSLLM: ) # https://bedrock-runtime.{region_name}.amazonaws.com aws_external_id: Final = optional_params.pop("aws_external_id", None) - credentials: Final[Credentials] = self.get_credentials( - aws_access_key_id=aws_access_key_id, - aws_secret_access_key=aws_secret_access_key, - aws_session_token=aws_session_token, - aws_region_name=aws_region_name, - aws_session_name=aws_session_name, - aws_profile_name=aws_profile_name, - aws_role_name=aws_role_name, - aws_web_identity_token=aws_web_identity_token, - aws_sts_endpoint=aws_sts_endpoint, - aws_external_id=aws_external_id, + credentials: Final[Credentials | None] = ( + None + if bearer_token is not None + else self.get_credentials( + aws_access_key_id=aws_access_key_id, + aws_secret_access_key=aws_secret_access_key, + aws_session_token=aws_session_token, + aws_region_name=aws_region_name, + aws_session_name=aws_session_name, + aws_profile_name=aws_profile_name, + aws_role_name=aws_role_name, + aws_web_identity_token=aws_web_identity_token, + aws_sts_endpoint=aws_sts_endpoint, + aws_external_id=aws_external_id, + ) ) return Boto3CredentialsInfo( @@ -1451,14 +1460,9 @@ class BaseAWSLLM: api_key: str | None = None, supports_bearer_token: bool = True, ) -> AWSPreparedRequest: - if not supports_bearer_token: - aws_bearer_token: str | None = None - elif api_key is not None: - aws_bearer_token = api_key - else: - aws_bearer_token = get_secret_str("AWS_BEARER_TOKEN_BEDROCK") + aws_bearer_token: Final = bedrock_bearer_token(api_key) if supports_bearer_token else None - if aws_bearer_token: + if aws_bearer_token is not None: try: from botocore.awsrequest import AWSRequest except ImportError: @@ -1555,13 +1559,9 @@ class BaseAWSLLM: Returns: Tuple[dict, Optional[str]]: A tuple containing the headers and the json str body of the request """ - if api_key is not None: - aws_bearer_token: str | None = api_key - else: - aws_bearer_token = get_secret_str("AWS_BEARER_TOKEN_BEDROCK") + aws_bearer_token: Final = bedrock_bearer_token(api_key) - # If aws bearer token is set, use it directly in the header - if aws_bearer_token: + if aws_bearer_token is not None: headers = headers or {} headers["Content-Type"] = "application/json" headers["Authorization"] = f"Bearer {aws_bearer_token}" diff --git a/litellm/llms/bedrock/chat/converse_handler.py b/litellm/llms/bedrock/chat/converse_handler.py index 7d5f99ca893..a75124325ae 100644 --- a/litellm/llms/bedrock/chat/converse_handler.py +++ b/litellm/llms/bedrock/chat/converse_handler.py @@ -21,7 +21,7 @@ from litellm.rust_bridge.chat_completions import rust_chat_completions_accepts from litellm.types.utils import ModelResponse from litellm.utils import CustomStreamWrapper -from ..base_aws_llm import BaseAWSLLM, Credentials +from ..base_aws_llm import BaseAWSLLM, Credentials, bedrock_bearer_token from ..common_utils import BedrockError, _get_all_bedrock_regions from .invoke_handler import AWSEventStreamDecoder, MockResponseIterator, make_call @@ -349,17 +349,21 @@ class BedrockConverseLLM(BaseAWSLLM): litellm_params["aws_region_name"] = aws_region_name # [DO NOT DELETE] important for async calls - credentials: Final[Credentials | None] = self.get_credentials( - aws_access_key_id=aws_access_key_id, - aws_secret_access_key=aws_secret_access_key, - aws_session_token=aws_session_token, - aws_region_name=aws_region_name, - aws_session_name=aws_session_name, - aws_profile_name=aws_profile_name, - aws_role_name=aws_role_name, - aws_web_identity_token=aws_web_identity_token, - aws_sts_endpoint=aws_sts_endpoint, - aws_external_id=aws_external_id, + credentials: Final[Credentials | None] = ( + None + if bedrock_bearer_token(api_key) is not None + else self.get_credentials( + aws_access_key_id=aws_access_key_id, + aws_secret_access_key=aws_secret_access_key, + aws_session_token=aws_session_token, + aws_region_name=aws_region_name, + aws_session_name=aws_session_name, + aws_profile_name=aws_profile_name, + aws_role_name=aws_role_name, + aws_web_identity_token=aws_web_identity_token, + aws_sts_endpoint=aws_sts_endpoint, + aws_external_id=aws_external_id, + ) ) ### SET RUNTIME ENDPOINT ### diff --git a/litellm/llms/bedrock/embed/embedding.py b/litellm/llms/bedrock/embed/embedding.py index c34ca7750e2..a7b74f3752a 100644 --- a/litellm/llms/bedrock/embed/embedding.py +++ b/litellm/llms/bedrock/embed/embedding.py @@ -6,7 +6,7 @@ import copy import json import urllib.parse from collections.abc import Callable -from typing import TYPE_CHECKING, Any, Final, get_args +from typing import TYPE_CHECKING, Final, get_args import httpx @@ -26,7 +26,7 @@ from litellm.types.llms.bedrock import ( ) from litellm.types.utils import EmbeddingResponse, LlmProviders -from ..base_aws_llm import BaseAWSLLM +from ..base_aws_llm import BaseAWSLLM, Credentials, bedrock_bearer_token from ..common_utils import BedrockError from .amazon_nova_transformation import AmazonNovaEmbeddingConfig from .amazon_titan_g1_transformation import AmazonTitanG1Config @@ -45,11 +45,8 @@ class BedrockEmbedding(BaseAWSLLM): def _load_credentials( self, optional_params: dict, - ) -> tuple[Any, str]: - try: - from botocore.credentials import Credentials - except ImportError: - raise ImportError("Missing boto3 to call bedrock. Run 'pip install boto3'.") + bearer_token: str | None = None, + ) -> tuple[Credentials | None, str]: ## CREDENTIALS ## # pop aws_secret_access_key, aws_access_key_id, aws_session_token, aws_region_name from kwargs, since completion calls fail with them aws_secret_access_key: Final = optional_params.pop("aws_secret_access_key", None) @@ -78,17 +75,21 @@ class BedrockEmbedding(BaseAWSLLM): if aws_region_name is None: aws_region_name = "us-west-2" - credentials: Final[Credentials] = self.get_credentials( - aws_access_key_id=aws_access_key_id, - aws_secret_access_key=aws_secret_access_key, - aws_session_token=aws_session_token, - aws_region_name=aws_region_name, - aws_session_name=aws_session_name, - aws_profile_name=aws_profile_name, - aws_role_name=aws_role_name, - aws_web_identity_token=aws_web_identity_token, - aws_sts_endpoint=aws_sts_endpoint, - aws_external_id=aws_external_id, + credentials: Final[Credentials | None] = ( + None + if bearer_token is not None + else self.get_credentials( + aws_access_key_id=aws_access_key_id, + aws_secret_access_key=aws_secret_access_key, + aws_session_token=aws_session_token, + aws_region_name=aws_region_name, + aws_session_name=aws_session_name, + aws_profile_name=aws_profile_name, + aws_role_name=aws_role_name, + aws_web_identity_token=aws_web_identity_token, + aws_sts_endpoint=aws_sts_endpoint, + aws_external_id=aws_external_id, + ) ) return credentials, aws_region_name @@ -233,7 +234,7 @@ class BedrockEmbedding(BaseAWSLLM): client: HTTPHandler | None, timeout: float | httpx.Timeout | None, batch_data: list[dict], - credentials: Any, + credentials: Credentials | None, extra_headers: dict | None, endpoint_url: str, aws_region_name: str, @@ -301,7 +302,7 @@ class BedrockEmbedding(BaseAWSLLM): client: AsyncHTTPHandler | None, timeout: float | httpx.Timeout | None, batch_data: list[dict], - credentials: Any, + credentials: Credentials | None, extra_headers: dict | None, endpoint_url: str, aws_region_name: str, @@ -383,7 +384,9 @@ class BedrockEmbedding(BaseAWSLLM): litellm_params: dict, api_key: str | None = None, ) -> EmbeddingResponse: - credentials, aws_region_name = self._load_credentials(optional_params) + credentials, aws_region_name = self._load_credentials( + optional_params, bearer_token=bedrock_bearer_token(api_key) + ) ### TRANSFORMATION ### unencoded_model_id: Final = optional_params.pop("model_id", None) or model # default to model if not passed @@ -595,8 +598,11 @@ class BedrockEmbedding(BaseAWSLLM): try: from botocore.auth import SigV4Auth from botocore.awsrequest import AWSRequest + from botocore.exceptions import NoCredentialsError except ImportError: raise ImportError("Missing boto3 to call bedrock. Run 'pip install boto3'.") + if credentials is None: + raise NoCredentialsError() # Create AWSRequest with GET method and encoded URL request: Final = AWSRequest( diff --git a/litellm/llms/bedrock/image_edit/handler.py b/litellm/llms/bedrock/image_edit/handler.py index 9d8631c7c26..5c517f2049c 100644 --- a/litellm/llms/bedrock/image_edit/handler.py +++ b/litellm/llms/bedrock/image_edit/handler.py @@ -29,7 +29,7 @@ from litellm.llms.custom_httpx.http_handler import ( ) from litellm.types.utils import ImageResponse -from ..base_aws_llm import BaseAWSLLM +from ..base_aws_llm import BaseAWSLLM, bedrock_bearer_token from ..common_utils import BedrockError if TYPE_CHECKING: @@ -198,7 +198,9 @@ class BedrockImageEdit(BaseAWSLLM): Returns: BedrockImageEditPreparedRequest: The prepared request object """ - boto3_credentials_info: Final = self._get_boto_credentials_from_optional_params(optional_params, model) + boto3_credentials_info: Final = self._get_boto_credentials_from_optional_params( + optional_params, model, bearer_token=bedrock_bearer_token(api_key) + ) # Use the existing ARN-aware provider detection method bedrock_provider: Final = self.get_bedrock_invoke_provider(model) diff --git a/litellm/llms/bedrock/image_generation/image_handler.py b/litellm/llms/bedrock/image_generation/image_handler.py index 6fac14a0dc3..c78e3c147cb 100644 --- a/litellm/llms/bedrock/image_generation/image_handler.py +++ b/litellm/llms/bedrock/image_generation/image_handler.py @@ -29,7 +29,7 @@ from litellm.llms.custom_httpx.http_handler import ( ) from litellm.types.utils import ImageResponse -from ..base_aws_llm import BaseAWSLLM +from ..base_aws_llm import BaseAWSLLM, bedrock_bearer_token from ..common_utils import BedrockError if TYPE_CHECKING: @@ -220,7 +220,9 @@ class BedrockImageGeneration(BaseAWSLLM): prepped (httpx.Request): The prepared request object body (bytes): The request body """ - boto3_credentials_info: Final = self._get_boto_credentials_from_optional_params(optional_params, model) + boto3_credentials_info: Final = self._get_boto_credentials_from_optional_params( + optional_params, model, bearer_token=bedrock_bearer_token(api_key) + ) # Use the existing ARN-aware provider detection method bedrock_provider: Final = self.get_bedrock_invoke_provider(model) diff --git a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py index 30526d30dc5..7f2616c2fb8 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py +++ b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py @@ -41,7 +41,7 @@ from litellm.llms.anthropic.chat.guardrail_translation.handler import AnthropicM from litellm.llms.base_llm.guardrail_translation.utils import ( effective_scan_only_tool_results_for_guardrail, ) -from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM +from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM, bedrock_bearer_token from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, httpxSpecialProvider, @@ -56,7 +56,6 @@ from litellm.proxy.guardrails.anthropic_sse import ( is_raw_sse_stream, model_response_text, ) -from litellm.secret_managers.main import get_secret_str from litellm.types.guardrails import ( BedrockChecksConfigModel, BedrockGuardrailStreamingParams, @@ -713,9 +712,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): # logic becomes shared across providers. #### CALL HOOKS - proxy only #### - def _load_credentials( - self, - ): + def _load_credentials(self, bearer_token: str | None = None): try: from botocore.credentials import Credentials except ImportError: @@ -737,17 +734,21 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): aws_region_name=aws_region_name, ) - credentials: Final[Credentials] = self.get_credentials( - aws_access_key_id=aws_access_key_id, - aws_secret_access_key=aws_secret_access_key, - aws_session_token=aws_session_token, - aws_region_name=aws_region_name, - aws_session_name=aws_session_name, - aws_profile_name=aws_profile_name, - aws_role_name=aws_role_name, - aws_web_identity_token=aws_web_identity_token, - aws_sts_endpoint=aws_sts_endpoint, - aws_external_id=aws_external_id, + credentials: Final[Credentials | None] = ( + None + if bearer_token is not None + else self.get_credentials( + aws_access_key_id=aws_access_key_id, + aws_secret_access_key=aws_secret_access_key, + aws_session_token=aws_session_token, + aws_region_name=aws_region_name, + aws_session_name=aws_session_name, + aws_profile_name=aws_profile_name, + aws_role_name=aws_role_name, + aws_web_identity_token=aws_web_identity_token, + aws_sts_endpoint=aws_sts_endpoint, + aws_external_id=aws_external_id, + ) ) return credentials, aws_region_name @@ -779,13 +780,9 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): proxy_endpoint_url = f"{proxy_endpoint_url}{request_path}" encoded_data: Final = json.dumps(data).encode("utf-8") - # first check api-key, if none, fall back to sigV4 - if api_key is not None: - aws_bearer_token: str | None = api_key - else: - aws_bearer_token = get_secret_str("AWS_BEARER_TOKEN_BEDROCK") + aws_bearer_token: Final = bedrock_bearer_token(api_key) - if aws_bearer_token: + if aws_bearer_token is not None: try: from botocore.awsrequest import AWSRequest except ImportError: @@ -916,7 +913,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): source, ) return BedrockGuardrailResponse() - credentials, aws_region_name = self._load_credentials() + credentials, aws_region_name = self._load_credentials(bearer_token=bedrock_bearer_token(api_key)) allow_chunking: Final = not self._content_uses_contextual_grounding(content) completed_chunk_usages: Final[list[BedrockGuardrailUsage]] = [] # mutable-ok: billed-chunk usage accumulator @@ -958,7 +955,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): self, content: Sequence[BedrockContentItem], base_request_data: Mapping[str, object], - credentials: "Credentials", + credentials: "Credentials | None", aws_region_name: str, api_key: str | None, request_data: dict | None, # mutable-ok: proxy request body dict, mutated by the logging helper @@ -1096,7 +1093,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): self, content: Sequence[BedrockContentItem], base_request_data: Mapping[str, object], - credentials: "Credentials", + credentials: "Credentials | None", aws_region_name: str, api_key: str | None, request_data: dict | None, # mutable-ok: proxy request body dict, mutated by the logging helper @@ -1146,7 +1143,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): self, content: Sequence[BedrockContentItem], base_request_data: Mapping[str, object], - credentials: "Credentials", + credentials: "Credentials | None", aws_region_name: str, api_key: str | None, request_data: dict | None, # mutable-ok: proxy request body dict, mutated by the logging helper @@ -1873,9 +1870,9 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): # Nothing to scan (e.g. tool-only turn) -> allow, like ApplyGuardrail does. return BedrockGuardrailResponse() - credentials, aws_region_name = self._load_credentials() - body: Final[dict[str, object]] = {"messages": checks_messages, "checks": self.checks} api_key: Final[str | None] = request_data.get("api_key") if request_data else None + credentials, aws_region_name = self._load_credentials(bearer_token=bedrock_bearer_token(api_key)) + body: Final[dict[str, object]] = {"messages": checks_messages, "checks": self.checks} prepared_request: Final = self._prepare_request( credentials=credentials, diff --git a/litellm/secret_managers/aws_secret_manager_v2.py b/litellm/secret_managers/aws_secret_manager_v2.py index 2c7f1f8389d..acdb83094e6 100644 --- a/litellm/secret_managers/aws_secret_manager_v2.py +++ b/litellm/secret_managers/aws_secret_manager_v2.py @@ -535,6 +535,7 @@ class AWSSecretsManagerV2(BaseAWSLLM, BaseSecretManager): try: from botocore.auth import SigV4Auth from botocore.awsrequest import AWSRequest + from botocore.exceptions import NoCredentialsError except ImportError: raise ImportError("Missing boto3 to call bedrock. Run 'pip install boto3'.") optional_params = optional_params or {} @@ -582,10 +583,14 @@ class AWSSecretsManagerV2(BaseAWSLLM, BaseSecretManager): "X-Amz-Target": f"secretsmanager.{action}", } + credentials: Final = boto3_credentials_info.credentials + if credentials is None: + raise NoCredentialsError() + # Sign request request: Final = AWSRequest(method="POST", url=endpoint_url, data=body, headers=headers) SigV4Auth( - boto3_credentials_info.credentials, + credentials, "secretsmanager", boto3_credentials_info.aws_region_name, ).add_auth(request) diff --git a/tests/test_litellm/llms/bedrock/chat/test_bedrock_converse_handler.py b/tests/test_litellm/llms/bedrock/chat/test_bedrock_converse_handler.py index 21e3239f623..c4d6896b17b 100644 --- a/tests/test_litellm/llms/bedrock/chat/test_bedrock_converse_handler.py +++ b/tests/test_litellm/llms/bedrock/chat/test_bedrock_converse_handler.py @@ -513,3 +513,27 @@ def test_the_rust_opt_in_needs_no_sigv4_principal(): assert not {"aws_access_key_id", "aws_secret_access_key", "aws_session_token"} & params.keys() assert params["aws_region_name"] == "us-east-1" assert seen["call"][0]["api_key"] == "bedrock-bearer-token" + + +@pytest.mark.parametrize("configured_through", ["env_var", "api_key"]) +def test_bearer_token_auth_never_runs_the_sigv4_credential_chain(monkeypatch, configured_through): + """The deployment's AWS profile does not exist, so resolving SigV4 credentials + raises; a bearer-token deployment must still serve the request, since the + bearer token alone signs it.""" + if configured_through == "env_var": + monkeypatch.setenv("AWS_BEARER_TOKEN_BEDROCK", "bedrock-bearer-token") + else: + monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) + client = _sync_client_returning_converse_response() + + response = BedrockConverseLLM().completion( + **_completion_kwargs( + optional_params={"maxTokens": 16, "aws_profile_name": "litellm-no-such-aws-profile"}, + litellm_params={}, + client=client, + api_key="bedrock-bearer-token" if configured_through == "api_key" else None, + ) + ) + + assert response.choices[0].message.content == "hi" + assert client.post.call_args.kwargs["headers"]["Authorization"] == "Bearer bedrock-bearer-token" diff --git a/tests/test_litellm/llms/bedrock/embed/test_bedrock_embedding.py b/tests/test_litellm/llms/bedrock/embed/test_bedrock_embedding.py index 08d01127eba..50f8bbcf584 100644 --- a/tests/test_litellm/llms/bedrock/embed/test_bedrock_embedding.py +++ b/tests/test_litellm/llms/bedrock/embed/test_bedrock_embedding.py @@ -1033,3 +1033,29 @@ def test_load_credentials_assumes_role_with_external_id(monkeypatch): assert credentials.token == "assumed-session-token" assert aws_region_name == "us-east-1" assert "aws_external_id" not in optional_params + + +def test_bedrock_embedding_bearer_token_never_runs_the_sigv4_credential_chain(monkeypatch): + """The deployment's AWS profile does not exist, so resolving SigV4 credentials + raises; a bearer-token deployment must still serve the request, since the + bearer token alone signs it.""" + monkeypatch.setenv("AWS_BEARER_TOKEN_BEDROCK", "env-bearer-token-12345") + client = HTTPHandler() + + with patch.object(client, "post") as mock_post: + mock_response = Mock() + mock_response.status_code = 200 + mock_response.text = json.dumps(titan_embedding_response) + mock_response.json = lambda: json.loads(mock_response.text) + mock_post.return_value = mock_response + + response = litellm.embedding( + model="bedrock/amazon.titan-embed-text-v1", + input=test_input, + client=client, + aws_region_name="us-west-2", + aws_profile_name="litellm-no-such-aws-profile", + ) + + assert response.data[0]["embedding"] == titan_embedding_response["embedding"] + assert mock_post.call_args.kwargs["headers"]["Authorization"] == "Bearer env-bearer-token-12345" diff --git a/tests/test_litellm/llms/bedrock/image/test_bedrock_image_bearer_token.py b/tests/test_litellm/llms/bedrock/image/test_bedrock_image_bearer_token.py index 7c36b2aa75f..0b11a66c100 100644 --- a/tests/test_litellm/llms/bedrock/image/test_bedrock_image_bearer_token.py +++ b/tests/test_litellm/llms/bedrock/image/test_bedrock_image_bearer_token.py @@ -135,3 +135,24 @@ class TestBedrockImageGeneration: assert response is not None assert len(response.data) > 0 mock_bedrock_image_gen.assert_called_once() + + +def test_image_generation_bearer_token_never_runs_the_sigv4_credential_chain(monkeypatch): + """The deployment's AWS profile does not exist, so resolving SigV4 credentials + raises; a bearer-token deployment must still sign the request with the + bearer token alone.""" + from litellm.llms.bedrock.image_generation.image_handler import BedrockImageGeneration + + monkeypatch.setenv("AWS_BEARER_TOKEN_BEDROCK", "env-bearer-token-12345") + + request = BedrockImageGeneration()._prepare_request( + model="amazon.nova-canvas-v1:0", + prompt="A cute baby sea otter", + optional_params={"aws_region_name": "us-west-2", "aws_profile_name": "litellm-no-such-aws-profile"}, + api_base=None, + extra_headers=None, + api_key=None, + logging_obj=Mock(), + ) + + assert request.prepped.headers["Authorization"] == "Bearer env-bearer-token-12345" diff --git a/tests/test_litellm/llms/bedrock/image_edit/test_amazon_nova_canvas_image_edit.py b/tests/test_litellm/llms/bedrock/image_edit/test_amazon_nova_canvas_image_edit.py index 020b8df1276..58411a9ae18 100644 --- a/tests/test_litellm/llms/bedrock/image_edit/test_amazon_nova_canvas_image_edit.py +++ b/tests/test_litellm/llms/bedrock/image_edit/test_amazon_nova_canvas_image_edit.py @@ -3,6 +3,7 @@ import base64 import io from typing import cast +from unittest.mock import Mock, patch import httpx import pytest @@ -655,3 +656,23 @@ def test_transform_response_empty_images_without_error_raises(): raw_response=resp, logging_obj=None, # type: ignore[arg-type] ) + + +def test_prepare_request_bearer_token_never_runs_the_sigv4_credential_chain(monkeypatch): + """The deployment's AWS profile does not exist, so resolving SigV4 credentials + raises; a bearer-token deployment must still sign the request with the + bearer token alone.""" + monkeypatch.setenv("AWS_BEARER_TOKEN_BEDROCK", "env-bearer-token-12345") + + request = BedrockImageEdit()._prepare_request( + model="amazon.nova-canvas-v1:0", + image=[io.BytesIO(b"fake-png")], + prompt="make it warmer", + optional_params={"aws_region_name": "us-west-2", "aws_profile_name": "litellm-no-such-aws-profile"}, + api_base=None, + extra_headers=None, + logging_obj=Mock(), + api_key=None, + ) + + assert request.prepped.headers["Authorization"] == "Bearer env-bearer-token-12345" diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py index 953e3de1519..14e124981c9 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py @@ -5792,3 +5792,25 @@ async def test_apply_guardrail_debug_log_masks_signed_request_headers(): assert header_lines, "expected the signed-request debug line to be logged" assert any("X-Amz-Security-Token" in message for message in header_lines) assert all(session_token not in message for message in rendered_messages) + + +@pytest.mark.asyncio +async def test_bearer_token_never_runs_the_sigv4_credential_chain(monkeypatch): + """The guardrail's AWS profile does not exist, so resolving SigV4 credentials + raises; with a bearer token configured the guardrail must still run, since + the bearer token alone signs the request.""" + monkeypatch.setenv("AWS_BEARER_TOKEN_BEDROCK", "env-bearer-token-12345") + guardrail = BedrockGuardrail( + guardrailIdentifier="test-guardrail", + guardrailVersion="DRAFT", + aws_profile_name="litellm-no-such-aws-profile", + ) + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.json.return_value = {"action": "NONE", "assessments": []} + + with patch.object(guardrail.async_handler, "post", new_callable=AsyncMock, return_value=mock_response) as mock_post: + response = await guardrail.make_bedrock_api_request(source="INPUT", messages=[{"role": "user", "content": "hello"}]) + + assert response["action"] == "NONE" + assert mock_post.call_args.kwargs["headers"]["Authorization"] == "Bearer env-bearer-token-12345" diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_invoke_guardrail_checks.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_invoke_guardrail_checks.py index d842a1ee5f9..f4af77d2e40 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_invoke_guardrail_checks.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_invoke_guardrail_checks.py @@ -833,3 +833,31 @@ async def test_many_blocks_scanned_at_request_level_and_can_block(): sent_texts = [c["text"] for m in body_messages for c in m["content"]] assert sent_texts == [f"b{i}" for i in range(25)] assert all(len(m["content"]) <= 10 for m in body_messages) + + +@pytest.mark.asyncio +async def test_checks_bearer_token_never_runs_the_sigv4_credential_chain(monkeypatch): + """Same bearer-token rule as ApplyGuardrail: the guardrail's AWS profile does + not exist, yet the InvokeGuardrailChecks call still goes out on the bearer + token and its verdict is enforced.""" + monkeypatch.setenv("AWS_BEARER_TOKEN_BEDROCK", "env-bearer-token-12345") + g = BedrockGuardrail( + checks=CONTENT_FILTER_CHECKS, + content_filter_threshold=0.5, + aws_profile_name="litellm-no-such-aws-profile", + ) + payload = {"results": {"contentFilter": {"results": [{"category": "VIOLENCE", "severityScore": 0.8}]}}} + post = AsyncMock(return_value=_mock_http_response(200, payload)) + + with patch.object(g.async_handler, "post", new=post): + with pytest.raises(HTTPException) as exc: + await g.make_bedrock_api_request( + source="INPUT", + messages=[{"role": "user", "content": "hi"}], + request_data={"messages": []}, + ) + + assert exc.value.detail["bedrock_guardrail_checks"] == [ + {"check": "contentFilter", "category": "VIOLENCE", "severityScore": 0.8} + ] + assert post.call_args.kwargs["headers"]["Authorization"] == "Bearer env-bearer-token-12345" From 0522110ddab688e3981fe43f3c972f9a595db074 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 2 Sep 2026 16:30:14 -0700 Subject: [PATCH 2/3] fix(bedrock): type the bearer path with overloads instead of None guards _get_boto_credentials_from_optional_params and BedrockEmbedding._load_credentials gain typed overloads, so callers that never pass a bearer token (rerank, the secrets manager, async-invoke status polling) keep a non-null Credentials and need no guard. The bearer branch returns a BearerRequestTarget instead of a Boto3CredentialsInfo holding None, and the secrets manager is back to its unchanged base version. The two guardrail-endpoint tests that patched the removed get_secret_str import now drive AWS_BEARER_TOKEN_BEDROCK through the environment. --- litellm/llms/bedrock/base_aws_llm.py | 61 +++++++++++++------ litellm/llms/bedrock/embed/embedding.py | 19 ++++-- .../secret_managers/aws_secret_manager_v2.py | 7 +-- .../guardrails/test_guardrail_endpoints.py | 21 ++----- 4 files changed, 64 insertions(+), 44 deletions(-) diff --git a/litellm/llms/bedrock/base_aws_llm.py b/litellm/llms/bedrock/base_aws_llm.py index 1f00bf7792e..c3da992a904 100644 --- a/litellm/llms/bedrock/base_aws_llm.py +++ b/litellm/llms/bedrock/base_aws_llm.py @@ -7,7 +7,7 @@ import urllib.parse from collections.abc import Callable from datetime import datetime from threading import Lock -from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, cast, get_args +from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, cast, get_args, overload import httpx from pydantic import BaseModel, ValidationError @@ -48,12 +48,19 @@ _STS_REGION_FROM_ENDPOINT_PATTERN: Final = re.compile( SIGV4_COMPUTED_HEADERS: Final = frozenset({"authorization", "x-amz-date", "x-amz-security-token", "date"}) -class Boto3CredentialsInfo(BaseModel): - credentials: Credentials | None +class BedrockRequestTarget(BaseModel): aws_region_name: str aws_bedrock_runtime_endpoint: str | None +class Boto3CredentialsInfo(BedrockRequestTarget): + credentials: Credentials + + +class BearerRequestTarget(BedrockRequestTarget): + credentials: None = None + + def bedrock_bearer_token(api_key: str | None) -> str | None: token: Final = api_key if api_key is not None else get_secret_str("AWS_BEARER_TOKEN_BEDROCK") return token or None @@ -1392,9 +1399,26 @@ class BaseAWSLLM: else: return f"https://bedrock-runtime.{aws_region_name}.{dns_suffix}" + @overload + def _get_boto_credentials_from_optional_params( + self, + optional_params: dict, # mutable-ok: the implementation pops the aws_* keys out of the caller's dict in place + model: str | None = None, + bearer_token: None = None, + ) -> Boto3CredentialsInfo: ... + + @overload + def _get_boto_credentials_from_optional_params( + self, + optional_params: dict, # mutable-ok: the implementation pops the aws_* keys out of the caller's dict in place + model: str | None = None, + *, + bearer_token: str, + ) -> BearerRequestTarget: ... + def _get_boto_credentials_from_optional_params( self, optional_params: dict, model: str | None = None, bearer_token: str | None = None - ) -> Boto3CredentialsInfo: + ) -> Boto3CredentialsInfo | BearerRequestTarget: """ Get boto3 credentials from optional params @@ -1425,23 +1449,24 @@ class BaseAWSLLM: ) # https://bedrock-runtime.{region_name}.amazonaws.com aws_external_id: Final = optional_params.pop("aws_external_id", None) - credentials: Final[Credentials | None] = ( - None - if bearer_token is not None - else self.get_credentials( - aws_access_key_id=aws_access_key_id, - aws_secret_access_key=aws_secret_access_key, - aws_session_token=aws_session_token, + if bearer_token is not None: + return BearerRequestTarget( aws_region_name=aws_region_name, - aws_session_name=aws_session_name, - aws_profile_name=aws_profile_name, - aws_role_name=aws_role_name, - aws_web_identity_token=aws_web_identity_token, - aws_sts_endpoint=aws_sts_endpoint, - aws_external_id=aws_external_id, + aws_bedrock_runtime_endpoint=aws_bedrock_runtime_endpoint, ) - ) + credentials: Final[Credentials] = self.get_credentials( + aws_access_key_id=aws_access_key_id, + aws_secret_access_key=aws_secret_access_key, + aws_session_token=aws_session_token, + aws_region_name=aws_region_name, + aws_session_name=aws_session_name, + aws_profile_name=aws_profile_name, + aws_role_name=aws_role_name, + aws_web_identity_token=aws_web_identity_token, + aws_sts_endpoint=aws_sts_endpoint, + aws_external_id=aws_external_id, + ) return Boto3CredentialsInfo( credentials=credentials, aws_region_name=aws_region_name, diff --git a/litellm/llms/bedrock/embed/embedding.py b/litellm/llms/bedrock/embed/embedding.py index a7b74f3752a..5fb86d476f4 100644 --- a/litellm/llms/bedrock/embed/embedding.py +++ b/litellm/llms/bedrock/embed/embedding.py @@ -6,7 +6,7 @@ import copy import json import urllib.parse from collections.abc import Callable -from typing import TYPE_CHECKING, Final, get_args +from typing import TYPE_CHECKING, Final, get_args, overload import httpx @@ -42,6 +42,20 @@ if TYPE_CHECKING: class BedrockEmbedding(BaseAWSLLM): + @overload + def _load_credentials( + self, + optional_params: dict, # mutable-ok: the implementation pops the aws_* keys out of the caller's dict in place + bearer_token: None = None, + ) -> tuple[Credentials, str]: ... + + @overload + def _load_credentials( + self, + optional_params: dict, # mutable-ok: the implementation pops the aws_* keys out of the caller's dict in place + bearer_token: str, + ) -> tuple[None, str]: ... + def _load_credentials( self, optional_params: dict, @@ -598,11 +612,8 @@ class BedrockEmbedding(BaseAWSLLM): try: from botocore.auth import SigV4Auth from botocore.awsrequest import AWSRequest - from botocore.exceptions import NoCredentialsError except ImportError: raise ImportError("Missing boto3 to call bedrock. Run 'pip install boto3'.") - if credentials is None: - raise NoCredentialsError() # Create AWSRequest with GET method and encoded URL request: Final = AWSRequest( diff --git a/litellm/secret_managers/aws_secret_manager_v2.py b/litellm/secret_managers/aws_secret_manager_v2.py index acdb83094e6..2c7f1f8389d 100644 --- a/litellm/secret_managers/aws_secret_manager_v2.py +++ b/litellm/secret_managers/aws_secret_manager_v2.py @@ -535,7 +535,6 @@ class AWSSecretsManagerV2(BaseAWSLLM, BaseSecretManager): try: from botocore.auth import SigV4Auth from botocore.awsrequest import AWSRequest - from botocore.exceptions import NoCredentialsError except ImportError: raise ImportError("Missing boto3 to call bedrock. Run 'pip install boto3'.") optional_params = optional_params or {} @@ -583,14 +582,10 @@ class AWSSecretsManagerV2(BaseAWSLLM, BaseSecretManager): "X-Amz-Target": f"secretsmanager.{action}", } - credentials: Final = boto3_credentials_info.credentials - if credentials is None: - raise NoCredentialsError() - # Sign request request: Final = AWSRequest(method="POST", url=endpoint_url, data=body, headers=headers) SigV4Auth( - credentials, + boto3_credentials_info.credentials, "secretsmanager", boto3_credentials_info.aws_region_name, ).add_auth(request) diff --git a/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py b/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py index 9511732fd50..a222e22f6d0 100644 --- a/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py +++ b/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py @@ -771,7 +771,7 @@ async def test_bedrock_guardrail_prepare_request_with_api_key(): @pytest.mark.asyncio -async def test_bedrock_guardrail_prepare_request_without_api_key(): +async def test_bedrock_guardrail_prepare_request_without_api_key(monkeypatch): """Test _prepare_request method falls back to SigV4 when no api_key is provided""" from unittest.mock import Mock, patch @@ -789,18 +789,13 @@ async def test_bedrock_guardrail_prepare_request_without_api_key(): # Test data without api_key test_data = {"source": "INPUT", "content": [{"text": {"text": "test content"}}]} + monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) with ( - patch( - "litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails.get_secret_str" - ) as mock_get_secret, patch("botocore.auth.SigV4Auth") as mock_sigv4_auth, patch("botocore.awsrequest.AWSRequest") as mock_aws_request, ): - # Mock no AWS_BEARER_TOKEN_BEDROCK - mock_get_secret.return_value = None - # Mock SigV4Auth mock_sigv4_instance = Mock() mock_sigv4_auth.return_value = mock_sigv4_instance @@ -826,7 +821,7 @@ async def test_bedrock_guardrail_prepare_request_without_api_key(): @pytest.mark.asyncio -async def test_bedrock_guardrail_prepare_request_with_bearer_token_env(): +async def test_bedrock_guardrail_prepare_request_with_bearer_token_env(monkeypatch): """Test _prepare_request method uses Bearer token from environment when available""" from unittest.mock import Mock, patch @@ -844,15 +839,9 @@ async def test_bedrock_guardrail_prepare_request_with_bearer_token_env(): # Test data without api_key test_data = {"source": "INPUT", "content": [{"text": {"text": "test content"}}]} + monkeypatch.setenv("AWS_BEARER_TOKEN_BEDROCK", "env-bearer-token-456") - with ( - patch( - "litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails.get_secret_str" - ) as mock_get_secret, - patch("botocore.awsrequest.AWSRequest") as mock_aws_request, - ): - - mock_get_secret.return_value = "env-bearer-token-456" + with patch("botocore.awsrequest.AWSRequest") as mock_aws_request: mock_request_instance = Mock() mock_request_instance.prepare.return_value = Mock() mock_aws_request.return_value = mock_request_instance From fa5a90e08e867dccf881c537302c55e62fb77134 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 3 Sep 2026 01:20:27 -0700 Subject: [PATCH 3/3] fix(bedrock): stop the Moonshot invoke transform from resolving AWS credentials AmazonMoonshotConfig.transform_request called _get_boto_credentials_from_optional_params purely for its side effect of popping the aws_* keys off optional_params, then threw the result away. On a box whose default AWS profile uses login_session without botocore[crt], that call raises, so a bearer-token bedrock/invoke/moonshot.* deployment still 500s with MissingDependencyException even after the rest of this branch skips the chain. It now filters the aws_* keys into a local dict the way the Qwen, OpenAI and Claude 3 invoke transformations already do, so no credentials are resolved and the caller's optional_params keeps the keys sign_request reads afterwards. --- .../amazon_moonshot_transformation.py | 8 +-- .../test_amazon_moonshot_transformation.py | 66 +++++++++++++++++++ 2 files changed, 68 insertions(+), 6 deletions(-) create mode 100644 tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_amazon_moonshot_transformation.py diff --git a/litellm/llms/bedrock/chat/invoke_transformations/amazon_moonshot_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/amazon_moonshot_transformation.py index 91c3a363c31..04c6ec86a13 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/amazon_moonshot_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/amazon_moonshot_transformation.py @@ -149,19 +149,15 @@ class AmazonMoonshotConfig(AmazonInvokeConfig, MoonshotChatConfig): - Temperature and parameter validation """ - # Filter out AWS credentials using the existing method from BaseAWSLLM - self._get_boto_credentials_from_optional_params(optional_params, model) + inference_params: Final = {k: v for k, v in optional_params.items() if k not in self.aws_authentication_params} - # Strip routing prefixes to get the actual model ID clean_model_id: Final = self._get_model_id(model) - # Use Moonshot's transform_request which handles message transformation - # and tool_choice="required" workaround return MoonshotChatConfig.transform_request( self, model=clean_model_id, messages=messages, - optional_params=optional_params, + optional_params=inference_params, litellm_params=litellm_params, headers=headers, ) diff --git a/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_amazon_moonshot_transformation.py b/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_amazon_moonshot_transformation.py new file mode 100644 index 00000000000..531f334e460 --- /dev/null +++ b/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_amazon_moonshot_transformation.py @@ -0,0 +1,66 @@ +import pytest + +from litellm.llms.bedrock.chat.invoke_transformations.amazon_moonshot_transformation import ( + AmazonMoonshotConfig, +) + +AWS_AUTH_PARAMS = { + "aws_access_key_id": "AKIAEXAMPLE", + "aws_secret_access_key": "secret", + "aws_session_token": "token", + "aws_region_name": "us-west-2", + "aws_session_name": "session", + "aws_role_name": "arn:aws:iam::000000000000:role/example", + "aws_web_identity_token": "web-identity", + "aws_sts_endpoint": "https://sts.us-west-2.amazonaws.com", + "aws_bedrock_runtime_endpoint": "https://bedrock-runtime.us-west-2.amazonaws.com", + "aws_external_id": "external", +} + + +def test_transform_request_never_resolves_aws_credentials(): + """A broken credential chain must not stop the request body from being built.""" + config = AmazonMoonshotConfig() + + transformed = config.transform_request( + model="bedrock/invoke/moonshot.kimi-k2-thinking", + messages=[{"role": "user", "content": "Hello"}], + optional_params={"aws_profile_name": "litellm-profile-that-does-not-exist", "max_tokens": 16}, + litellm_params={}, + headers={}, + ) + + assert transformed["model"] == "moonshot.kimi-k2-thinking" + assert transformed["max_tokens"] == 16 + assert "aws_profile_name" not in transformed + + +@pytest.mark.parametrize("aws_param", sorted(AWS_AUTH_PARAMS)) +def test_transform_request_keeps_aws_params_out_of_the_body(aws_param: str): + config = AmazonMoonshotConfig() + + transformed = config.transform_request( + model="bedrock/invoke/moonshot.kimi-k2-thinking", + messages=[{"role": "user", "content": "Hello"}], + optional_params={aws_param: AWS_AUTH_PARAMS[aws_param]}, + litellm_params={}, + headers={}, + ) + + assert aws_param not in transformed + + +def test_transform_request_leaves_the_caller_aws_params_in_place_for_signing(): + """sign_request reads the aws_* keys off optional_params after transform_request runs.""" + config = AmazonMoonshotConfig() + optional_params = dict(AWS_AUTH_PARAMS) + + config.transform_request( + model="bedrock/invoke/moonshot.kimi-k2-thinking", + messages=[{"role": "user", "content": "Hello"}], + optional_params=optional_params, + litellm_params={}, + headers={}, + ) + + assert optional_params == AWS_AUTH_PARAMS