Merge pull request #39362 from BerriAI/litellm_lit5443_mantle_chat_aws_creds

fix(bedrock_mantle): carry per-request AWS credentials into chat completions SigV4 signing

litellm/llms/custom_httpx/llm_http_handler.py

@@ -28,6 +28,7 @@ from litellm.litellm_core_utils.audio_utils.subtitle_utils import (
     SUBTITLE_RESPONSE_FORMATS,
     synthesize_subtitle_document,
 )
+from litellm.litellm_core_utils.get_litellm_params import AWS_CREDENTIAL_KWARGS_KEYS
 from litellm.litellm_core_utils.llm_request_utils import serialize_multipart_form_fields
 from litellm.litellm_core_utils.realtime_errors import realtime_error_event, websocket_close_reason
 from litellm.litellm_core_utils.realtime_streaming import RealTimeStreaming
@@ -274,6 +275,16 @@ def _has_pre_call_deployment_hook(logging_obj: LiteLLMLoggingObj) -> bool:
     return False
 
 
+def _aws_signing_overrides(optional_params: Mapping[str, Any], litellm_params: Mapping[str, Any]) -> Mapping[str, Any]:
+    return MappingProxyType(
+        {
+            key: litellm_params[key]
+            for key in AWS_CREDENTIAL_KWARGS_KEYS
+            if optional_params.get(key) is None and litellm_params.get(key) is not None
+        }
+    )
+
+
 def _collect_ws_project_quota_callbacks() -> tuple[ProjectQuotaCallback, ...]:
     """Duck-type discover proxy hooks exposing per-frame project ITPM/OTPM
     enforcement, so the Responses WebSocket loop can charge every
@@ -538,7 +549,10 @@ class BaseLLMHTTPHandler:
 
         headers, signed_json_body = provider_config.sign_request(
             headers=headers,
-            optional_params=optional_params,
+            optional_params={
+                **optional_params,
+                **_aws_signing_overrides(optional_params, litellm_params),
+            },
             request_data=data,
             api_base=api_base,
             api_key=api_key,

tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_transformation.py

@@ -486,6 +486,71 @@ class TestBedrockMantleChatAuth:
         assert "/us-east-2/bedrock/aws4_request" in authorization
         assert requests[0]["url"].startswith("https://bedrock-mantle.us-east-2.api.aws")
 
+    def test_completion_per_request_role_reaches_signer_and_not_the_body(self, monkeypatch):
+        from unittest.mock import MagicMock, Mock
+
+        from botocore.credentials import Credentials
+
+        from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
+        from litellm.llms.custom_httpx.http_handler import HTTPHandler
+        from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler
+        from litellm.types.utils import ModelResponse
+
+        for var in ("BEDROCK_MANTLE_API_KEY", "AWS_BEARER_TOKEN_BEDROCK", "BEDROCK_MANTLE_API_BASE"):
+            monkeypatch.delenv(var, raising=False)
+
+        signer = BaseAWSLLM()
+        signer.get_credentials = MagicMock(
+            return_value=Credentials(
+                access_key="ASIAEXAMPLE",
+                secret_key="YXNzdW1lZC1yb2xlLXNlY3JldC1hc3N1bWVk",
+                token="assumed-session-token",
+            )
+        )
+        url = "https://bedrock-mantle.us-east-1.api.aws/openai/v1/chat/completions"
+        client = HTTPHandler(client=httpx.Client())
+        client.post = Mock(
+            return_value=httpx.Response(
+                status_code=200,
+                json={
+                    "id": "chatcmpl-test",
+                    "object": "chat.completion",
+                    "created": 1733529600,
+                    "model": "google.gemma-4-31b",
+                    "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}],
+                    "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2},
+                },
+                request=httpx.Request("POST", url),
+            )
+        )
+
+        BaseLLMHTTPHandler().completion(
+            model="google.gemma-4-31b",
+            messages=[{"role": "user", "content": "hello"}],
+            api_base=None,
+            custom_llm_provider="bedrock_mantle",
+            model_response=ModelResponse(),
+            encoding=None,
+            logging_obj=Mock(),
+            optional_params={},
+            timeout=10,
+            litellm_params={
+                "aws_role_name": "arn:aws:iam::000000000000:role/attributed-role",
+                "aws_session_name": "user-123",
+                "aws_region_name": "us-east-1",
+            },
+            acompletion=False,
+            client=client,
+            provider_config=BedrockMantleChatConfig(aws_signer=signer),
+        )
+
+        credential_kwargs = signer.get_credentials.call_args.kwargs
+        assert credential_kwargs["aws_role_name"] == "arn:aws:iam::000000000000:role/attributed-role"
+        assert credential_kwargs["aws_session_name"] == "user-123"
+        sent = client.post.call_args.kwargs
+        assert sent["headers"]["Authorization"].startswith("AWS4-HMAC-SHA256")
+        assert not [key for key in json.loads(sent["data"]) if key.startswith("aws_")]
+
 
 class TestBedrockMantleProjectHeader:
     def test_validate_environment_sets_openai_project_header(self):

tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py

@@ -2295,6 +2295,25 @@ async def test_anthropic_invalid_thinking_signature_retry_resigns_bedrock_reques
     assert retry_authorization != first_attempt_headers["Authorization"]
 
 
+def test_aws_signing_overrides_only_fills_missing_credentials():
+    from litellm.llms.custom_httpx.llm_http_handler import _aws_signing_overrides
+
+    overrides = _aws_signing_overrides(
+        {"temperature": 0.2, "aws_region_name": "us-west-2"},
+        {
+            "aws_role_name": "arn:aws:iam::000000000000:role/attributed",
+            "aws_session_name": "user-123",
+            "aws_region_name": "us-east-1",
+            "api_key": "not-an-aws-param",
+        },
+    )
+
+    assert dict(overrides) == {
+        "aws_role_name": "arn:aws:iam::000000000000:role/attributed",
+        "aws_session_name": "user-123",
+    }
+
+
 class TestServerFulfilledToolsInRequest:
     """_server_fulfilled_tools_in_request gates the buffered (non-leaking) streaming
     mode for server-fulfilled tools like headroom_retrieve."""