tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-08-13 02:22:50 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(bedrock_mantle): add SigV4/IAM auth to Responses API route (#29788)

Browse Source 
Applied as the squash diff of PR #29788 (head 9800b2f17c), which landed
upstream inside the litellm_oss_staging_080626 sync (32c88ca74f, #29932)
and has no standalone commit to cherry-pick.
This commit is contained in:
Kent
2026-06-11 03:03:39 +00:00
committed by mateo-berri
parent dad0894dff
commit 63edbf3fa2
5 changed files with 833 additions and 48 deletions
Download Patch File Download Diff File Expand all files Collapse all files
litellm/llms/base_llm/responses/transformation.py
+20
View File
@@ -62,6 +62,26 @@ class BaseResponsesAPIConfig(ABC):
"""
return False
def sign_request(
self,
headers: dict,
optional_params: dict,
request_data: dict,
api_base: str,
api_key: Optional[str] = None,
model: Optional[str] = None,
stream: Optional[bool] = None,
fake_stream: Optional[bool] = None,
) -> Tuple[dict, Optional[bytes]]:
"""Sign the request after the body is finalized.
Default is a no-op (returns headers unchanged, no signed body). Providers
whose endpoint requires request signing (e.g. Bedrock Mantle SigV4)
override this and return the signed body bytes so the handler sends those
exact bytes.
"""
return headers, None
@abstractmethod
def get_supported_openai_params(self, model: str) -> list:
pass
litellm/llms/bedrock_mantle/responses/transformation.py
+105 -15
View File
@@ -4,14 +4,26 @@ Amazon Bedrock Mantle - Responses API backend.
gpt-5.5 / gpt-5.4 on Mantle are exposed ONLY on the `/openai/v1/responses`
path (not the standard `/v1/responses`). Payloads and SSE follow the OpenAI
Responses spec, so this config inherits OpenAIResponsesAPIConfig and overrides
only the endpoint URL and Bearer authentication.
only the endpoint URL and authentication.
Auth: AWS Bedrock API key as Bearer token (BEDROCK_MANTLE_API_KEY or the
standard AWS_BEARER_TOKEN_BEDROCK), NOT SigV4.
Auth: Bearer token (BEDROCK_MANTLE_API_KEY or the standard
AWS_BEARER_TOKEN_BEDROCK, or litellm_params.api_key) when present; otherwise
AWS SigV4 (service name "bedrock") using the standard credential chain (IAM
role / access key / profile / web identity), signed via the shared
BaseAWSLLM._sign_request after the request body is finalized.
"""
from typing import Optional
import re
from typing import Optional, Tuple
from botocore.exceptions import (
CredentialRetrievalError,
NoCredentialsError,
PartialCredentialsError,
ProfileNotFound,
)
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig
from litellm.secret_managers.main import get_secret_str
from litellm.types.router import GenericLiteLLMParams
@@ -29,22 +41,44 @@ _BASE_SUFFIXES_TO_STRIP = (
"/v1",
)
# Standard Mantle host: https://bedrock-mantle.<region>.api.aws (group 1 = region).
_MANTLE_HOST_RE = re.compile(
r"^https?://bedrock-mantle\.([^/.]+)\.api\.aws", re.IGNORECASE
)
class BedrockMantleResponsesAPIConfig(OpenAIResponsesAPIConfig):
def __init__(self, aws_signer: Optional[BaseAWSLLM] = None):
super().__init__()
self._aws_signer = aws_signer or BaseAWSLLM()
@property
def custom_llm_provider(self) -> LlmProviders:
return LlmProviders.BEDROCK_MANTLE
@staticmethod
def _resolve_region(params: dict) -> str:
region = params.get("aws_region_name")
if region:
return region
base = params.get("api_base") or get_secret_str("BEDROCK_MANTLE_API_BASE")
if base:
match = _MANTLE_HOST_RE.match(base.rstrip("/"))
if match:
return match.group(1)
return (
get_secret_str("BEDROCK_MANTLE_REGION")
or get_secret_str("AWS_REGION_NAME")
or get_secret_str("AWS_REGION")
or BEDROCK_MANTLE_DEFAULT_REGION
)
def get_complete_url(
self,
api_base: Optional[str],
litellm_params: dict,
) -> str:
region = (
get_secret_str("BEDROCK_MANTLE_REGION")
or get_secret_str("AWS_REGION")
or BEDROCK_MANTLE_DEFAULT_REGION
)
region = self._resolve_region({**litellm_params, "api_base": api_base})
base = (
api_base
or get_secret_str("BEDROCK_MANTLE_API_BASE")
@@ -55,6 +89,11 @@ class BedrockMantleResponsesAPIConfig(OpenAIResponsesAPIConfig):
if base.endswith(suffix):
base = base[: -len(suffix)]
break
# For the standard Mantle host (including the default-region base that
# responses/main.py auto-injects into litellm_params.api_base), pin to the
# single resolved region so aws_region_name wins; preserve custom proxy hosts.
if _MANTLE_HOST_RE.match(base):
base = f"https://bedrock-mantle.{region}.api.aws"
return f"{base}/openai/v1/responses"
def validate_environment(
@@ -66,12 +105,8 @@ class BedrockMantleResponsesAPIConfig(OpenAIResponsesAPIConfig):
or get_secret_str("BEDROCK_MANTLE_API_KEY")
or get_secret_str("AWS_BEARER_TOKEN_BEDROCK")
)
if not api_key:
raise ValueError(
"Bedrock Mantle API key is required. Set BEDROCK_MANTLE_API_KEY "
"(or AWS_BEARER_TOKEN_BEDROCK) or pass api_key."
)
headers["Authorization"] = f"Bearer {api_key}"
if api_key:
headers["Authorization"] = f"Bearer {api_key}"
return headers
def supports_native_file_search(self) -> bool:
@@ -79,3 +114,58 @@ class BedrockMantleResponsesAPIConfig(OpenAIResponsesAPIConfig):
def supports_native_websocket(self) -> bool:
return False
def sign_request(
self,
headers: dict,
optional_params: dict,
request_data: dict,
api_base: str,
api_key: Optional[str] = None,
model: Optional[str] = None,
stream: Optional[bool] = None,
fake_stream: Optional[bool] = None,
) -> Tuple[dict, Optional[bytes]]:
bearer = (
api_key
or get_secret_str("BEDROCK_MANTLE_API_KEY")
or get_secret_str("AWS_BEARER_TOKEN_BEDROCK")
)
if not bearer:
# SigV4 path. Pin the credential-scope region to the region of the actual
# signing URL (api_base, already region-resolved by get_complete_url) so the
# SigV4 scope and the URL host can never disagree. Resolve from api_base first,
# then fall back to the regular precedence. Also drop any caller Authorization
# so _sign_request's restore-original-Authorization step cannot override the
# SigV4 header.
optional_params = {
**optional_params,
"aws_region_name": self._resolve_region(
{**optional_params, "api_base": api_base}
),
}
headers = {k: v for k, v in headers.items() if k.lower() != "authorization"}
try:
return self._aws_signer._sign_request(
service_name="bedrock",
headers=headers,
optional_params=optional_params,
request_data=request_data,
api_base=api_base,
api_key=bearer,
model=model,
stream=stream,
fake_stream=fake_stream,
)
except (
NoCredentialsError,
PartialCredentialsError,
ProfileNotFound,
CredentialRetrievalError,
) as e:
raise ValueError(
"Bedrock Mantle auth failed: no Bearer token and no usable AWS "
"credentials. Set BEDROCK_MANTLE_API_KEY (or AWS_BEARER_TOKEN_BEDROCK) "
"or pass api_key for Bearer auth, or provide AWS credentials "
"(IAM role / access key / profile / web identity) for SigV4."
) from e
litellm/llms/custom_httpx/llm_http_handler.py
+79 -26
View File
@@ -2315,6 +2315,31 @@ class BaseLLMHTTPHandler:
# but never included in the outbound provider payload.
request_context["litellm_params"] = dict(litellm_params)
is_stream_request = bool(stream)
if is_stream_request and fake_stream is True:
stream, data = self._prepare_fake_stream_request(
stream=stream,
data=data,
fake_stream=fake_stream,
)
# Sign after the body is final (post-transform/normalize/extra_body and post
# fake-stream prep) so signed bytes match what we send. No-op for providers
# that inherit the default sign_request.
headers, signed_body = responses_api_provider_config.sign_request(
headers=headers,
optional_params=dict(litellm_params),
request_data=data,
api_base=api_base,
api_key=litellm_params.api_key,
model=model,
stream=stream,
fake_stream=fake_stream,
)
body_kwargs: Dict[str, Any] = (
{"data": signed_body} if signed_body is not None else {"json": data}
)
## LOGGING
logging_obj.pre_call(
input=input,
@@ -2327,22 +2352,14 @@ class BaseLLMHTTPHandler:
)
try:
if stream:
# For streaming, use stream=True in the request
if fake_stream is True:
stream, data = self._prepare_fake_stream_request(
stream=stream,
data=data,
fake_stream=fake_stream,
)
if is_stream_request:
response = sync_httpx_client.post(
url=api_base,
headers=headers,
json=data,
timeout=timeout
or float(response_api_optional_request_params.get("timeout", 0)),
stream=stream,
**body_kwargs,
)
if fake_stream is True:
return MockResponsesAPIStreamingIterator(
@@ -2367,13 +2384,12 @@ class BaseLLMHTTPHandler:
call_type=CallTypes.responses.value,
)
else:
# For non-streaming requests
response = sync_httpx_client.post(
url=api_base,
headers=headers,
json=data,
timeout=timeout
or float(response_api_optional_request_params.get("timeout", 0)),
**body_kwargs,
)
except Exception as e:
raise self._handle_error(
@@ -2461,6 +2477,28 @@ class BaseLLMHTTPHandler:
# but never included in the outbound provider payload.
request_context["litellm_params"] = dict(litellm_params)
is_stream_request = bool(stream)
if is_stream_request and fake_stream is True:
stream, data = self._prepare_fake_stream_request(
stream=stream,
data=data,
fake_stream=fake_stream,
)
headers, signed_body = responses_api_provider_config.sign_request(
headers=headers,
optional_params=dict(litellm_params),
request_data=data,
api_base=api_base,
api_key=litellm_params.api_key,
model=model,
stream=stream,
fake_stream=fake_stream,
)
body_kwargs: Dict[str, Any] = (
{"data": signed_body} if signed_body is not None else {"json": data}
)
## LOGGING
logging_obj.pre_call(
input=input,
@@ -2473,22 +2511,14 @@ class BaseLLMHTTPHandler:
)
try:
if stream:
# For streaming, we need to use stream=True in the request
if fake_stream is True:
stream, data = self._prepare_fake_stream_request(
stream=stream,
data=data,
fake_stream=fake_stream,
)
if is_stream_request:
response = await async_httpx_client.post(
url=api_base,
headers=headers,
json=data,
timeout=timeout
or float(response_api_optional_request_params.get("timeout", 0)),
stream=stream,
**body_kwargs,
)
if fake_stream is True:
@@ -2515,13 +2545,12 @@ class BaseLLMHTTPHandler:
call_type=CallTypes.responses.value,
)
else:
# For non-streaming, proceed as before
response = await async_httpx_client.post(
url=api_base,
headers=headers,
json=data,
timeout=timeout
or float(response_api_optional_request_params.get("timeout", 0)),
**body_kwargs,
)
except Exception as e:
@@ -3998,6 +4027,18 @@ class BaseLLMHTTPHandler:
)
data = BaseResponsesAPIConfig.normalize_responses_api_request_dict(data)
headers, signed_body = responses_api_provider_config.sign_request(
headers=headers,
optional_params=dict(litellm_params),
request_data=data,
api_base=url,
api_key=litellm_params.api_key,
model=model,
)
body_kwargs: Dict[str, Any] = (
{"data": signed_body} if signed_body is not None else {"json": data}
)
## LOGGING
logging_obj.pre_call(
input=input,
@@ -4011,7 +4052,7 @@ class BaseLLMHTTPHandler:
try:
response = sync_httpx_client.post(
url=url, headers=headers, json=data, timeout=timeout
url=url, headers=headers, timeout=timeout, **body_kwargs
)
except Exception as e:
@@ -4081,6 +4122,18 @@ class BaseLLMHTTPHandler:
)
data = BaseResponsesAPIConfig.normalize_responses_api_request_dict(data)
headers, signed_body = responses_api_provider_config.sign_request(
headers=headers,
optional_params=dict(litellm_params),
request_data=data,
api_base=url,
api_key=litellm_params.api_key,
model=model,
)
body_kwargs: Dict[str, Any] = (
{"data": signed_body} if signed_body is not None else {"json": data}
)
## LOGGING
logging_obj.pre_call(
input=input,
@@ -4094,7 +4147,7 @@ class BaseLLMHTTPHandler:
try:
response = await async_httpx_client.post(
url=url, headers=headers, json=data, timeout=timeout
url=url, headers=headers, timeout=timeout, **body_kwargs
)
except Exception as e:
tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py
+391 -7
View File
@@ -12,6 +12,11 @@ import sys
sys.path.insert(0, os.path.abspath("../../../../.."))
import pytest
from botocore.exceptions import (
ConnectTimeoutError,
PartialCredentialsError,
ProfileNotFound,
)
import litellm
from litellm.llms.bedrock_mantle.responses.transformation import (
@@ -114,16 +119,15 @@ class TestBedrockMantleResponsesAuth:
)
assert headers["Authorization"] == "Bearer bearer-key"
def test_missing_key_raises(self, monkeypatch):
def test_missing_bearer_does_not_raise_in_validate_environment(self, monkeypatch):
# SigV4 may still apply, so validate_environment must defer instead of raising.
monkeypatch.delenv("BEDROCK_MANTLE_API_KEY", raising=False)
monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False)
cfg = BedrockMantleResponsesAPIConfig()
with pytest.raises(ValueError, match="Bedrock Mantle API key"):
cfg.validate_environment(
headers={},
model="openai.gpt-5.5",
litellm_params=GenericLiteLLMParams(),
)
headers = cfg.validate_environment(
headers={}, model="openai.gpt-5.5", litellm_params=GenericLiteLLMParams()
)
assert "Authorization" not in headers
def test_custom_llm_provider(self):
cfg = BedrockMantleResponsesAPIConfig()
@@ -261,6 +265,386 @@ def local_cost_map(monkeypatch):
litellm.get_model_info.cache_clear()
class TestBedrockMantleResponsesSigV4:
def test_bearer_short_circuits_without_credentials(self, monkeypatch):
from unittest.mock import MagicMock
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False)
monkeypatch.delenv("BEDROCK_MANTLE_API_KEY", raising=False)
signer = BaseAWSLLM()
signer.get_credentials = MagicMock(
side_effect=AssertionError("get_credentials must not run for bearer auth")
)
cfg = BedrockMantleResponsesAPIConfig(aws_signer=signer)
headers, signed_body = cfg.sign_request(
headers={},
optional_params={},
request_data={"input": "hi"},
api_base="https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses",
api_key="bearer-from-config",
)
assert headers["Authorization"] == "Bearer bearer-from-config"
assert signed_body == b'{"input": "hi"}'
signer.get_credentials.assert_not_called()
def test_bearer_resolved_from_mantle_env_key(self, monkeypatch):
from unittest.mock import MagicMock
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False)
monkeypatch.setenv("BEDROCK_MANTLE_API_KEY", "env-bearer")
signer = BaseAWSLLM()
signer.get_credentials = MagicMock(
side_effect=AssertionError("get_credentials must not run for bearer auth")
)
cfg = BedrockMantleResponsesAPIConfig(aws_signer=signer)
headers, _ = cfg.sign_request(
headers={},
optional_params={},
request_data={"input": "hi"},
api_base="https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses",
api_key=None,
)
assert headers["Authorization"] == "Bearer env-bearer"
def test_bearer_arg_takes_priority_over_mantle_env_key(self, monkeypatch):
# The passed api_key (e.g. litellm_params.api_key) must win over the env
# bearer; a reordered precedence chain would silently use the wrong token.
from unittest.mock import MagicMock
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False)
monkeypatch.setenv("BEDROCK_MANTLE_API_KEY", "env-bearer")
signer = BaseAWSLLM()
signer.get_credentials = MagicMock(
side_effect=AssertionError("get_credentials must not run for bearer auth")
)
cfg = BedrockMantleResponsesAPIConfig(aws_signer=signer)
headers, _ = cfg.sign_request(
headers={},
optional_params={},
request_data={"input": "hi"},
api_base="https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses",
api_key="arg-bearer",
)
assert headers["Authorization"] == "Bearer arg-bearer"
signer.get_credentials.assert_not_called()
def test_access_key_produces_sigv4_headers(self, monkeypatch):
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False)
monkeypatch.delenv("BEDROCK_MANTLE_API_KEY", raising=False)
cfg = BedrockMantleResponsesAPIConfig(aws_signer=BaseAWSLLM())
headers, signed_body = cfg.sign_request(
headers={},
optional_params={
"aws_access_key_id": "AKIAEXAMPLE",
"aws_secret_access_key": "c2VjcmV0LXRlc3Qtc2VjcmV0LXRlc3Qtc2VjcmV0",
"aws_session_token": "session-token-test",
"aws_region_name": "us-east-2",
},
request_data={"input": "hi"},
api_base="https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses",
api_key=None,
)
assert headers["Authorization"].startswith("AWS4-HMAC-SHA256")
assert "Credential=AKIAEXAMPLE/" in headers["Authorization"]
assert "/us-east-2/bedrock/aws4_request" in headers["Authorization"]
assert "X-Amz-Date" in headers
assert headers["X-Amz-Security-Token"] == "session-token-test"
assert signed_body == b'{"input": "hi"}'
def test_assume_role_path_produces_sigv4_headers(self, monkeypatch):
from unittest.mock import MagicMock
from botocore.credentials import Credentials
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False)
monkeypatch.delenv("BEDROCK_MANTLE_API_KEY", raising=False)
signer = BaseAWSLLM()
signer.get_credentials = MagicMock(
return_value=Credentials(
access_key="ASIAEXAMPLE",
secret_key="YXNzdW1lZC1yb2xlLXNlY3JldC1hc3N1bWVk",
token="assumed-session-token",
)
)
cfg = BedrockMantleResponsesAPIConfig(aws_signer=signer)
headers, _ = cfg.sign_request(
headers={},
optional_params={
"aws_role_name": "arn:aws:iam::000000000000:role/test-role",
"aws_session_name": "litellm-test",
"aws_region_name": "us-east-2",
},
request_data={"input": "hi"},
api_base="https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses",
api_key=None,
)
signer.get_credentials.assert_called_once()
call = signer.get_credentials.call_args.kwargs
assert call["aws_role_name"] == "arn:aws:iam::000000000000:role/test-role"
assert call["aws_session_name"] == "litellm-test"
assert headers["Authorization"].startswith("AWS4-HMAC-SHA256")
assert "/us-east-2/bedrock/aws4_request" in headers["Authorization"]
def test_signed_body_matches_final_data_after_normalize(self, monkeypatch):
"""Core regression: the signed bytes must equal the bytes actually sent.
Sign the *final* data dict and assert the returned signed_body decodes to
exactly that dict, so a later change to the data would break the SigV4 hash.
"""
import json
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False)
monkeypatch.delenv("BEDROCK_MANTLE_API_KEY", raising=False)
final_data = {"model": "openai.gpt-5.5", "input": "hi", "max_output_tokens": 16}
cfg = BedrockMantleResponsesAPIConfig(aws_signer=BaseAWSLLM())
_, signed_body = cfg.sign_request(
headers={},
optional_params={
"aws_access_key_id": "AKIAEXAMPLE",
"aws_secret_access_key": "c2VjcmV0LXRlc3Qtc2VjcmV0LXRlc3Qtc2VjcmV0",
"aws_region_name": "us-east-2",
},
request_data=final_data,
api_base="https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses",
api_key=None,
)
assert signed_body is not None
assert json.loads(signed_body) == final_data
def test_region_comes_from_optional_params(self, monkeypatch):
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False)
monkeypatch.delenv("BEDROCK_MANTLE_API_KEY", raising=False)
monkeypatch.delenv("AWS_REGION", raising=False)
monkeypatch.delenv("AWS_REGION_NAME", raising=False)
cfg = BedrockMantleResponsesAPIConfig(aws_signer=BaseAWSLLM())
headers, _ = cfg.sign_request(
headers={},
optional_params={
"aws_access_key_id": "AKIAEXAMPLE",
"aws_secret_access_key": "c2VjcmV0LXRlc3Qtc2VjcmV0LXRlc3Qtc2VjcmV0",
"aws_region_name": "eu-west-1",
},
request_data={"input": "hi"},
api_base="https://bedrock-mantle.eu-west-1.api.aws/openai/v1/responses",
api_key=None,
)
assert "/eu-west-1/bedrock/aws4_request" in headers["Authorization"]
def test_url_region_and_sigv4_region_agree_from_litellm_params(self, monkeypatch):
"""Adversarial-review regression: a caller-supplied aws_region_name (no region
env set) must shape BOTH the URL host and the SigV4 credential scope, or the
request is signed for one region and sent to another -> 401.
"""
monkeypatch.delenv("BEDROCK_MANTLE_REGION", raising=False)
monkeypatch.delenv("BEDROCK_MANTLE_API_BASE", raising=False)
monkeypatch.delenv("AWS_REGION", raising=False)
monkeypatch.delenv("AWS_REGION_NAME", raising=False)
monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False)
monkeypatch.delenv("BEDROCK_MANTLE_API_KEY", raising=False)
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
params = {
"aws_region_name": "ap-southeast-2",
"aws_access_key_id": "AKIAEXAMPLE",
"aws_secret_access_key": "c2VjcmV0LXRlc3Qtc2VjcmV0LXRlc3Qtc2VjcmV0",
}
cfg = BedrockMantleResponsesAPIConfig(aws_signer=BaseAWSLLM())
url = cfg.get_complete_url(api_base=None, litellm_params=params)
assert (
url == "https://bedrock-mantle.ap-southeast-2.api.aws/openai/v1/responses"
)
headers, _ = cfg.sign_request(
headers={},
optional_params=params,
request_data={"input": "hi"},
api_base=url,
api_key=None,
)
assert "/ap-southeast-2/bedrock/aws4_request" in headers["Authorization"]
def test_injected_default_region_base_does_not_override_aws_region_name(
self, monkeypatch
):
"""2nd-round adversarial regression: responses/main.py auto-injects
litellm_params.api_base = https://bedrock-mantle.<DEFAULT>.api.aws/v1 (default
region, ignoring aws_region_name). The config must still pin BOTH the URL host
and the SigV4 scope to aws_region_name, or the IAM deployment 401s. A naive
'resolve region only when api_base is None' fix would fail this test.
"""
monkeypatch.delenv("BEDROCK_MANTLE_REGION", raising=False)
monkeypatch.delenv("BEDROCK_MANTLE_API_BASE", raising=False)
monkeypatch.delenv("AWS_REGION", raising=False)
monkeypatch.delenv("AWS_REGION_NAME", raising=False)
monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False)
monkeypatch.delenv("BEDROCK_MANTLE_API_KEY", raising=False)
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
injected_base = "https://bedrock-mantle.us-east-1.api.aws/v1" # default region
params = {
"aws_region_name": "us-east-2", # what the caller actually wants
"api_base": injected_base,
"aws_access_key_id": "AKIAEXAMPLE",
"aws_secret_access_key": "c2VjcmV0LXRlc3Qtc2VjcmV0LXRlc3Qtc2VjcmV0",
}
cfg = BedrockMantleResponsesAPIConfig(aws_signer=BaseAWSLLM())
url = cfg.get_complete_url(api_base=injected_base, litellm_params=params)
assert url == "https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses"
headers, _ = cfg.sign_request(
headers={},
optional_params=params,
request_data={"input": "hi"},
api_base=url,
api_key=None,
)
assert "/us-east-2/bedrock/aws4_request" in headers["Authorization"]
assert "us-east-1" not in headers["Authorization"]
def test_custom_proxy_host_is_preserved(self, monkeypatch):
"""A genuinely custom (non-Mantle) api_base host must be preserved, not rewritten
to a bedrock-mantle host. Only standard Mantle hosts are region-pinned.
"""
monkeypatch.delenv("BEDROCK_MANTLE_API_BASE", raising=False)
cfg = BedrockMantleResponsesAPIConfig()
url = cfg.get_complete_url(
api_base="https://mantle-proxy.internal.example/openai/v1",
litellm_params={"aws_region_name": "us-east-2"},
)
assert url == "https://mantle-proxy.internal.example/openai/v1/responses"
def test_caller_authorization_does_not_override_sigv4(self, monkeypatch):
"""Adversarial-review regression: a caller-supplied Authorization header (e.g.
from extra_headers, surviving the relaxed validate_environment) must not clobber
the SigV4 Authorization that _sign_request would otherwise restore.
"""
monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False)
monkeypatch.delenv("BEDROCK_MANTLE_API_KEY", raising=False)
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
cfg = BedrockMantleResponsesAPIConfig(aws_signer=BaseAWSLLM())
headers, _ = cfg.sign_request(
headers={"Authorization": "Bearer stale-caller-token"},
optional_params={
"aws_access_key_id": "AKIAEXAMPLE",
"aws_secret_access_key": "c2VjcmV0LXRlc3Qtc2VjcmV0LXRlc3Qtc2VjcmV0",
"aws_region_name": "us-east-2",
},
request_data={"input": "hi"},
api_base="https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses",
api_key=None,
)
assert headers["Authorization"].startswith("AWS4-HMAC-SHA256")
assert "Bearer stale-caller-token" not in headers["Authorization"]
def test_no_bearer_and_no_credentials_raises_both_paths(self, monkeypatch):
from unittest.mock import MagicMock
from botocore.exceptions import NoCredentialsError
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
monkeypatch.delenv("BEDROCK_MANTLE_API_KEY", raising=False)
monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False)
signer = BaseAWSLLM()
signer.get_credentials = MagicMock(side_effect=NoCredentialsError())
cfg = BedrockMantleResponsesAPIConfig(aws_signer=signer)
with pytest.raises(ValueError) as exc:
cfg.sign_request(
headers={},
optional_params={"aws_region_name": "us-east-2"},
request_data={"input": "hi"},
api_base="https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses",
api_key=None,
)
msg = str(exc.value)
assert "Bearer" in msg
assert "SigV4" in msg or "IAM" in msg
@pytest.mark.parametrize(
"cred_error",
[
PartialCredentialsError(provider="env", cred_var="aws_secret_access_key"),
ProfileNotFound(profile="missing-profile"),
],
)
def test_partial_credentials_raises_both_paths(self, monkeypatch, cred_error):
from unittest.mock import MagicMock
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
monkeypatch.delenv("BEDROCK_MANTLE_API_KEY", raising=False)
monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False)
signer = BaseAWSLLM()
signer.get_credentials = MagicMock(side_effect=cred_error)
cfg = BedrockMantleResponsesAPIConfig(aws_signer=signer)
with pytest.raises(ValueError) as exc:
cfg.sign_request(
headers={},
optional_params={"aws_region_name": "us-east-2"},
request_data={"input": "hi"},
api_base="https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses",
api_key=None,
)
msg = str(exc.value)
assert "Bearer" in msg
assert "SigV4" in msg or "IAM" in msg
def test_sts_transport_error_is_not_masked_as_credentials(self, monkeypatch):
# An AssumeRole / web-identity flow hits STS over the network, so a transient
# connection error must surface as itself, not be rewritten into the
# "no usable AWS credentials" message that would send the user to fix the
# wrong thing.
from unittest.mock import MagicMock
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
monkeypatch.delenv("BEDROCK_MANTLE_API_KEY", raising=False)
monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False)
signer = BaseAWSLLM()
signer.get_credentials = MagicMock(
side_effect=ConnectTimeoutError(
endpoint_url="https://sts.us-east-2.amazonaws.com"
)
)
cfg = BedrockMantleResponsesAPIConfig(aws_signer=signer)
with pytest.raises(ConnectTimeoutError):
cfg.sign_request(
headers={},
optional_params={
"aws_role_name": "arn:aws:iam::000000000000:role/test-role",
"aws_region_name": "us-east-2",
},
request_data={"input": "hi"},
api_base="https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses",
api_key=None,
)
class TestBedrockMantleResponsesPricing:
def test_gpt_5_5_pricing_and_mode(self, local_cost_map):
info = litellm.get_model_info("bedrock_mantle/openai.gpt-5.5")
tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py
+238
View File
@@ -629,3 +629,241 @@ async def test_anthropic_post_retry_reserializes_mutated_body():
assert first_sent == prebuilt # attempt 0 used prebuilt
assert second_sent == _json.dumps(request_body) # attempt 1 re-serialized
assert "MUTATED" in second_sent # ... the mutated body
def test_base_responses_config_sign_request_is_noop_by_default():
"""Default responses sign_request must be a no-op: unchanged headers, no signed body.
Guards the 15 existing responses providers from accidental signing when the
handler starts calling sign_request.
"""
from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig
cfg = OpenAIResponsesAPIConfig()
headers = {"Authorization": "Bearer sk-existing"}
out_headers, signed_body = cfg.sign_request(
headers=headers,
optional_params={},
request_data={"input": "hi"},
api_base="https://api.openai.com/v1/responses",
)
assert out_headers == {"Authorization": "Bearer sk-existing"}
assert signed_body is None
def _make_responses_handler_call(signed_body):
"""Drive BaseLLMHTTPHandler.response_api_handler with a fully mocked provider
config + sync client, returning the kwargs the client.post was called with.
signed_body=None simulates a no-op (non-signing) provider; bytes simulates a
signing provider (e.g. Bedrock Mantle).
"""
from unittest.mock import MagicMock
from litellm.llms.custom_httpx.http_handler import HTTPHandler
from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler
from litellm.types.router import GenericLiteLLMParams
provider_config = MagicMock()
provider_config.validate_environment.return_value = {}
provider_config.get_complete_url.return_value = (
"https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses"
)
provider_config.transform_responses_api_request.return_value = {"input": "hi"}
provider_config.should_fake_stream.return_value = False
provider_config.sign_request.return_value = ({"X-Signed": "1"}, signed_body)
mock_client = MagicMock(spec=HTTPHandler)
mock_client.post.return_value = MagicMock()
handler = BaseLLMHTTPHandler()
handler.response_api_handler(
model="openai.gpt-5.5",
input="hi",
responses_api_provider_config=provider_config,
response_api_optional_request_params={},
custom_llm_provider="bedrock_mantle",
litellm_params=GenericLiteLLMParams(aws_region_name="us-east-2"),
logging_obj=MagicMock(),
client=mock_client,
_is_async=False,
)
return mock_client.post.call_args.kwargs
def test_responses_handler_sends_json_when_not_signed():
"""No-op provider (signed_body is None) -> handler posts json=data, no data= bytes."""
kwargs = _make_responses_handler_call(signed_body=None)
assert kwargs.get("json") == {"input": "hi"}
assert "data" not in kwargs
def test_responses_handler_sends_signed_bytes_when_signed():
"""Signing provider -> handler posts the exact signed bytes via data=, not json=."""
kwargs = _make_responses_handler_call(signed_body=b'{"input": "hi"}')
assert kwargs.get("data") == b'{"input": "hi"}'
assert "json" not in kwargs
assert kwargs["headers"] == {"X-Signed": "1"}
def test_responses_handler_signs_after_fake_stream_prep_strips_stream():
"""Fake-stream signing-order invariant: the bytes SIGNED must equal the bytes SENT.
In the streaming + fake-stream path the handler first runs
_prepare_fake_stream_request, which pops "stream" out of the body, and only
then calls sign_request. If signing ran before that pop, the signed body
would still carry "stream" while the body sent over the wire would not,
producing a SigV4 payload-hash mismatch (401) for a real Mantle deployment.
We snapshot request_data at sign time and assert "stream" is already gone.
"""
from unittest.mock import MagicMock
from litellm.llms.custom_httpx.http_handler import HTTPHandler
from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler
from litellm.types.llms.openai import ResponsesAPIResponse
from litellm.types.router import GenericLiteLLMParams
provider_config = MagicMock()
provider_config.validate_environment.return_value = {}
provider_config.get_complete_url.return_value = (
"https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses"
)
provider_config.transform_responses_api_request.return_value = {
"input": "hi",
"stream": True,
}
provider_config.should_fake_stream.return_value = True
provider_config.transform_response_api_response.return_value = ResponsesAPIResponse(
id="resp_1",
created_at=0,
output=[],
status="completed",
model="openai.gpt-5.5",
)
captured = {}
def _capture_sign(**kwargs):
captured["request_data"] = dict(kwargs["request_data"])
return ({"X-Signed": "1"}, b'{"input": "hi"}')
provider_config.sign_request.side_effect = _capture_sign
mock_client = MagicMock(spec=HTTPHandler)
mock_client.post.return_value = MagicMock()
handler = BaseLLMHTTPHandler()
handler.response_api_handler(
model="openai.gpt-5.5",
input="hi",
responses_api_provider_config=provider_config,
response_api_optional_request_params={"stream": True},
custom_llm_provider="bedrock_mantle",
litellm_params=GenericLiteLLMParams(aws_region_name="us-east-2"),
logging_obj=MagicMock(),
client=mock_client,
_is_async=False,
fake_stream=True,
)
assert "stream" not in captured["request_data"]
assert "input" in captured["request_data"]
post_kwargs = mock_client.post.call_args.kwargs
assert post_kwargs.get("data") == b'{"input": "hi"}'
assert "json" not in post_kwargs
assert "stream" in post_kwargs
def _make_compact_handler_call(signed_body, is_async):
"""Drive (async_)compact_response_api_handler with a fully mocked provider config
+ client, returning the kwargs the client.post was called with.
signed_body=None simulates a no-op (non-signing) provider; bytes simulates a
signing provider (e.g. Bedrock Mantle SigV4 / bearer).
"""
from unittest.mock import MagicMock
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler
from litellm.types.router import GenericLiteLLMParams
compact_url = "https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses/compact"
provider_config = MagicMock()
provider_config.validate_environment.return_value = {}
provider_config.get_complete_url.return_value = (
"https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses"
)
provider_config.transform_compact_response_api_request.return_value = (
compact_url,
{"model": "openai.gpt-5.5", "input": "hi"},
)
provider_config.sign_request.return_value = ({"X-Signed": "1"}, signed_body)
provider_config.transform_compact_response_api_response.return_value = "ok"
spec = AsyncHTTPHandler if is_async else HTTPHandler
mock_client = MagicMock(spec=spec)
if is_async:
mock_client.post = AsyncMock(return_value=MagicMock())
else:
mock_client.post.return_value = MagicMock()
handler = BaseLLMHTTPHandler()
result = handler.compact_response_api_handler(
model="openai.gpt-5.5",
input="hi",
responses_api_provider_config=provider_config,
response_api_optional_request_params={},
custom_llm_provider="bedrock_mantle",
litellm_params=GenericLiteLLMParams(aws_region_name="us-east-2"),
logging_obj=MagicMock(),
client=mock_client,
_is_async=is_async,
)
if is_async:
asyncio.run(result)
return provider_config, mock_client.post.call_args.kwargs
def test_compact_handler_sends_json_when_not_signed():
"""No-op provider on compact (signed_body is None) -> posts json=data, no data= bytes."""
provider_config, kwargs = _make_compact_handler_call(
signed_body=None, is_async=False
)
provider_config.sign_request.assert_called_once()
assert kwargs.get("json") == {"model": "openai.gpt-5.5", "input": "hi"}
assert "data" not in kwargs
def test_compact_handler_sends_signed_bytes_when_signed():
"""Signing provider on compact -> posts the signed bytes via data=, not json=.
Regression for the adversarial-review finding that /responses/compact bypassed
the SigV4 signing hook, so IAM-only Mantle callers sent unsigned bodies.
"""
provider_config, kwargs = _make_compact_handler_call(
signed_body=b'{"model": "openai.gpt-5.5", "input": "hi"}', is_async=False
)
assert kwargs.get("data") == b'{"model": "openai.gpt-5.5", "input": "hi"}'
assert "json" not in kwargs
assert kwargs["headers"] == {"X-Signed": "1"}
# signing must use the compact endpoint as api_base, not the create URL
assert provider_config.sign_request.call_args.kwargs["api_base"].endswith(
"/openai/v1/responses/compact"
)
def test_async_compact_handler_sends_signed_bytes_when_signed():
"""Async compact must sign identically to sync (same omission in the async twin)."""
provider_config, kwargs = _make_compact_handler_call(
signed_body=b'{"model": "openai.gpt-5.5", "input": "hi"}', is_async=True
)
assert kwargs.get("data") == b'{"model": "openai.gpt-5.5", "input": "hi"}'
assert "json" not in kwargs
assert kwargs["headers"] == {"X-Signed": "1"}
def test_async_compact_handler_sends_json_when_not_signed():
"""Async no-op provider on compact -> posts json=data, no data= bytes."""
_provider_config, kwargs = _make_compact_handler_call(
signed_body=None, is_async=True
)
assert kwargs.get("json") == {"model": "openai.gpt-5.5", "input": "hi"}
assert "data" not in kwargs