feat(llm): assume an AWS role for cross-account Bedrock

Merged
JU
1d ago
+318-2

Description

Adds cross-account Bedrock. Onyx can run in one AWS account and call Bedrock in another. The Bedrock provider gets two optional fields: Role ARN to assume and External ID.

  • Chat calls: AWS_ROLE_NAME / AWS_EXTERNAL_ID in custom_config map to LiteLLM's aws_role_name / aws_external_id. LiteLLM 1.93 does the STS AssumeRole, including from IRSA, and caches the credentials.
  • Model fetch: /admin/llm/bedrock/available-models assumes the role with the submitted credentials (IAM or access key), then lists models with the temporary credentials.
  • Admin modal: the two fields show for the IAM and access key auth methods. They are cleared for the long-term API key method, because a bearer token is bound to one account.
  • Cloud: multi-tenant deployments reject these keys on provider upsert, on provider test, and on model fetch. The pods' shared identity would let any tenant assume any role that trusts it. The modal still shows the fields on cloud; the backend rejects them (the same pattern as Vertex Workload Identity).

The env var fallback (AWS_ROLE_NAME on the pods) works today without this change. This PR makes the setting per-provider and fixes the model fetch, which ignored the role.

Screenshots

Access key + role. The role fields show below the credentials. Both are optional.

Access key auth with role ARN and external ID

Environment IAM role + role. The pod's own role (for example IRSA) assumes the role in the other account.

IAM auth with role ARN and external ID

Long-term API key. The role fields are hidden and cleared, because a bearer token works in only one account.

Long-term API key auth hides the role fields

Model fetch with a role. With dummy keys, the error comes from the AssumeRole call. This shows the fetch assumes the role before it lists models.

Model fetch fails at AssumeRole with dummy keys

How Has This Been Tested?

  • Unit: the role keys map to LiteLLM kwargs; the model fetch calls sts.assume_role with the base credentials and lists models with the temporary credentials; multi-tenant rejection in the model fetch and in the write-path gate.
  • tsc, oxlint, oxfmt run by hand (no bun locally); CI runs the real hooks.
  • Local run of this branch (fresh DB): the modal shows and clears the fields for each auth method; the model fetch returns 400 CREDENTIAL_INVALID from AssumeRole with dummy keys.
  • Not yet tested against two real AWS accounts.

Additional Options

  • [Optional] Please cherry-pick this PR to the latest release version.

🤖 Generated with Claude Code


Summary by cubic

Adds cross-account Bedrock support so Onyx can run in one AWS account and call Bedrock in another, with the Bedrock provider assuming a role.

  • Adds optional Role ARN and External ID fields to the Bedrock provider.
  • Chat calls map the fields to LiteLLM's aws_role_name and aws_external_id, which performs the STS AssumeRole and caches credentials.
  • Only the uppercase AWS_ROLE_NAME/AWS_EXTERNAL_ID spellings are accepted; lowercase spellings are dropped when an admin edits the provider.
  • /admin/llm/bedrock/available-models now assumes the role with the submitted credentials before listing models; previously it ignored the role.
  • The admin modal shows the two fields for the IAM and access key auth methods, and clears them when the long-term API key method is selected, since a bearer token is bound to one account.

Security

  • Multi-tenant deployments reject these keys on provider upsert, provider test, and model fetch, because the shared server identity would let any tenant assume roles that trust it.

Summary by CodeRabbit

  • New Features

    • Amazon Bedrock configurations now support cross-account IAM role assumption with an optional external ID.
    • Administrators can enter these settings in the language-model setup and use them when discovering available models.
    • Saved role settings are displayed for existing configurations and cleared when bearer-token authentication is selected.
    • Role assumption is unavailable in multi-tenant deployments and unsupported for bearer-token authentication.
  • Documentation

    • Added setup guidance and localized interface text for the new Bedrock options.