From 18d9c7aa21e1308c5ecf05254b29bc0715965bde Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 29 Jul 2026 22:09:49 +0000 Subject: [PATCH 1/3] fix(bedrock): pass SSE-KMS key through to the batch input-file S3 upload --- .../llms/bedrock/batches/transformation.py | 8 ++- litellm/llms/bedrock/common_utils.py | 19 ++++- litellm/llms/bedrock/files/transformation.py | 16 ++++- litellm/types/router.py | 1 + .../bedrock/batches/test_transformation.py | 2 +- .../test_bedrock_files_transformation.py | 72 ++++++++++++++++++- tests/test_litellm/test_router.py | 31 ++++++++ 7 files changed, 141 insertions(+), 8 deletions(-) diff --git a/litellm/llms/bedrock/batches/transformation.py b/litellm/llms/bedrock/batches/transformation.py index a4ff1c78467b..7500531b81ae 100644 --- a/litellm/llms/bedrock/batches/transformation.py +++ b/litellm/llms/bedrock/batches/transformation.py @@ -12,7 +12,6 @@ from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.llms.base_llm.batches.transformation import BaseBatchesConfig from litellm.llms.base_llm.chat.transformation import BaseLLMException -from litellm.secret_managers.main import get_secret_str from litellm.types.llms.bedrock import ( BedrockCreateBatchRequest, BedrockCreateBatchResponse, @@ -29,7 +28,7 @@ from litellm.types.utils import LiteLLMBatch, LlmProviders from ..base_aws_llm import BaseAWSLLM -from ..common_utils import CommonBatchFilesUtils +from ..common_utils import CommonBatchFilesUtils, resolve_s3_encryption_key_id # Bedrock batch input files are uploaded as # s3://bucket/litellm-bedrock-files-{model, ":" -> "-"}-{uuid4}.jsonl (see @@ -200,7 +199,10 @@ def transform_create_batch_request( ) # Add optional KMS encryption key ID if provided - s3_encryption_key_id = litellm_params.get("s3_encryption_key_id") or get_secret_str("AWS_S3_ENCRYPTION_KEY_ID") + s3_encryption_key_id = resolve_s3_encryption_key_id( + litellm_params=litellm_params, + optional_params=optional_params, + ) if s3_encryption_key_id: s3_output_config["s3EncryptionKeyId"] = s3_encryption_key_id diff --git a/litellm/llms/bedrock/common_utils.py b/litellm/llms/bedrock/common_utils.py index 5114677ffc0c..9d427fa6f12b 100644 --- a/litellm/llms/bedrock/common_utils.py +++ b/litellm/llms/bedrock/common_utils.py @@ -35,7 +35,7 @@ ) from litellm.llms.base_llm.base_utils import BaseLLMModelInfo, BaseTokenCounter from litellm.llms.base_llm.chat.transformation import BaseLLMException -from litellm.secret_managers.main import get_secret +from litellm.secret_managers.main import get_secret, get_secret_str if TYPE_CHECKING: from litellm.types.llms.openai import AllMessageValues @@ -1313,6 +1313,23 @@ def get_anthropic_beta_from_headers(headers: dict) -> List[str]: return [] +def resolve_s3_encryption_key_id( + litellm_params: Mapping[str, Any], + optional_params: Mapping[str, Any] | None = None, +) -> str | None: + """ + Resolve the SSE-KMS key configured for Bedrock batch/file S3 objects. + + Precedence: `s3_encryption_key_id` in litellm_params, then optional_params + (client-side / request params), then the AWS_S3_ENCRYPTION_KEY_ID env var. + """ + for source in (litellm_params, optional_params or {}): + value = source.get("s3_encryption_key_id") + if isinstance(value, str) and value: + return value + return get_secret_str("AWS_S3_ENCRYPTION_KEY_ID") + + class CommonBatchFilesUtils: """ Common utilities for Bedrock batch and file operations. diff --git a/litellm/llms/bedrock/files/transformation.py b/litellm/llms/bedrock/files/transformation.py index d4865a1c87a1..d1674b260b45 100644 --- a/litellm/llms/bedrock/files/transformation.py +++ b/litellm/llms/bedrock/files/transformation.py @@ -53,7 +53,7 @@ from litellm.utils import get_llm_provider from ..base_aws_llm import BaseAWSLLM -from ..common_utils import BedrockError +from ..common_utils import BedrockError, 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 @@ -741,6 +741,10 @@ def transform_create_file_request( content=file_content, api_base=api_base, optional_params=optional_params, + s3_encryption_key_id=resolve_s3_encryption_key_id( + litellm_params=litellm_params, + optional_params=optional_params, + ), ) litellm_params["upload_url"] = api_base @@ -758,6 +762,7 @@ def _sign_s3_request( content: str, api_base: str, optional_params: dict, + s3_encryption_key_id: str | None = None, ) -> Tuple[dict, str]: """ Sign S3 PUT request using the same proven logic as S3Logger. @@ -790,11 +795,20 @@ def _sign_s3_request( content_hash = hashlib.sha256(content.encode("utf-8")).hexdigest() # Prepare headers with required S3 headers (same as s3_v2.py) + sse_headers = ( + { + "x-amz-server-side-encryption": "aws:kms", + "x-amz-server-side-encryption-aws-kms-key-id": s3_encryption_key_id, + } + if s3_encryption_key_id + else {} + ) request_headers = { "Content-Type": "application/json", # JSONL files are JSON content "x-amz-content-sha256": content_hash, # REQUIRED by S3 "Content-Language": "en", "Cache-Control": "private, immutable, max-age=31536000, s-maxage=0", + **sse_headers, } # Use requests.Request to prepare the request (same pattern as s3_v2.py) diff --git a/litellm/types/router.py b/litellm/types/router.py index 28e4a8272e87..c4d679a2500f 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -211,6 +211,7 @@ class CredentialLiteLLMParams(BaseModel): aws_bedrock_runtime_endpoint: Optional[str] = None aws_bedrock_project_id: Optional[str] = None s3_bucket_name: Optional[str] = None + s3_encryption_key_id: Optional[str] = None ## IBM WATSONX ## watsonx_region_name: Optional[str] = None diff --git a/tests/test_litellm/llms/bedrock/batches/test_transformation.py b/tests/test_litellm/llms/bedrock/batches/test_transformation.py index 3681daffe5e7..01420eb10df6 100644 --- a/tests/test_litellm/llms/bedrock/batches/test_transformation.py +++ b/tests/test_litellm/llms/bedrock/batches/test_transformation.py @@ -172,7 +172,7 @@ def test_create_request_omits_kms_key_when_absent(config): "generate_unique_job_name", return_value="litellm-batch-1", ), patch.object(config.common_utils, "sign_aws_request") as mock_sign, patch( - "litellm.llms.bedrock.batches.transformation.get_secret_str", + "litellm.llms.bedrock.common_utils.get_secret_str", return_value=None, ): mock_sign.return_value = ({}, b"{}") diff --git a/tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py b/tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py index c548fe53e156..a57e5801327d 100644 --- a/tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py +++ b/tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py @@ -442,7 +442,7 @@ def test_transform_create_file_request_injects_s3_region_for_signing(self): captured_optional_params: dict = {} - def fake_sign(content, api_base, optional_params): + def fake_sign(content, api_base, optional_params, s3_encryption_key_id=None): captured_optional_params.update(optional_params) return {"Authorization": "fake"}, content @@ -498,7 +498,7 @@ def test_s3_region_name_wins_over_aws_region_name_for_signing(self): captured_optional_params: dict = {} - def fake_sign(content, api_base, optional_params): + def fake_sign(content, api_base, optional_params, s3_encryption_key_id=None): captured_optional_params.update(optional_params) return {"Authorization": "fake"}, content @@ -514,6 +514,74 @@ def fake_sign(content, api_base, optional_params): captured_optional_params.get("aws_region_name") == "us-gov-west-1" ), "s3_region_name must override aws_region_name for SigV4 signing" + def _signed_upload_request(self, litellm_params: dict) -> dict: + from litellm.llms.bedrock.files.transformation import BedrockFilesConfig + + config = BedrockFilesConfig() + jsonl_content = json.dumps( + { + "custom_id": "req-1", + "method": "POST", + "url": "/v1/chat/completions", + "body": { + "model": "bedrock/amazon.nova-pro-v1:0", + "messages": [{"role": "user", "content": "Hello"}], + "max_tokens": 10, + }, + } + ).encode() + + request = config.transform_create_file_request( + model="amazon.nova-pro-v1:0", + create_file_data={ + "file": ("batch.jsonl", jsonl_content, "application/jsonl"), + "purpose": "batch", + }, + optional_params={ + "aws_access_key_id": "test-key-id", + "aws_secret_access_key": "test-secret", + "aws_region_name": "us-west-2", + }, + litellm_params={"s3_bucket_name": "litellm-batch-bucket", **litellm_params}, + ) + assert isinstance(request, dict) + return request + + def test_upload_signs_sse_kms_headers_when_key_configured(self, monkeypatch): + """ + Buckets whose policy requires SSE-KMS reject the batch input-file PutObject + unless the upload carries the aws:kms encryption headers; they must also be + covered by SigV4 SignedHeaders or S3 answers SignatureDoesNotMatch. + """ + monkeypatch.delenv("AWS_S3_ENCRYPTION_KEY_ID", raising=False) + kms_key = "arn:aws:kms:us-west-2:1234:key/abcd" + + request = self._signed_upload_request({"s3_encryption_key_id": kms_key}) + + headers = {key.lower(): value for key, value in request["headers"].items()} + assert headers["x-amz-server-side-encryption"] == "aws:kms" + assert headers["x-amz-server-side-encryption-aws-kms-key-id"] == kms_key + signed_headers = headers["authorization"].split("SignedHeaders=")[1].split(",")[0] + assert "x-amz-server-side-encryption" in signed_headers + assert "x-amz-server-side-encryption-aws-kms-key-id" in signed_headers + + def test_upload_reads_sse_kms_key_from_env(self, monkeypatch): + monkeypatch.setenv("AWS_S3_ENCRYPTION_KEY_ID", "env-kms-key") + + request = self._signed_upload_request({}) + + headers = {key.lower(): value for key, value in request["headers"].items()} + assert headers["x-amz-server-side-encryption-aws-kms-key-id"] == "env-kms-key" + + def test_upload_omits_sse_headers_when_no_key_configured(self, monkeypatch): + monkeypatch.delenv("AWS_S3_ENCRYPTION_KEY_ID", raising=False) + + request = self._signed_upload_request({}) + + headers = {key.lower() for key in request["headers"]} + assert "x-amz-server-side-encryption" not in headers + assert "x-amz-server-side-encryption-aws-kms-key-id" not in headers + def test_openai_passthrough_still_works(self): """ Regression test: ensure OpenAI-compatible models (e.g. gpt-oss) diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 46b5ce65c3fd..fa047d7ee46a 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -3666,6 +3666,37 @@ def test_get_deployment_credentials_with_provider_includes_bucket_name(): assert credentials["custom_llm_provider"] == "vertex_ai" +def test_get_deployment_credentials_with_provider_includes_s3_encryption_key_id(): + """ + Regression: s3_encryption_key_id must survive the CredentialLiteLLMParams filter, + otherwise the Bedrock batch input-file upload loses the SSE-KMS key and S3 rejects + the PutObject on buckets whose policy requires aws:kms encryption. + """ + router = litellm.Router( + model_list=[ + { + "model_name": "bedrock-batch", + "litellm_params": { + "model": "bedrock/anthropic.claude-sonnet-4-20250514-v1:0", + "aws_region_name": "us-west-2", + "s3_bucket_name": "my-batch-bucket", + "s3_encryption_key_id": "arn:aws:kms:us-west-2:1234:key/abcd", + }, + } + ], + ) + + credentials = router.get_deployment_credentials_with_provider( + model_id="bedrock-batch" + ) + + assert credentials is not None + assert ( + credentials["s3_encryption_key_id"] + == "arn:aws:kms:us-west-2:1234:key/abcd" + ) + + def test_get_deployment_credentials_with_provider_resolves_credential_name(): """ Test that get_deployment_credentials_with_provider correctly resolves From 6b2ac7cc5b79bc173e5aea73f963a8b34089f885 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 6 Aug 2026 19:25:47 -0700 Subject: [PATCH 2/3] refactor(bedrock): freeze the SSE-KMS header and key-source collections The staging merge tightened the LIT002 ceiling, so the three mutable dict literals this branch added now breach it. Build the S3 request headers as MappingProxyType and resolve the encryption key from a tuple of sources. --- litellm/llms/bedrock/common_utils.py | 10 ++++---- litellm/llms/bedrock/files/transformation.py | 26 +++++++++++--------- 2 files changed, 20 insertions(+), 16 deletions(-) diff --git a/litellm/llms/bedrock/common_utils.py b/litellm/llms/bedrock/common_utils.py index c5ef78b44f3e..d18cb7d87343 100644 --- a/litellm/llms/bedrock/common_utils.py +++ b/litellm/llms/bedrock/common_utils.py @@ -1314,11 +1314,11 @@ def resolve_s3_encryption_key_id( Precedence: `s3_encryption_key_id` in litellm_params, then optional_params (client-side / request params), then the AWS_S3_ENCRYPTION_KEY_ID env var. """ - for source in (litellm_params, optional_params or {}): - value = source.get("s3_encryption_key_id") - if isinstance(value, str) and value: - return value - return get_secret_str("AWS_S3_ENCRYPTION_KEY_ID") + candidates: Final = tuple( + source.get("s3_encryption_key_id") for source in (litellm_params, optional_params) if source is not None + ) + explicit: Final = next((value for value in candidates if isinstance(value, str) and value), None) + return explicit or get_secret_str("AWS_S3_ENCRYPTION_KEY_ID") class CommonBatchFilesUtils: diff --git a/litellm/llms/bedrock/files/transformation.py b/litellm/llms/bedrock/files/transformation.py index b4c4a270bd7c..728fd02f0019 100644 --- a/litellm/llms/bedrock/files/transformation.py +++ b/litellm/llms/bedrock/files/transformation.py @@ -788,20 +788,24 @@ def _sign_s3_request( # Prepare headers with required S3 headers (same as s3_v2.py) sse_headers: Final = ( + MappingProxyType( + { + "x-amz-server-side-encryption": "aws:kms", + "x-amz-server-side-encryption-aws-kms-key-id": s3_encryption_key_id, + } + ) + if s3_encryption_key_id + else MappingProxyType({}) + ) + request_headers: Final = MappingProxyType( { - "x-amz-server-side-encryption": "aws:kms", - "x-amz-server-side-encryption-aws-kms-key-id": s3_encryption_key_id, + "Content-Type": "application/json", # JSONL files are JSON content + "x-amz-content-sha256": content_hash, # REQUIRED by S3 + "Content-Language": "en", + "Cache-Control": "private, immutable, max-age=31536000, s-maxage=0", + **sse_headers, } - if s3_encryption_key_id - else {} ) - request_headers: Final = { - "Content-Type": "application/json", # JSONL files are JSON content - "x-amz-content-sha256": content_hash, # REQUIRED by S3 - "Content-Language": "en", - "Cache-Control": "private, immutable, max-age=31536000, s-maxage=0", - **sse_headers, - } # Use requests.Request to prepare the request (same pattern as s3_v2.py) req: Final = requests.Request("PUT", api_base, data=content, headers=request_headers) From a2806d430f55c2d372dbd633687d84ad207b5131 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 6 Aug 2026 19:58:30 -0700 Subject: [PATCH 3/3] test(router): drop the duplicate s3_encryption_key_id credential test test_get_deployment_credentials_with_provider_bedrock_batch_fields already covers s3_encryption_key_id on the base branch, and the new test passes with every production file in this branch reverted, so it guards nothing. --- tests/test_litellm/test_router.py | 31 ------------------------------- 1 file changed, 31 deletions(-) diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 1576fa3dbcdf..4a3395a7d3fa 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -3976,37 +3976,6 @@ def test_get_deployment_credentials_with_provider_includes_bucket_name(): assert credentials["custom_llm_provider"] == "vertex_ai" -def test_get_deployment_credentials_with_provider_includes_s3_encryption_key_id(): - """ - Regression: s3_encryption_key_id must survive the CredentialLiteLLMParams filter, - otherwise the Bedrock batch input-file upload loses the SSE-KMS key and S3 rejects - the PutObject on buckets whose policy requires aws:kms encryption. - """ - router = litellm.Router( - model_list=[ - { - "model_name": "bedrock-batch", - "litellm_params": { - "model": "bedrock/anthropic.claude-sonnet-4-20250514-v1:0", - "aws_region_name": "us-west-2", - "s3_bucket_name": "my-batch-bucket", - "s3_encryption_key_id": "arn:aws:kms:us-west-2:1234:key/abcd", - }, - } - ], - ) - - credentials = router.get_deployment_credentials_with_provider( - model_id="bedrock-batch" - ) - - assert credentials is not None - assert ( - credentials["s3_encryption_key_id"] - == "arn:aws:kms:us-west-2:1234:key/abcd" - ) - - def test_get_deployment_credentials_with_provider_resolves_credential_name(): """ Test that get_deployment_credentials_with_provider correctly resolves