diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 3cd6ee54069..26c085a95e6 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/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, diff --git a/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_transformation.py b/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_transformation.py index cd775abf136..56b111f294e 100644 --- a/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_transformation.py +++ b/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): diff --git a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py index 26f841c1146..1d583c16ad7 100644 --- a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py +++ b/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."""