Merge pull request #36160 from daleselaji-dev/codex/bedrock-s3-credentials-36155

fix(bedrock): use deployment credentials for AWS requests
This commit is contained in:
Mateo Wang 2026-08-10 22:46:37 -07:00 committed by GitHub
commit c4656d806c
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
11 changed files with 435 additions and 14 deletions

View file

@ -28,7 +28,11 @@ from litellm.types.llms.openai import (
from litellm.types.utils import LiteLLMBatch, LlmProviders
from ..base_aws_llm import BaseAWSLLM
from ..common_utils import CommonBatchFilesUtils, resolve_s3_encryption_key_id
from ..common_utils import (
CommonBatchFilesUtils,
merge_bedrock_aws_request_params,
resolve_s3_encryption_key_id,
)
# Bedrock batch input files are uploaded as
# s3://bucket/litellm-bedrock-files-{model, ":" -> "-"}-{uuid4}.jsonl (see
@ -130,7 +134,8 @@ class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig):
Get the complete URL for Bedrock batch creation.
Bedrock batch jobs are created via the model invocation job API.
"""
aws_region_name: Final = self._get_aws_region_name(optional_params, model)
request_params: Final = merge_bedrock_aws_request_params(litellm_params, optional_params)
aws_region_name: Final = self._get_aws_region_name(request_params, model)
# Bedrock model invocation job endpoint
# Format: https://bedrock.{region}.amazonaws.com/model-invocation-job
@ -232,14 +237,15 @@ class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig):
# For Bedrock, we need to return a pre-signed request with AWS auth headers
# Use common utility for AWS signing
request_params: Final = merge_bedrock_aws_request_params(litellm_params, optional_params)
endpoint_url: Final = (
f"https://bedrock.{self._get_aws_region_name(optional_params, model)}.amazonaws.com/model-invocation-job"
f"https://bedrock.{self._get_aws_region_name(request_params, model)}.amazonaws.com/model-invocation-job"
)
signed_headers, signed_data = self.common_utils.sign_aws_request(
service_name="bedrock",
data=bedrock_request,
endpoint_url=endpoint_url,
optional_params=optional_params,
optional_params=request_params,
method="POST",
)
@ -387,11 +393,12 @@ class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig):
endpoint_url: Final = f"https://bedrock.{region}.amazonaws.com/model-invocation-job/{encoded_arn}"
# Use common utility for AWS signing
request_params: Final = merge_bedrock_aws_request_params(litellm_params, optional_params)
signed_headers, _ = self.common_utils.sign_aws_request(
service_name="bedrock",
data={}, # GET request has no body
endpoint_url=endpoint_url,
optional_params=optional_params,
optional_params=request_params,
method="GET",
)

View file

@ -36,6 +36,44 @@ class BedrockError(BaseLLMException):
pass
_BEDROCK_AWS_AUTH_PARAMETER_KEYS: Final[tuple[str, ...]] = (
"aws_access_key_id",
"aws_secret_access_key",
"aws_session_token",
"aws_region_name",
"aws_session_name",
"aws_profile_name",
"aws_role_name",
"aws_web_identity_token",
"aws_sts_endpoint",
"aws_external_id",
)
def merge_bedrock_aws_request_params(
litellm_params: Mapping[str, Any],
optional_params: Mapping[str, Any],
) -> dict[str, Any]:
"""Merge deployment and request parameters without allowing auth escalation.
Deployment configuration is authoritative for AWS authentication. When a
deployment supplies static credentials, caller-supplied profile/role/token
selectors must not redirect signing to another identity available on the
server. Requests may still provide AWS credentials when the deployment has
no static credentials configured.
"""
request_params: Final = {**optional_params, **litellm_params} # mutable-ok: AWS helpers require a plain dict
has_static_deployment_credentials: Final = all(
isinstance(litellm_params.get(key), str) and bool(litellm_params.get(key))
for key in ("aws_access_key_id", "aws_secret_access_key", "aws_region_name")
)
if has_static_deployment_credentials:
for key in _BEDROCK_AWS_AUTH_PARAMETER_KEYS:
if key not in litellm_params:
request_params.pop(key, None)
return request_params
# Lazy import cache to avoid circular imports and performance impact
_get_model_info = None

View file

@ -54,7 +54,7 @@ from litellm.types.utils import ExtractedFileData, LlmProviders, SpecialEnums
from litellm.utils import get_llm_provider
from ..base_aws_llm import BaseAWSLLM
from ..common_utils import BedrockError, resolve_s3_encryption_key_id
from ..common_utils import BedrockError, merge_bedrock_aws_request_params, resolve_s3_encryption_key_id
# litellm_params key used to hand the SigV4-signed GET headers from
# `transform_file_content_request` to `validate_environment` (the only hook
@ -285,6 +285,7 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig):
"""
Get the complete S3 URL for the file upload request
"""
request_params: Final = merge_bedrock_aws_request_params(litellm_params, optional_params)
bucket_name = litellm_params.get("s3_bucket_name") or os.getenv("AWS_S3_BUCKET_NAME")
if not bucket_name:
raise ValueError(
@ -293,7 +294,7 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig):
bucket_name, object_prefix = split_configured_cloud_bucket_name(bucket_name)
s3_region_name: Final = litellm_params.get("s3_region_name") or optional_params.get("s3_region_name")
aws_region_name: Final = s3_region_name or self._get_aws_region_name(optional_params, model)
aws_region_name: Final = s3_region_name or self._get_aws_region_name(request_params, model)
file_data: Final = data.get("file")
purpose: Final = data.get("purpose")
@ -309,7 +310,7 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig):
# S3 endpoint URL format
s3_endpoint_url: Final = (
optional_params.get("s3_endpoint_url") or f"https://s3.{aws_region_name}.amazonaws.com"
request_params.get("s3_endpoint_url") or f"https://s3.{aws_region_name}.amazonaws.com"
).rstrip("/")
return f"{s3_endpoint_url}/{bucket_name}/{encoded_object_name}"
@ -843,20 +844,23 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig):
)
# s3_region_name always wins for S3 operations (same priority as in
# get_complete_file_url above). Overwrite aws_region_name unconditionally
# so the SigV4 region matches the URL region, avoiding SignatureDoesNotMatch.
# get_complete_file_url above). Overwrite aws_region_name unconditionally,
# after the deployment-credential merge, so the SigV4 region matches the
# URL region, avoiding SignatureDoesNotMatch.
merged_params: Final = merge_bedrock_aws_request_params(litellm_params, optional_params)
s3_region_name: Final = litellm_params.get("s3_region_name") or optional_params.get("s3_region_name")
if s3_region_name:
optional_params = {**optional_params, "aws_region_name": s3_region_name}
request_params: Final = (
{**merged_params, "aws_region_name": s3_region_name} if s3_region_name else merged_params
)
# Sign the request and return a pre-signed request object
signed_headers, signed_body = self._sign_s3_request(
content=file_content,
api_base=api_base,
optional_params=optional_params,
optional_params=request_params,
s3_encryption_key_id=resolve_s3_encryption_key_id(
litellm_params=litellm_params,
optional_params=optional_params,
optional_params=request_params,
),
)

View file

@ -262,6 +262,14 @@ _BANNED_REQUEST_BODY_PARAMS: Final[tuple[str, ...]] = (
"aws_sts_endpoint",
"aws_web_identity_token",
"aws_role_name",
# Remaining AWS identity selectors. ``get_credentials`` prefers a named
# profile over the deployment's static keys, so a caller-supplied
# ``aws_profile_name`` signs Bedrock and S3 requests as any profile
# present on the proxy host; the two AssumeRole knobs are banned with it
# so the whole identity-selection family lives behind the same opt-in.
"aws_profile_name",
"aws_session_name",
"aws_external_id",
"vertex_credentials",
# Azure managed-identity / federated-auth token. The Azure provider
# transformer reads ``azure_ad_token`` (top-level or via

View file

@ -237,7 +237,14 @@ class CredentialLiteLLMParams(BaseModel):
## AWS BEDROCK / SAGEMAKER ##
aws_access_key_id: str | None = None
aws_secret_access_key: str | None = None
aws_session_token: str | None = None
aws_region_name: str | None = None
aws_session_name: str | None = None
aws_profile_name: str | None = None
aws_role_name: str | None = None
aws_web_identity_token: str | None = None
aws_sts_endpoint: str | None = None
aws_external_id: str | None = None
aws_bedrock_runtime_endpoint: str | None = None
aws_bedrock_project_id: str | None = None
s3_bucket_name: str | None = None

View file

@ -389,3 +389,170 @@ def test_bedrock_batch_with_encryption_key_in_post_request():
)
print("SUCCESS: s3_encryption_key_id properly included in AWS POST request")
def test_bedrock_file_upload_signing_uses_deployment_credentials(monkeypatch):
from litellm.llms.bedrock.files.transformation import BedrockFilesConfig
config = BedrockFilesConfig()
captured = {}
def capture_signing(**kwargs):
captured.update(kwargs)
return {}, ""
monkeypatch.setattr(config, "_sign_s3_request", capture_signing)
result = config.transform_create_file_request(
model="",
create_file_data={
"file": (
"batch.jsonl",
b'{"custom_id":"req-1","body":{"model":"bedrock/model"}}\n',
"application/jsonl",
),
"purpose": "batch",
},
optional_params={},
litellm_params={
"s3_bucket_name": "deployment-bucket",
"aws_access_key_id": "deployment-access-key",
"aws_secret_access_key": "deployment-secret",
"aws_region_name": "eu-west-1",
},
)
assert "eu-west-1" in result["url"]
assert captured["optional_params"]["aws_access_key_id"] == "deployment-access-key"
assert captured["optional_params"]["aws_secret_access_key"] == "deployment-secret"
assert captured["optional_params"]["aws_region_name"] == "eu-west-1"
def test_bedrock_batch_signing_uses_deployment_credentials(monkeypatch):
from litellm.llms.bedrock.batches.transformation import BedrockBatchesConfig
config = BedrockBatchesConfig()
captured = {}
def capture_signing(**kwargs):
captured.update(kwargs)
return {}, b"{}"
monkeypatch.setattr(config.common_utils, "sign_aws_request", capture_signing)
result = config.transform_create_batch_request(
model="us.anthropic.claude-haiku-4-5-20251001-v1:0",
create_batch_data={
"input_file_id": "s3://deployment-bucket/input.jsonl",
"completion_window": "24h",
"endpoint": "/v1/chat/completions",
},
optional_params={},
litellm_params={
"aws_access_key_id": "deployment-access-key",
"aws_secret_access_key": "deployment-secret",
"aws_region_name": "eu-west-1",
"aws_batch_role_arn": "arn:aws:iam::123456789012:role/bedrock-batch",
},
)
assert result["url"].startswith("https://bedrock.eu-west-1.amazonaws.com/")
assert captured["optional_params"]["aws_access_key_id"] == "deployment-access-key"
assert captured["optional_params"]["aws_secret_access_key"] == "deployment-secret"
assert captured["optional_params"]["aws_region_name"] == "eu-west-1"
def test_bedrock_batch_retrieval_signing_uses_deployment_credentials(monkeypatch):
from litellm.llms.bedrock.batches.transformation import BedrockBatchesConfig
config = BedrockBatchesConfig()
captured = {}
def capture_signing(**kwargs):
captured.update(kwargs)
return {}, b""
monkeypatch.setattr(config.common_utils, "sign_aws_request", capture_signing)
result = config.transform_retrieve_batch_request(
batch_id="arn:aws:bedrock:eu-west-1:123456789012:model-invocation-job/job-1",
optional_params={},
litellm_params={
"aws_access_key_id": "deployment-access-key",
"aws_secret_access_key": "deployment-secret",
"aws_region_name": "eu-west-1",
},
)
assert result["url"].startswith("https://bedrock.eu-west-1.amazonaws.com/")
assert captured["optional_params"]["aws_access_key_id"] == "deployment-access-key"
assert captured["optional_params"]["aws_secret_access_key"] == "deployment-secret"
assert captured["optional_params"]["aws_region_name"] == "eu-west-1"
def test_bedrock_deployment_credentials_block_caller_profile_override(monkeypatch):
from litellm.llms.bedrock.batches.transformation import BedrockBatchesConfig
config = BedrockBatchesConfig()
captured = {}
def capture_signing(**kwargs):
captured.update(kwargs)
return {}, b"{}"
monkeypatch.setattr(config.common_utils, "sign_aws_request", capture_signing)
config.transform_create_batch_request(
model="us.anthropic.claude-haiku-4-5-20251001-v1:0",
create_batch_data={
"input_file_id": "s3://deployment-bucket/input.jsonl",
"completion_window": "24h",
},
optional_params={"aws_profile_name": "caller-controlled-profile"},
litellm_params={
"aws_access_key_id": "deployment-access-key",
"aws_secret_access_key": "deployment-secret",
"aws_region_name": "eu-west-1",
"aws_batch_role_arn": "arn:aws:iam::123456789012:role/bedrock-batch",
},
)
assert "aws_profile_name" not in captured["optional_params"]
assert captured["optional_params"]["aws_access_key_id"] == "deployment-access-key"
def test_bedrock_file_upload_s3_region_survives_deployment_region_merge(monkeypatch):
from litellm.llms.bedrock.files.transformation import BedrockFilesConfig
config = BedrockFilesConfig()
captured = {}
def capture_signing(**kwargs):
captured.update(kwargs)
return {}, ""
monkeypatch.setattr(config, "_sign_s3_request", capture_signing)
result = config.transform_create_file_request(
model="",
create_file_data={
"file": (
"batch.jsonl",
b'{"custom_id":"req-1","body":{"model":"bedrock/model"}}\n',
"application/jsonl",
),
"purpose": "batch",
},
optional_params={},
litellm_params={
"s3_bucket_name": "deployment-bucket",
"s3_region_name": "eu-central-1",
"aws_access_key_id": "deployment-access-key",
"aws_secret_access_key": "deployment-secret",
"aws_region_name": "us-east-1",
},
)
assert "s3.eu-central-1.amazonaws.com" in result["url"]
assert captured["optional_params"]["aws_region_name"] == "eu-central-1"
assert captured["optional_params"]["aws_access_key_id"] == "deployment-access-key"

View file

@ -473,3 +473,55 @@ def test_capability_lookups_fall_back_to_base_model_when_regional_entry_lacks_fi
assert is_claude_4_5_on_bedrock(regional) is True
assert bedrock_converse_supports_parallel_tool_use_config(regional) is True
def test_merge_bedrock_aws_request_params_strips_caller_identity_when_deployment_has_static_credentials():
from litellm.llms.bedrock.common_utils import merge_bedrock_aws_request_params
merged = merge_bedrock_aws_request_params(
litellm_params={
"aws_access_key_id": "deployment-key",
"aws_secret_access_key": "deployment-secret",
"aws_region_name": "us-west-2",
"s3_bucket_name": "deployment-bucket",
},
optional_params={
"aws_access_key_id": "caller-key",
"aws_profile_name": "caller-profile",
"aws_role_name": "arn:aws:iam::123456789012:role/caller",
"aws_session_token": "caller-token",
"aws_web_identity_token": "caller-web-identity",
"timeout": 600,
},
)
assert merged["aws_access_key_id"] == "deployment-key"
assert merged["aws_secret_access_key"] == "deployment-secret"
assert merged["aws_region_name"] == "us-west-2"
assert merged["s3_bucket_name"] == "deployment-bucket"
assert merged["timeout"] == 600
for stripped in (
"aws_profile_name",
"aws_role_name",
"aws_session_token",
"aws_web_identity_token",
):
assert stripped not in merged
def test_merge_bedrock_aws_request_params_keeps_caller_credentials_without_static_deployment_credentials():
from litellm.llms.bedrock.common_utils import merge_bedrock_aws_request_params
merged = merge_bedrock_aws_request_params(
litellm_params={"aws_region_name": "us-west-2"},
optional_params={
"aws_access_key_id": "caller-key",
"aws_secret_access_key": "caller-secret",
"aws_session_token": "caller-token",
},
)
assert merged["aws_access_key_id"] == "caller-key"
assert merged["aws_secret_access_key"] == "caller-secret"
assert merged["aws_session_token"] == "caller-token"
assert merged["aws_region_name"] == "us-west-2"

View file

@ -3131,3 +3131,72 @@ class TestHasUserSetupSso:
monkeypatch.setenv("SAML_IDP_METADATA_XML", "<EntityDescriptor/>")
assert _has_user_setup_sso() is True
class TestIsRequestBodySafeBlocksAwsIdentitySelectors:
"""A caller must not be able to redirect Bedrock signing to another identity
reachable from the proxy host. ``get_credentials`` prefers a named profile
and the AssumeRole knobs over the deployment's static keys, and the file /
batch endpoints fold the request body and the deployment credentials into a
single params dict, so these have to be rejected at the boundary (#36155).
"""
@pytest.mark.parametrize(
"selector",
["aws_profile_name", "aws_session_name", "aws_external_id"],
)
def test_aws_identity_selector_in_batch_body_is_rejected(self, selector):
with pytest.raises(ValueError, match=selector):
is_request_body_safe(
request_body={
"input_file_id": "file-abc123",
"endpoint": "/v1/chat/completions",
"completion_window": "24h",
"model": "bedrock-batch-model",
selector: "attacker-chosen",
},
general_settings={},
llm_router=None,
model="bedrock-batch-model",
)
@pytest.mark.parametrize(
"selector",
["aws_profile_name", "aws_session_name", "aws_external_id"],
)
def test_aws_identity_selector_under_extra_body_is_rejected(self, selector):
with pytest.raises(ValueError, match=selector):
is_request_body_safe(
request_body={
"model": "bedrock-batch-model",
"extra_body": {selector: "attacker-chosen"},
},
general_settings={},
llm_router=None,
model="bedrock-batch-model",
)
def test_aws_identity_selector_allowed_under_proxy_wide_opt_in(self):
assert (
is_request_body_safe(
request_body={
"model": "bedrock-batch-model",
"aws_profile_name": "admin-approved-profile",
},
general_settings={"allow_client_side_credentials": True},
llm_router=None,
model="bedrock-batch-model",
)
is True
)
def test_upload_body_without_identity_selectors_is_accepted(self):
assert (
is_request_body_safe(
request_body={"purpose": "batch", "model": "bedrock-batch-model"},
general_settings={},
llm_router=None,
model="bedrock-batch-model",
)
is True
)

View file

@ -23,6 +23,7 @@ from litellm.proxy.auth.auth_utils import is_request_body_safe # noqa: E402
"aws_web_identity_token",
"aws_sts_endpoint",
"aws_role_name",
"aws_profile_name",
"api_base",
"base_url",
"vertex_credentials",

View file

@ -4062,6 +4062,46 @@ def test_get_deployment_credentials_with_provider_bedrock_batch_fields():
assert credentials["aws_batch_role_arn"] == "arn:aws:iam::123:role/batch-role"
def test_get_deployment_credentials_with_provider_preserves_aws_auth_params():
"""
Test that get_deployment_credentials_with_provider preserves every AWS auth
selector (session token, assume-role, web identity, profile) so bedrock
files/batches deployments using temporary or role-based credentials do not
silently fall back to the server's ambient identity (#36155).
"""
aws_auth_params = {
"aws_access_key_id": "deployment-access-key",
"aws_secret_access_key": "deployment-secret",
"aws_session_token": "deployment-session-token",
"aws_region_name": "us-west-2",
"aws_session_name": "deployment-session",
"aws_profile_name": "deployment-profile",
"aws_role_name": "arn:aws:iam::123:role/deployment-role",
"aws_web_identity_token": "deployment-web-identity",
"aws_sts_endpoint": "https://sts.us-west-2.amazonaws.com",
"aws_external_id": "deployment-external-id",
}
router = litellm.Router(
model_list=[
{
"model_name": "bedrock-batch-model",
"litellm_params": {
"model": "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0",
**aws_auth_params,
},
}
],
)
credentials = router.get_deployment_credentials_with_provider(
model_id="bedrock-batch-model"
)
assert credentials is not None
for key, value in aws_auth_params.items():
assert credentials.get(key) == value, key
def _team_wildcard_model(api_key: str, model_id: str = "team-wildcard-id") -> dict:
return {
"model_name": f"model_name_team-1_{model_id}",

View file

@ -26789,10 +26789,24 @@ export interface components {
aws_bedrock_project_id?: string | null;
/** Aws Bedrock Runtime Endpoint */
aws_bedrock_runtime_endpoint?: string | null;
/** Aws External Id */
aws_external_id?: string | null;
/** Aws Profile Name */
aws_profile_name?: string | null;
/** Aws Region Name */
aws_region_name?: string | null;
/** Aws Role Name */
aws_role_name?: string | null;
/** Aws Secret Access Key */
aws_secret_access_key?: string | null;
/** Aws Session Name */
aws_session_name?: string | null;
/** Aws Session Token */
aws_session_token?: string | null;
/** Aws Sts Endpoint */
aws_sts_endpoint?: string | null;
/** Aws Web Identity Token */
aws_web_identity_token?: string | null;
/** Azure Ad Token */
azure_ad_token?: string | null;
/** Budget Duration */
@ -35467,10 +35481,24 @@ export interface components {
aws_bedrock_project_id?: string | null;
/** Aws Bedrock Runtime Endpoint */
aws_bedrock_runtime_endpoint?: string | null;
/** Aws External Id */
aws_external_id?: string | null;
/** Aws Profile Name */
aws_profile_name?: string | null;
/** Aws Region Name */
aws_region_name?: string | null;
/** Aws Role Name */
aws_role_name?: string | null;
/** Aws Secret Access Key */
aws_secret_access_key?: string | null;
/** Aws Session Name */
aws_session_name?: string | null;
/** Aws Session Token */
aws_session_token?: string | null;
/** Aws Sts Endpoint */
aws_sts_endpoint?: string | null;
/** Aws Web Identity Token */
aws_web_identity_token?: string | null;
/** Azure Ad Token */
azure_ad_token?: string | null;
/** Budget Duration */