Skip to content

fix(bedrock): pass SSE-KMS key through to the batch input-file S3 upload - #35148

Merged
mateo-berri merged 5 commits into
litellm_internal_stagingfrom
litellm_bedrock_batch_sse_kms
Aug 7, 2026
Merged

fix(bedrock): pass SSE-KMS key through to the batch input-file S3 upload#35148
mateo-berri merged 5 commits into
litellm_internal_stagingfrom
litellm_bedrock_batch_sse_kms

Conversation

@devin-ai-integration

@devin-ai-integration devin-ai-integration Bot commented Jul 29, 2026

Copy link
Copy Markdown
Contributor

TLDR

Problem this solves:

  • Bedrock batch input-file upload never sent SSE-KMS to S3
  • KMS-enforced buckets reject the PutObject with a 403
  • So POST /v1/files fails and no batch can be created

How it solves it:

  • Sign the S3 PUT with the aws:kms encryption headers
  • Share one KMS-key resolver between batches and files

User Flow

A platform team runs Bedrock batch jobs through the gateway, and every S3 bucket they are allowed to use denies writes unless the object is encrypted with a specific KMS key.

Before

The upload that has to happen before any batch exists is rejected by S3, so the team never gets a file id and never reaches the batch endpoint at all.

  1. The admin adds a bedrock/... deployment to the gateway config, pointing it at the locked-down bucket and setting the KMS key ARN on it, then starts the gateway
  2. A user posts the batch input file: POST https://litellm-domain/v1/files with purpose=batch, model=bedrock-batch, and a JSONL body of chat-completion requests
  3. The call comes back 403 with S3's AccessDenied text saying the bucket policy requires server-side encryption with the bucket's KMS key
  4. The user retries with different file sizes and content types, and every attempt returns the same 403
  5. There is no file id to pass on, so POST https://litellm-domain/v1/batches is never reachable and the team cannot run batch inference against any of their buckets
  6. The only way through is to drop the bucket's encryption requirement, which their security policy forbids

After

The same upload succeeds, the object lands encrypted with the configured key, and the batch runs to completion.

  1. The admin adds a bedrock/... deployment to the gateway config, pointing it at the locked-down bucket and setting the KMS key ARN on it, then starts the gateway
  2. A user posts the batch input file: POST https://litellm-domain/v1/files with purpose=batch, model=bedrock-batch, and a JSONL body of chat-completion requests
  3. The call comes back 200 with a file- id, the long scrambled kind the gateway hands back for model-routed uploads
  4. Running aws s3api head-object on the bucket shows the new object with ServerSideEncryption: aws:kms and SSEKMSKeyId equal to the ARN configured on the deployment
  5. The user starts the job: POST https://litellm-domain/v1/batches with that file id, a 24h completion window, and the /v1/chat/completions endpoint, and gets back 200 with a batch_ id
  6. The job walks from validating through in progress to completed, and the results land in the bucket, one line per request in the original file

Relevant issues

Fixes #35135

Linear ticket

Pre-Submission checklist

Please complete all items before asking a LiteLLM maintainer to review your PR

  • I have added meaningful tests
  • My PR passes all CI/CD checks (e.g., lint, format, unit tests)
  • My PR's scope is as isolated as possible; it only solves 1 specific problem
  • I have received a Greptile Confidence Score of at least 4/5 before requesting a maintainer review (Greptile reviews automatically once the PR is opened; only comment @greptileai to re-request a review after pushing changes)

Delays in PR merge?

If you're seeing a delay in your PR being merged, ping the LiteLLM Team on Slack (#pr-review).

Screenshots / Proof of Fix

s3_encryption_key_id was only ever used for the batch job's outputDataConfig; the input-file upload that LiteLLM performs first went out with no encryption headers at all, which is what a KMS-enforced bucket rejects.

Two proxies on one machine, same config file, same curl, only the code differs. BEFORE is litellm_internal_staging at f6587fa on port 4917, AFTER is this branch at a2806d4 on port 4918. Everything downstream is real: a customer-managed KMS key, a bucket whose policy denies PutObject unless the object carries that exact key, real s3.us-west-2.amazonaws.com, and a real Bedrock batch of 110 records on us.anthropic.claude-haiku-4-5-20251001-v1:0 billing real $. The steps below are numbered to match the User Flow steps above, one for one.

Step 1 of both flows is the deployment the admin sets up. The bucket policy is what makes this a real test rather than a header-shape check:

