From e0463a38fffdeb32dec9297039ede25a113ae543 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 11 Jul 2026 16:10:39 -0700 Subject: [PATCH] fix(completion): forward aws credential kwargs into litellm_params so the responses bridge keeps WIF auth Chat-completions requests to responses-only Bedrock Mantle models are bridged to the Responses API, but completion() forwarded only aws_bedrock_project_id into get_litellm_params, so aws_role_name, aws_web_identity_token, aws_session_name and the other SigV4 credential kwargs never reached sign_request and botocore fell back to the default credential chain ("Bedrock Mantle auth failed: no Bearer token and no usable AWS credentials"). Forward the whole AWS credential kwarg family, extracted from the OPTIONAL_KWARGS_KEYS set get_litellm_params already supports. --- .../litellm_core_utils/get_litellm_params.py | 56 ++++++++++------- litellm/main.py | 7 ++- tests/test_litellm/test_main.py | 62 +++++++++++++++++++ 3 files changed, 99 insertions(+), 26 deletions(-) diff --git a/litellm/litellm_core_utils/get_litellm_params.py b/litellm/litellm_core_utils/get_litellm_params.py index 352e55e9c23..b8ef9d8cca7 100644 --- a/litellm/litellm_core_utils/get_litellm_params.py +++ b/litellm/litellm_core_utils/get_litellm_params.py @@ -2,26 +2,8 @@ from typing import Optional from litellm.llms.openai.data_residency import infer_openai_data_residency -# Pre-define optional kwargs keys as frozenset for O(1) lookups -# These are extracted from kwargs only if present, avoiding unnecessary .get() calls -OPTIONAL_KWARGS_KEYS = frozenset( +AWS_CREDENTIAL_KWARGS_KEYS = frozenset( { - "azure_ad_token", - "tenant_id", - "client_id", - "client_secret", - "azure_username", - "azure_password", - "azure_scope", - "timeout", - "gcs_bucket_name", - "bucket_name", - "vertex_credentials", - "vertex_project", - "vertex_location", - "vertex_ai_project", - "vertex_ai_location", - "vertex_ai_credentials", "aws_region_name", "aws_access_key_id", "aws_secret_access_key", @@ -34,14 +16,40 @@ OPTIONAL_KWARGS_KEYS = frozenset( "aws_external_id", "aws_bedrock_runtime_endpoint", "aws_bedrock_project_id", - "tpm", - "rpm", - "itpm", - "otpm", - "use_xai_oauth", } ) +# Pre-define optional kwargs keys as frozenset for O(1) lookups +# These are extracted from kwargs only if present, avoiding unnecessary .get() calls +OPTIONAL_KWARGS_KEYS = ( + frozenset( + { + "azure_ad_token", + "tenant_id", + "client_id", + "client_secret", + "azure_username", + "azure_password", + "azure_scope", + "timeout", + "gcs_bucket_name", + "bucket_name", + "vertex_credentials", + "vertex_project", + "vertex_location", + "vertex_ai_project", + "vertex_ai_location", + "vertex_ai_credentials", + "tpm", + "rpm", + "itpm", + "otpm", + "use_xai_oauth", + } + ) + | AWS_CREDENTIAL_KWARGS_KEYS +) + # Backward-compatible alias for existing imports/tests. _OPTIONAL_KWARGS_KEYS = OPTIONAL_KWARGS_KEYS diff --git a/litellm/main.py b/litellm/main.py index 7d457d9cdd1..6fd68921fb0 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -92,7 +92,10 @@ from litellm.litellm_core_utils.completion_timeout import CompletionTimeout from litellm.litellm_core_utils.request_timeout_resolver import ( get_configured_request_timeout, ) -from litellm.litellm_core_utils.get_litellm_params import OPTIONAL_KWARGS_KEYS +from litellm.litellm_core_utils.get_litellm_params import ( + AWS_CREDENTIAL_KWARGS_KEYS, + OPTIONAL_KWARGS_KEYS, +) from litellm.litellm_core_utils.dd_tracing import tracer from litellm.litellm_core_utils.get_provider_specific_headers import ( ProviderSpecificHeaderUtils, @@ -5322,7 +5325,7 @@ def completion( # type: ignore tpm=kwargs.get("tpm"), rpm=kwargs.get("rpm"), use_xai_oauth=kwargs.get("use_xai_oauth", False), - aws_bedrock_project_id=kwargs.get("aws_bedrock_project_id"), + **{key: kwargs[key] for key in AWS_CREDENTIAL_KWARGS_KEYS if key in kwargs}, ) cast(LiteLLMLoggingObj, logging).update_environment_variables( model=model, diff --git a/tests/test_litellm/test_main.py b/tests/test_litellm/test_main.py index 28cf4fa0744..c5e70e0fabf 100644 --- a/tests/test_litellm/test_main.py +++ b/tests/test_litellm/test_main.py @@ -2081,3 +2081,65 @@ def test_stream_chunk_builder_text_completion_combines_text_and_usage(): assert response.usage.prompt_tokens > 0 assert response.usage.completion_tokens > 0 assert response.usage.total_tokens == response.usage.prompt_tokens + response.usage.completion_tokens + + +@pytest.mark.asyncio +async def test_acompletion_forwards_aws_credentials_through_responses_bridge( + respx_mock: respx.MockRouter, monkeypatch +): + from botocore.credentials import Credentials + + from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM + + original_disable_aiohttp = litellm.disable_aiohttp_transport + try: + litellm.disable_aiohttp_transport = True + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + litellm.in_memory_llm_clients_cache.flush_cache() + monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) + monkeypatch.delenv("BEDROCK_MANTLE_API_KEY", raising=False) + + get_credentials_mock = MagicMock(return_value=Credentials("fake-key", "fake-secret")) + monkeypatch.setattr(BaseAWSLLM, "get_credentials", get_credentials_mock) + + respx_mock.post("https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses").respond( + json={ + "id": "resp_123", + "object": "response", + "created_at": 1760144904, + "status": "completed", + "model": "openai.gpt-5.4", + "output": [ + { + "type": "message", + "id": "msg_1", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": "ok", "annotations": []}], + } + ], + } + ) + + response = await litellm.acompletion( + model="bedrock_mantle/openai.gpt-5.4", + messages=[{"role": "user", "content": "hi"}], + api_base="https://bedrock-mantle.us-east-2.api.aws/v1", + aws_region_name="us-east-2", + aws_session_name="litellm-gcp", + aws_role_name="arn:aws:iam::123456789012:role/litellm-bedrock-role", + aws_web_identity_token="oidc/google/108963886734710037768", + num_retries=0, + ) + + assert response.choices[0].message.content == "ok" + credential_kwargs = get_credentials_mock.call_args.kwargs + assert credential_kwargs["aws_role_name"] == "arn:aws:iam::123456789012:role/litellm-bedrock-role" + assert credential_kwargs["aws_web_identity_token"] == "oidc/google/108963886734710037768" + assert credential_kwargs["aws_session_name"] == "litellm-gcp" + authorization = respx_mock.calls.last.request.headers["Authorization"] + assert authorization.startswith("AWS4-HMAC-SHA256") + assert "fake-key" in authorization + finally: + litellm.disable_aiohttp_transport = original_disable_aiohttp + litellm.in_memory_llm_clients_cache.flush_cache()