Skip to content

feat(bedrock): add AWS SigV4 and STS web identity authentication#5388

Open
skamenan7 wants to merge 2 commits intollamastack:mainfrom
skamenan7:fix/4730-bedrock-sts-auth-v2
Open

feat(bedrock): add AWS SigV4 and STS web identity authentication#5388
skamenan7 wants to merge 2 commits intollamastack:mainfrom
skamenan7:fix/4730-bedrock-sts-auth-v2

Conversation

@skamenan7
Copy link
Copy Markdown
Contributor

What does this PR do?

The Bedrock inference provider previously depended on a pre-signed bearer token (AWS_BEARER_TOKEN_BEDROCK). That works for short-lived manual setups, but it is a poor fit for environments that already rely on AWS identity, like Kubernetes/OpenShift with IRSA, GitHub Actions with OIDC, EC2, ECS, and Lambda.

This PR adds AWS SigV4 support for the Bedrock OpenAI-compatible endpoint using standard AWS credential sources. When no api_key is configured, requests are signed with SigV4 via botocore. Static credentials, profiles, IAM roles, and web identity are supported. When aws_role_arn is configured, the explicit assume-role and web-identity path uses RefreshableBotoSession to refresh temporary credentials automatically. Bearer token mode is unchanged: if api_key is set in config or passed via x-llamastack-provider-data, it takes precedence.

This also updates the endpoint URL from the legacy bedrock-mantle hostname to bedrock-runtime.<region>.amazonaws.com/openai/v1.

Closes #4730

Auth modes

Bearer token, unchanged:

providers:
  inference:
    - provider_type: remote::bedrock
      config:
        api_key: ${env.AWS_BEARER_TOKEN_BEDROCK}
        region_name: us-west-2

SigV4 via AWS credentials, new:

providers:
  inference:
    - provider_type: remote::bedrock
      config:
        region_name: us-west-2
        # api_key intentionally omitted

STS web identity / IRSA, new:

providers:
  inference:
    - provider_type: remote::bedrock
      config:
        region_name: us-west-2
        aws_role_arn: ${env.AWS_ROLE_ARN}
        aws_web_identity_token_file: ${env.AWS_WEB_IDENTITY_TOKEN_FILE}

Per-request bearer override still works on a SigV4-mode server. If x-llamastack-provider-data includes {"aws_bearer_token_bedrock": "<token>"}, that request uses the bearer path. Empty, whitespace-only, or null values fall back to SigV4.

Implementation notes

  • BedrockSigV4Auth is an httpx.Auth implementation that removes the OpenAI SDK placeholder Authorization header and replaces it with a SigV4 signature generated through botocore.
  • The signed httpx.AsyncClient is built once in initialize() and reused, so get_extra_client_params() can stay synchronous and the adapter does not create a new client per request.
  • asyncio.to_thread() is used so signing and credential resolution do not block the event loop.
  • asyncio.shield() is used around signing and shutdown cleanup so cancellation does not interrupt those operations halfway through.

Test plan

Unit tests:

uv run pytest \
  tests/unit/providers/inference/bedrock/ \
  tests/unit/providers/inference/test_bedrock_adapter.py \
  tests/unit/providers/inference/test_bedrock_config.py \
  tests/unit/providers/inference/test_bedrock_sts.py \
  tests/unit/providers/safety/test_bedrock_safety_adapter.py \
  -v --tb=short

Focused SigV4 unit tests:

uv run pytest tests/unit/providers/inference/bedrock/test_sigv4_auth.py -x --tb=short

Local live validation in SigV4 mode:

Server config with api_key intentionally omitted:

providers:
  inference:
    - provider_type: remote::bedrock
      config:
        region_name: us-west-2

Validated locally for:

  • non-streaming requests
  • streaming requests
  • bearer override on a SigV4-mode server
  • empty, whitespace, and null bearer values falling back to SigV4
  • malformed provider-data handling
  • invalid bearer rejection without leaking token data
  • concurrent mixed-auth request isolation

@meta-cla meta-cla bot added the CLA Signed This label is managed by the Meta Open Source bot. label Mar 30, 2026
@mergify
Copy link
Copy Markdown
Contributor

mergify bot commented Mar 30, 2026

This pull request has merge conflicts that must be resolved before it can be merged. @skamenan7 please rebase it. https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

@mergify mergify bot added the needs-rebase label Mar 30, 2026
@skamenan7 skamenan7 force-pushed the fix/4730-bedrock-sts-auth-v2 branch 2 times, most recently from 742fe2a to 4302e2f Compare March 30, 2026 22:15
@mergify mergify bot removed the needs-rebase label Mar 30, 2026
@skamenan7 skamenan7 force-pushed the fix/4730-bedrock-sts-auth-v2 branch 8 times, most recently from d8c9c75 to c1851ac Compare March 31, 2026 17:16
@skamenan7 skamenan7 marked this pull request as ready for review March 31, 2026 17:37
@skamenan7 skamenan7 force-pushed the fix/4730-bedrock-sts-auth-v2 branch 3 times, most recently from 43847c4 to 47fdbc5 Compare April 1, 2026 15:16
@mergify
Copy link
Copy Markdown
Contributor

mergify bot commented Apr 1, 2026

This pull request has merge conflicts that must be resolved before it can be merged. @skamenan7 please rebase it. https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

@mergify mergify bot added the needs-rebase label Apr 1, 2026
@skamenan7 skamenan7 force-pushed the fix/4730-bedrock-sts-auth-v2 branch from 47fdbc5 to 21a1f1b Compare April 1, 2026 16:01
@mergify mergify bot removed the needs-rebase label Apr 1, 2026
The Bedrock inference provider previously required a pre-signed bearer
token (AWS_BEARER_TOKEN_BEDROCK). This PR adds full AWS credential chain
support so Bedrock works natively in EKS/IRSA, GitHub Actions OIDC,
EC2, ECS, and Lambda without managing long-lived credentials.

When no api_key is configured, requests are signed using AWS SigV4 via
botocore. STS role assumption and web identity federation are supported
through RefreshableBotoSession, which refreshes credentials
automatically. Bearer token mode is unchanged — if api_key is set in
config or passed via x-llamastack-provider-data, it takes precedence.

Also corrects the endpoint URL from bedrock-mantle to
bedrock-runtime.<region>.amazonaws.com/openai/v1, and gates the
bedrock model in ci-tests on AWS_DEFAULT_REGION (works for both
bearer and SigV4 modes) instead of AWS_BEARER_TOKEN_BEDROCK.

Closes llamastack#4730

Signed-off-by: skamenan7 <skamenan@redhat.com>
the provider-compat-matrix pre-commit hook generates this file from
recording data; existing bedrock recordings used bedrock-mantle
(the old endpoint) so the hook output reflects that correctly

Signed-off-by: skamenan7 <skamenan@redhat.com>
@skamenan7 skamenan7 force-pushed the fix/4730-bedrock-sts-auth-v2 branch from 21a1f1b to 99db583 Compare April 1, 2026 17:18
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

CLA Signed This label is managed by the Meta Open Source bot.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Add support for AWS Security Token Service (STS) for AWS Bedrock inference provider

1 participant