{
  "Version": "2012-10-17",
  "Statement": [
    {
      "Sid": "DenyUnencryptedUploads",
      "Effect": "Deny",
      "Principal": "*",
      "Action": "s3:PutObject",
      "Resource": "arn:aws:s3:::litellm-kms-qa-35148/*",
      "Condition": {"StringNotEquals": {"s3:x-amz-server-side-encryption": "aws:kms"}}
    },
    {
      "Sid": "DenyWrongKmsKey",
      "Effect": "Deny",
      "Principal": "*",
      "Action": "s3:PutObject",
      "Resource": "arn:aws:s3:::litellm-kms-qa-35148/*",
      "Condition": {"StringNotEquals": {"s3:x-amz-server-side-encryption-aws-kms-key-id": "arn:aws:kms:us-west-2:439158074652:key/fa719dd3-a738-4b20-940a-a8df971041cd"}}
    }
  ]
}

Config, identical for both proxies:

model_list:
  - model_name: bedrock-batch
    litellm_params:
      model: bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0
      aws_region_name: us-west-2
      s3_bucket_name: litellm-kms-qa-35148
      s3_region_name: us-west-2
      s3_encryption_key_id: arn:aws:kms:us-west-2:439158074652:key/fa719dd3-a738-4b20-940a-a8df971041cd
      aws_batch_role_arn: arn:aws:iam::439158074652:role/litellm-bedrock-batch-role
    model_info:
      mode: batch

general_settings:
  master_key: sk-1234

Before

Steps 2 and 3, the 110-record upload, refused by S3 with the bucket's encryption requirement quoted back:

$ curl -s -w "\nHTTP %{http_code}\n" -X POST http://localhost:4917/v1/files \
    -H "Authorization: Bearer sk-1234" \
    -F purpose=batch -F model=bedrock-batch \
    -F "file=@batch.jsonl;type=application/jsonl"
{"error":{"message":"<?xml version=\"1.0\" encoding=\"UTF-8\"?>\n<Error><Code>AccessDenied</Code><Message>User: arn:aws:iam::439158074652:user/Mateo_Wang is not authorized to perform: s3:PutObject on resource: \"arn:aws:s3:::litellm-kms-qa-35148/litellm-bedrock-files-us.anthropic.claude-haiku-4-5-20251001-v1-0-27a04378-8ae1-4609-9478-fd68e26bb97a.jsonl\" with an explicit deny in a resource-based policy</Message><RequestId>64J0MA00Z4APZ5TF</RequestId><HostId>iOzO9gocfAfO9iRQcj0Vc41AaQOg/MU36Nl6hIiDsaC8Bb4jRvRQ741ntMGIWzdLPelNvUPmjk8YmPkgD3LX/dAGaoCD5su4</HostId></Error>","type":"None","param":"None","code":"403"}}
HTTP 403

Step 4, a smaller file and two other content types, same answer every time:

