Merge pull request #32956 from BerriAI/litellm_fix_lit3859_wif_bridge

fix(completion): forward aws credential kwargs into litellm_params so the responses bridge keeps WIF auth

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
 

litellm/llms/bedrock/base_aws_llm.py

@@ -877,6 +877,15 @@ class BaseAWSLLM:
                     "Resource": "*",
                     "Condition": {"Bool": {"aws:SecureTransport": "true"}},
                 },
+                {
+                    "Sid": "BedrockMantleLiteLLM",
+                    "Effect": "Allow",
+                    "Action": [
+                        "bedrock-mantle:CreateInference",
+                    ],
+                    "Resource": "*",
+                    "Condition": {"Bool": {"aws:SecureTransport": "true"}},
+                },
             ],
         }
         assume_role_params = {

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,

tests/test_litellm/llms/bedrock/test_web_identity_session_policy.py

@@ -158,6 +158,54 @@ class TestClaudePlatformActionsCovered:
         )
 
 
+class TestBedrockMantleActionsCovered:
+    """LIT-3859: bedrock_mantle inference authorizes against the
+    ``bedrock-mantle`` action namespace, so the session-policy ceiling
+    must include it or every Mantle request via OIDC/WIF auth denies
+    with "no session policy allows the bedrock-mantle:CreateInference
+    action" even when the role's identity policy grants it."""
+
+    def test_bedrock_mantle_create_inference_present(self):
+        policy = _captured_policy()
+        all_actions: set = set()
+        for stmt in policy["Statement"]:
+            stmt_actions = stmt.get("Action")
+            if isinstance(stmt_actions, str):
+                all_actions.add(stmt_actions)
+            elif isinstance(stmt_actions, list):
+                all_actions.update(stmt_actions)
+        assert "bedrock-mantle:CreateInference" in all_actions, (
+            "bedrock-mantle:CreateInference missing from session policy — "
+            "bedrock_mantle/* requests will 403 on OIDC/WIF auth"
+        )
+
+    def test_bedrock_mantle_statement_allows(self):
+        policy = _captured_policy()
+        stmt = _statement_by_sid(policy, "BedrockMantleLiteLLM")
+        assert stmt["Effect"] == "Allow"
+        assert stmt["Resource"] == "*"
+
+    def test_no_bedrock_mantle_wildcard(self):
+        policy = _captured_policy()
+        stmt = _statement_by_sid(policy, "BedrockMantleLiteLLM")
+        actions = stmt["Action"]
+        if isinstance(actions, str):
+            actions = [actions]
+        assert "bedrock-mantle:*" not in actions, (
+            "session policy must not grant bedrock-mantle:* — "
+            "the ceiling should match the documented action set"
+        )
+
+    def test_bedrock_mantle_statement_carries_secure_transport_condition(self):
+        policy = _captured_policy()
+        stmt = _statement_by_sid(policy, "BedrockMantleLiteLLM")
+        cond = stmt.get("Condition") or {}
+        assert cond.get("Bool", {}).get("aws:SecureTransport") == "true", (
+            "BedrockMantleLiteLLM must require aws:SecureTransport=true "
+            "to keep parity with the bedrock statement"
+        )
+
+
 def _make_jwt(payload: dict) -> str:
     def _segment(data: dict) -> str:
         return base64.urlsafe_b64encode(json.dumps(data).encode()).rstrip(b"=").decode()

tests/test_litellm/test_main.py

@@ -2081,3 +2081,79 @@ 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
+@pytest.mark.parametrize(
+    "aws_credential_kwargs",
+    [
+        {
+            "aws_session_name": "litellm-gcp",
+            "aws_role_name": "arn:aws:iam::123456789012:role/litellm-bedrock-role",
+            "aws_web_identity_token": "oidc/google/108963886734710037768",
+        },
+        {
+            "aws_access_key_id": "AKIASTATICKEYFORTEST",
+            "aws_secret_access_key": "static-secret-key",
+            "aws_session_token": "static-session-token",
+        },
+    ],
+    ids=["web_identity", "static_keys"],
+)
+async def test_acompletion_forwards_aws_credentials_through_responses_bridge(
+    respx_mock: respx.MockRouter, monkeypatch, aws_credential_kwargs: dict
+):
+    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",
+            num_retries=0,
+            **aws_credential_kwargs,
+        )
+
+        assert response.choices[0].message.content == "ok"
+        credential_kwargs = get_credentials_mock.call_args.kwargs
+        assert credential_kwargs["aws_region_name"] == "us-east-2"
+        for key, value in aws_credential_kwargs.items():
+            assert credential_kwargs[key] == value
+        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()