$ for spec in "batch-small.jsonl application/jsonl" "batch.jsonl application/json" "batch.jsonl text/plain"; do
    f=${spec%% *}; t=${spec##* }
    printf '%s (%s, %s bytes) -> HTTP %s %s\n' "$f" "$t" "$(wc -c < $f | tr -d ' ')" \
      "$(curl -s -o /tmp/r.json -w '%{http_code}' -X POST http://localhost:4917/v1/files \
          -H 'Authorization: Bearer sk-1234' \
          -F purpose=batch -F model=bedrock-batch -F "file=@$f;type=$t")" \
      "$(grep -o '<Code>[^<]*</Code>' /tmp/r.json | head -1)"
  done
batch-small.jsonl (application/jsonl, 504 bytes) -> HTTP 403 <Code>AccessDenied</Code>
batch.jsonl (application/json, 27944 bytes) -> HTTP 403 <Code>AccessDenied</Code>
batch.jsonl (text/plain, 27944 bytes) -> HTTP 403 <Code>AccessDenied</Code>

Step 5, with no file id in hand there is nothing to start a batch from:

$ curl -s -w "\nHTTP %{http_code}\n" -X POST http://localhost:4917/v1/batches \
    -H "Authorization: Bearer sk-1234" -H "Content-Type: application/json" \
    -d '{"input_file_id":"","endpoint":"/v1/chat/completions","completion_window":"24h","model":"bedrock-batch"}'
{"error":{"message":"input_file_id is required for Bedrock batch creation","type":"internal_server_error","param":"None","code":"500"}}
HTTP 500

Step 6 has no command to run: dropping the bucket's encryption requirement is the one move the team's security policy rules out.

After

Steps 2 and 3, the same file, the same bucket, the same curl, now a managed file id:

$ curl -s -w "\nHTTP %{http_code}\n" -X POST http://localhost:4918/v1/files \
    -H "Authorization: Bearer sk-1234" \
    -F purpose=batch -F model=bedrock-batch \
    -F "file=@batch.jsonl;type=application/jsonl"
{"id":"file-bGl0ZWxsbTpzMzovL2xpdGVsbG0ta21zLXFhLTM1MTQ4L2xpdGVsbG0tYmVkcm9jay1maWxlcy11cy5hbnRocm9waWMuY2xhdWRlLWhhaWt1LTQtNS0yMDI1MTAwMS12MS0wLTEzNjA0ZWNkLTJlNjEtNDRiMC05ZjdlLTg1YTk5ODNmYjYxNC5qc29ubDttb2RlbCxiZWRyb2NrLWJhdGNo","bytes":0,"created_at":1786074798,"filename":"litellm-bedrock-files-us.anthropic.claude-haiku-4-5-20251001-v1-0-13604ecd-2e61-44b0-9f7e-85a9983fb614.jsonl","object":"file","purpose":"batch","status":"uploaded","expires_at":null,"status_details":null}
HTTP 200

Step 4, the object really landed encrypted with the configured key, not with a bucket default:

$ aws s3api head-object --bucket litellm-kms-qa-35148 \
    --key litellm-bedrock-files-us.anthropic.claude-haiku-4-5-20251001-v1-0-13604ecd-2e61-44b0-9f7e-85a9983fb614.jsonl \
    --query '{ServerSideEncryption:ServerSideEncryption,SSEKMSKeyId:SSEKMSKeyId,ContentLength:ContentLength}'
{
    "ServerSideEncryption": "aws:kms",
    "SSEKMSKeyId": "arn:aws:kms:us-west-2:439158074652:key/fa719dd3-a738-4b20-940a-a8df971041cd",
    "ContentLength": 23873
}

Step 5, the batch the user came for, created from that file id:

$ curl -s -w "\nHTTP %{http_code}\n" -X POST http://localhost:4918/v1/batches \
    -H "Authorization: Bearer sk-1234" -H "Content-Type: application/json" \
    -d '{"input_file_id":"file-bGl0ZWxsbTpzMzovL2xpdGVsbG0ta21zLXFhLTM1MTQ4L2xpdGVsbG0tYmVkcm9jay1maWxlcy11cy5hbnRocm9waWMuY2xhdWRlLWhhaWt1LTQtNS0yMDI1MTAwMS12MS0wLTEzNjA0ZWNkLTJlNjEtNDRiMC05ZjdlLTg1YTk5ODNmYjYxNC5qc29ubDttb2RlbCxiZWRyb2NrLWJhdGNo","endpoint":"/v1/chat/completions","completion_window":"24h","model":"bedrock-batch"}'
{"id":"batch_bGl0ZWxsbTphcm46YXdzOmJlZHJvY2s6dXMtd2VzdC0yOjQzOTE1ODA3NDY1Mjptb2RlbC1pbnZvY2F0aW9uLWpvYi9oaHlpYmgxYnEyaTY7bW9kZWwsYmVkcm9jay1iYXRjaA","completion_window":"24h","created_at":1786074815,"endpoint":"/v1/chat/completions","input_file_id":"file-bGl0ZWxsbTpzMzovL2xpdGVsbG0ta21zLXFhLTM1MTQ4L2xpdGVsbG0tYmVkcm9jay1maWxlcy11cy5hbnRocm9waWMuY2xhdWRlLWhhaWt1LTQtNS0yMDI1MTAwMS12MS0wLTEzNjA0ZWNkLTJlNjEtNDRiMC05ZjdlLTg1YTk5ODNmYjYxNC5qc29ubDttb2RlbCxiZWRyb2NrLWJhdGNo","object":"batch","status":"validating",...}
HTTP 200

Step 6, the job walks all the way to completion. Polled every 30s, first appearance of each status:

$ while :; do
    printf '%s  %s\n' "$(date -u +%H:%M:%SZ)" \
      "$(aws bedrock get-model-invocation-job --job-identifier $JOB --query status --output text)"
    sleep 30
  done
03:53:49Z  Submitted
03:54:50Z  Validating
03:57:23Z  Scheduled
03:58:25Z  InProgress
04:35:37Z  Completed

And the results land in the bucket, 110 for 110, one line per request in the original file, real tokens billed:

$ aws s3 cp s3://litellm-kms-qa-35148/litellm-batch-outputs/litellm-batch-0c2fa8ae/hhyibh1bq2i6/manifest.json.out -
{"totalRecordCount":110,"processedRecordCount":110,"successRecordCount":110,"errorRecordCount":0,"inputTokenCount":1980,"outputTokenCount":550,...}

$ aws s3 cp s3://litellm-kms-qa-35148/litellm-batch-outputs/litellm-batch-0c2fa8ae/hhyibh1bq2i6/litellm-bedrock-files-us.anthropic.claude-haiku-4-5-20251001-v1-0-13604ecd-2e61-44b0-9f7e-85a9983fb614.jsonl.out - | head -1
{"modelInput":{"messages":[{"role":"user","content":[{"type":"text","text":"Reply with the number 6 and nothing else."}]}],"max_tokens":16,"anthropic_version":"bedrock-2023-05-31"},"modelOutput":{"model":"claude-haiku-4-5-20251001","id":"msg_bdrk_017k9nrrj8x4d9BQcqJgxM5o","type":"message","role":"assistant","content":[{"type":"text","text":"6"}],"stop_reason":"end_turn","usage":{"input_tokens":18,"output_tokens":5}},"recordId":"req-6"}

The status walk above is read straight from Bedrock because the gateway's own GET /v1/batches/{batch_id}, which is where a user would normally watch it, answers AccessDeniedException ... The provided resource ARN is from a different account for a job the proxy itself just created. That is #36155, pre-existing and reproducing on litellm_internal_staging:

$ curl -s -w "\nHTTP %{http_code}\n" "http://localhost:4918/v1/batches/$BID?model=bedrock-batch" \
    -H "Authorization: Bearer sk-1234"
{"error":{"message":"An error occurred (AccessDeniedException) when calling the GetModelInvocationJob operation: The provided resource ARN is from a different account.","type":"internal_server_error","param":"None","code":"500"}}
HTTP 500

The other bug this run surfaced is #36156: each record's body.model has to be the raw bedrock/... id, because an alias makes the upload write a modelInput with no anthropic_version and every record then fails at invoke time behind a job Bedrock still reports as Completed. Neither is touched by this diff and both reproduce on litellm_internal_staging, so both are filed rather than fixed here.

Type

🐛 Bug Fix

Changes

BedrockFilesConfig._sign_s3_request now takes the resolved KMS key and, when there is one, adds x-amz-server-side-encryption: aws:kms plus x-amz-server-side-encryption-aws-kms-key-id to the headers it signs, so both headers land inside SigV4 SignedHeaders rather than being appended afterwards (which S3 would answer with SignatureDoesNotMatch).

The key itself is resolved by a new resolve_s3_encryption_key_id in litellm/llms/bedrock/common_utils.py, shared with the batch-creation path that previously inlined the same precedence chain: litellm_params, then optional_params, then AWS_S3_ENCRYPTION_KEY_ID.

That is the whole diff. s3_encryption_key_id already survives CredentialLiteLLMParams on litellm_internal_staging thanks to #24548, so nothing here touches litellm/types/router.py.

Docs for the widened meaning of the key are in BerriAI/litellm-docs#810: it now encrypts the input file LiteLLM uploads as well as the batch output, so the uploading identity needs kms:GenerateDataKey on it.

Caveats

  • The optional_params tier of resolve_s3_encryption_key_id is unreachable in production: BaseLLMHTTPHandler.create_file and create_batch both pass optional_params={} to the transform, so "the batch path now reads optional_params too" cannot change any outcome, and deleting that argument leaves the whole bedrock suite green
    • Status: won't fix, the decision came out negative
    • Solution: drop the parameter and resolve from litellm_params then the env var only
      • Pro(s) of solution: nothing is left to maintain or test on a path that cannot run
      • Pro(s) mag: +0.2
      • Con(s) of solution: aws_batch_role_arn and bedrock_tags are read from optional_params the same unreachable way a few lines above, so singling this one out makes the file inconsistent, and the tier becomes live the moment the caller stops passing an empty dict
      • Con(s) mag: -0.4
      • Decision: -0.20, so it stays
  • The deployment's own aws_access_key_id and aws_region_name never reach the S3 PUT, because _sign_s3_request reads credentials and region exclusively from that same empty optional_params; only s3_region_name is honoured from litellm_params, so an upload configured with a deployment-specific key pair is signed by the server environment's identity instead

QA runbook

Needs an S3 bucket whose policy denies PutObject unless x-amz-server-side-encryption is aws:kms with a specific key ARN, plus credentials with kms:GenerateDataKey on that key.

  • Bedrock batch input-file upload against an SSE-KMS-enforced bucket
    • Point a bedrock/... deployment at that bucket with s3_encryption_key_id set to the key ARN
    • POST /v1/files with purpose=batch, model=<deployment> and a batch JSONL body; expect a 200 with a managed file id (pre-fix this is a 403 AccessDenied naming the bucket's encryption requirement)
    • Retry the failing upload with other file sizes and content types; expect the same 403 pre-fix, since nothing about the payload is what S3 objects to
    • Confirm the object's encryption with aws s3api head-object; expect ServerSideEncryption: aws:kms and SSEKMSKeyId equal to the configured ARN
    • POST /v1/batches with that file id and expect the job to be created, walk from validating to completed, and write results back to the bucket

Final Attestation

  • The tests check the right things, including the edge cases, and regressions in the respective real-world customer use-cases are not possible after this PR

Link to Devin session: https://app.devin.ai/sessions/050f0554f9bc486ab14918b6ab10aea1

@devin-ai-integration

Copy link
Copy Markdown
Contributor Author

🤖 Devin AI Engineer

I'll be helping with this pull request! Here's what you should know:

✅ I will automatically:

  • Address comments on this PR. Add '(aside)' to your comment to have me ignore it.
  • Look at CI failures and help fix them

Note: I can only respond to comments from users who have write access to this repository.

⚙️ Control Options:

  • Disable automatic comment, CI, and merge conflict monitoring

@CLAassistant

CLAassistant commented Jul 29, 2026

Copy link
Copy Markdown

CLA assistant check
Thank you for your submission! We really appreciate it. Like many open source projects, we ask that you all sign our Contributor License Agreement before we can accept your contribution.
1 out of 2 committers have signed the CLA.

✅ mateo-berri
❌ devin-ai-integration[bot]
You have signed the CLA already but the status is still pending? Let us recheck it.

@greptile-apps

greptile-apps Bot commented Jul 29, 2026

Copy link
Copy Markdown
Contributor

Greptile Summary

Adds SSE-KMS support for Bedrock batch input-file uploads.

  • Centralizes KMS key resolution across Bedrock file uploads and batch creation
  • Includes KMS encryption headers in the S3 SigV4 signature when a key is configured
  • Adds regression coverage for configured, environment-provided, and absent KMS keys

Confidence Score: 5/5

The PR appears safe to merge

No blocking failure remains

Important Files Changed

Filename Overview
litellm/llms/bedrock/common_utils.py Adds a shared resolver that preserves deployment, optional-parameter, and environment precedence for the S3 encryption key
litellm/llms/bedrock/batches/transformation.py Reuses the shared KMS key resolver when constructing Bedrock batch output configuration
litellm/llms/bedrock/files/transformation.py Adds configured SSE-KMS headers before signing Bedrock batch input-file S3 uploads
tests/test_litellm/llms/bedrock/batches/test_transformation.py Updates the secret resolver patch target after moving key resolution into the shared utility
tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py Covers signed SSE-KMS headers, environment fallback, and omission when no key is configured

Reviews (3): Last reviewed commit: "test(router): drop the duplicate s3_encr..." | Re-trigger Greptile

@codecov

codecov Bot commented Jul 29, 2026

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 80.00000% with 3 lines in your changes missing coverage. Please review.

Files with missing lines Patch % Lines
litellm/llms/bedrock/common_utils.py 75.00% 2 Missing ⚠️
litellm/llms/bedrock/files/transformation.py 80.00% 1 Missing ⚠️

📢 Thoughts on this report? Let us know!

@codspeed-hq

codspeed-hq Bot commented Jul 29, 2026

Copy link
Copy Markdown
Contributor

Merging this PR will not alter performance

✅ 31 untouched benchmarks


Comparing litellm_bedrock_batch_sse_kms (a2806d4) with litellm_internal_staging (4e5495e)

Open in CodSpeed

@mateo-berri

Copy link
Copy Markdown
Contributor

@greptileai

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.
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.
@mateo-berri

Copy link
Copy Markdown
Contributor

@greptileai

@mateo-berri mateo-berri left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM

@mateo-berri
mateo-berri merged commit e46721a into litellm_internal_staging Aug 7, 2026
79 checks passed
@mateo-berri
mateo-berri deleted the litellm_bedrock_batch_sse_kms branch August 7, 2026 03:52
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[Bug]: Bedrock batch file upload has no SSE-KMS support, blocking batch creation against KMS-enforced S3 buckets

2 participants