feat: support cross-account Bedrock AssumeRole in AI Bridge - #26527
Conversation
Docs preview📖 View docs preview for |
This comment has been minimized.
This comment has been minimized.
a63a21b to
dbf927d
Compare
8e7868d to
4ce147a
Compare
4ce147a to
65cb77c
Compare
|
/coder-agents-review |
|
Chat: Review posted | View chat Review historydeep-review v0.9.0 | Round 2 | Last posted: Round 2, 19 findings (2 P2, 11 P3, 6 Nit), COMMENT. Review Finding inventoryFindings
Contested and acknowledgedCRF-4 (P3, aibridge/provider/bedrock.go:18) - bedrockSessionName naming
CRF-6 (P3, codersdk/aiproviders_bedrock.go:20) - PATCH clears RoleARN
CRF-7 (P3, aibridge/provider/bedrock.go:71) - No ExternalId
CRF-9 (P3, aibridge/provider/bedrock.go:20) - Doc comment bloat
CRF-10 (P3, aibridge/provider/anthropic.go:60) - Inline comment restates code
CRF-13 (Nit, aibridge/provider/anthropic.go:48) - Missing doc comment on NewAnthropic
CRF-14 (Nit, aibridge/provider/bedrock.go:27) - Mid-sentence capitalization
Round logRound 1Panel (12 reviewers). 2 P2, 8 P3, 5 Nit. Reviewed against d5ec26b..511e209. Round 2Churn guard: 8 addressed, 7 contested. PROCEED. Panel (8 reviewers). All 7 contested findings closed by panel. 3 P3, 1 Nit new. Reviewed against d5ec26b..ed11c8c. About deep-reviewCRF = Coder Review Finding (P0-P4, Nit, Note)
|
Documentation CheckUpdates Needed
Automated review via Coder Agents |
There was a problem hiding this comment.
The credential resolution refactoring is well-structured: lifting AWS config loading from per-request into provider construction eliminates repeated SDK config parsing and makes credential failures visible at startup. The aws.NewCredentialsCache wrapping stscreds.AssumeRoleProvider is the standard AWS pattern, and the test suite covers the new paths thoroughly, including a mock STS endpoint for AssumeRole.
Severity summary: 2 P2, 8 P3, 5 Nit.
The two P2s are a region regression (the new manual aws.Config assembly loses the SDK-resolved region that LoadDefaultConfig previously carried) and a misleading error message ("no AWS credentials found" now also wraps STS AssumeRole denials, sending operators down the wrong debugging path).
"Someone searching CloudTrail for this session name will grep the codebase and find nothing; someone grepping for 'aigateway' will find only this constant."
-- Gon, on the naming outlier
🤖 This review was automatically generated with Coder Agents.
| require.NotNil(t, persisted.Bedrock) | ||
| require.Equal(t, "arn:aws:iam::123456789012:role/target", persisted.Bedrock.RoleARN) | ||
| require.NotNil(t, persisted.Bedrock.AccessKey) | ||
| require.Equal(t, "", *persisted.Bedrock.AccessKey) |
There was a problem hiding this comment.
Nit [CRF-11] MigrateStaticToRole sends AccessKeySecret: ptr.Ref("") in the PATCH and reads the DB row, but only asserts AccessKey == "". For a test titled "MigrateStaticToRole" the full migration should be shown:
require.NotNil(t, persisted.Bedrock.AccessKeySecret)
require.Equal(t, "", *persisted.Bedrock.AccessKeySecret)(Bisky)
🤖
There was a problem hiding this comment.
Verified fixed. Both AccessKey and AccessKeySecret are now asserted cleared.
🤖
| func (*interceptionBase) withAWSBedrockOptions(ctx context.Context, cfg *aibconfig.AWSBedrock) ([]option.RequestOption, error) { | ||
| if cfg == nil { | ||
| return nil, xerrors.New("nil config given") | ||
| // Credentials come from i.bedrock.Creds, it is a shared credentials cache, so the per-request |
There was a problem hiding this comment.
Nit [CRF-12] Comma splice: "Credentials come from i.bedrock.Creds, it is a shared credentials cache" is two independent clauses joined by a comma. Use a semicolon or split into two sentences.
(Leorio)
🤖
| } | ||
|
|
||
| func NewAnthropic(cfg config.Anthropic, bedrockCfg *config.AWSBedrock) *Anthropic { | ||
| func NewAnthropic(ctx context.Context, cfg config.Anthropic, bedrockCfg *config.AWSBedrock) (*Anthropic, error) { |
There was a problem hiding this comment.
Nit [CRF-13] Exported constructor NewAnthropic has no doc comment. The signature change (added ctx, now returns error) makes this the moment to document when it errors, what ctx is used for, and why bedrockCfg is a pointer (nil = non-Bedrock).
(Leorio)
🤖
There was a problem hiding this comment.
It existed before PR, I think it's okay.
There was a problem hiding this comment.
Panel evaluated (3/3 close). The missing doc comment pre-dates this PR and is not a regression introduced by the signature change. Closing.
🤖
| // | ||
| // The result is wrapped in aws.NewCredentialsCache, which caches and rotates | ||
| // the resolved temporary credentials. buildBedrockCredentials should be called | ||
| // once when the Bedrock provider is constructed, and the returned Credential |
There was a problem hiding this comment.
Nit [CRF-14] "Credential Provider" and "Bedrock Provider" are capitalized mid-sentence. Go doc convention is lowercase unless it is a proper noun or a type reference. Write "credentials provider" and "Bedrock provider".
(Leorio)
🤖
There was a problem hiding this comment.
I think it's fine.
There was a problem hiding this comment.
Panel evaluated (2/2 close). The capitalization matches the AWS SDK type names (aws.CredentialsProvider), which is defensible under Go doc convention. Closing.
🤖
| ) | ||
|
|
||
| var bridgeTestTracer = otel.Tracer("bridge_test") | ||
|
|
There was a problem hiding this comment.
Nit [CRF-15] Two names for the same helper across files: mustNewAnthropicProvider in bridge_test.go, aibridged_test.go, aibridged_integration_test.go vs mustNewAnthropic in anthropic_internal_test.go, setupbridge.go. mustNewAnthropic matches the constructor NewAnthropic; the Provider suffix is redundant. Also, passthrough_internal_test.go avoids the helper and uses require.NoError(t, err) directly, which is the better pattern when t is available.
(Gon, Meruem)
🤖
There was a problem hiding this comment.
I think it's outdated.
dannykopping
left a comment
There was a problem hiding this comment.
Looking good, but added a few minor points.
Have you tested this using a real AWS account?
| enabledCount++ | ||
| } | ||
| prov, err := buildAIProviderFromRow(row, keysByProvider[row.ID], cfg, metrics) | ||
| prov, err := buildAIProviderFromRow(ctx, row, keysByProvider[row.ID], cfg, metrics) |
There was a problem hiding this comment.
Given we're now passing ctx through, do we need to reconsider timeouts etc?
There was a problem hiding this comment.
buildBedrockCredentials is supposed to be lazy, expensive credential resolution (e.g. STS AssumeRole call) happens in withAWSBedrockOptions function.
But I'll consider adding 30secs timeout.
There was a problem hiding this comment.
I thought about it more, looks like buildBedrockCredentials and LoadDefaultConfig don't make network calls, unless specifically configured to do so.
So it should be fast, lazy construction calls. So I think timeout is unnecessary.
| "github.com/coder/quartz" | ||
| ) | ||
|
|
||
| func mustNewAnthropicProvider(cfg aibridge.AnthropicConfig, bedrockCfg *aibridge.AWSBedrockConfig) aibridge.Provider { |
There was a problem hiding this comment.
Repeated pattern in a few places; can we abstract?
There was a problem hiding this comment.
It requires creating a new pkg to avoid import cycles, e.g.: #26636.
I don't mind to merge this example, if you're okay with that.
Helper usage is split across multiple places:
┌─────┬──────────────────────────────────────────────────┬─────────────────┬───────────────────┬──────────────────────┐
│ # │ File │ Pkg │ Returns │ Consolidatable? │
├─────┼──────────────────────────────────────────────────┼─────────────────┼───────────────────┼──────────────────────┤
│ 1 │ aibridge/bridge_test.go │ aibridge_test │ aibridge.Provider │ ✅ │
├─────┼──────────────────────────────────────────────────┼─────────────────┼───────────────────┼──────────────────────┤
│ 2 │ coderd/aibridged/aibridged_test.go │ aibridged_test │ aibridge.Provider │ ✅ │
├─────┼──────────────────────────────────────────────────┼─────────────────┼───────────────────┼──────────────────────┤
│ 3 │ enterprise/aibridged_integration_test.go │ enterprise_test │ aibridge.Provider │ ✅ │
├─────┼──────────────────────────────────────────────────┼─────────────────┼───────────────────┼──────────────────────┤
│ 4 │ aibridge/internal/integrationtest/setupbridge.go │ integrationtest │ aibridge.Provider │ ✅ │
├─────┼──────────────────────────────────────────────────┼─────────────────┼───────────────────┼──────────────────────┤
│ 5 │ aibridge/provider/anthropic_internal_test.go │ provider │ *Anthropic │ ❌ │
└─────┴──────────────────────────────────────────────────┴─────────────────┴───────────────────┴──────────────────────┘
Our existing testutil pkg is under internal path.
UPD:
- I included this fix for now, let me know if I need to revert it.
- Also Note we still have one
mustNewAnthropicduplicate helper in provider package. There is no good way to remove it. Happy to discuss this part.
| }, | ||
| expectKind: intercept.CredentialKindCentralized, | ||
| expectAuthHeader: "", | ||
| expectHint: "<aws chain credentials>", |
There was a problem hiding this comment.
Why this change? It was more clear before.
There was a problem hiding this comment.
Hint should be under 15 chars, according to DB constraint to prevent accidentally storing real credentials.
So this bug existed before.
But I think this bug wasn't released yet, probably I need to add test for it as well.
There was a problem hiding this comment.
Updated test here: 5ef859f to prevent it from happening in the future.
Yes, but I'll double check before merge, considering I made changes. I tested with only one AWS account. Both User (Base Creds) and Role in the same account. But it should work the same way for cross-account. From Coder PoV it doesn't matter, it only concerns TrustPolicy of IAM Role. |
|
/coder-agents-review |
There was a problem hiding this comment.
Round 2. All 8 R1 fixes verified by panel. All 7 contested findings evaluated and closed by panel vote.
R1 P2 fixes (CRF-2 region loss, CRF-3 misleading error) are solid. The ResolvedRegion fallback and neutral error message both address root causes. The new ARN validation (CRF-5 fix) covers standard, govcloud, and China partitions with clean test coverage. The AssumeRole caching, refresh, and error tests are genuine, including the atomic-counter verification that 5 retrievals produce exactly 1 STS call.
Contested findings: CRF-4 (naming) closed because the external-facing brand is "AI Gateway" and the session name is for CloudTrail consumers, not code readers. CRF-6 (PATCH clears RoleARN) closed; consistent with Region/Model as non-secret readable fields. CRF-7 (no ExternalId) closed; confused-deputy doesn't apply to first-party self-hosted. CRF-9, CRF-10 (comment style) closed by majority; the comments document non-obvious contracts that the code alone doesn't reveal.
4 new findings below (3 P3, 1 Nit).
"This code earned [silence]."
-- Hisoka, on the credential lifecycle design
🤖 This review was automatically generated with Coder Agents.
|
|
||
| base := &interceptionBase{} | ||
| opts, err := base.withAWSBedrockOptions(context.Background(), tt.cfg) | ||
| base := &interceptionBase{ |
There was a problem hiding this comment.
P3 [CRF-16] TestAWSBedrockValidation hardcodes a successful credential provider in every row, leaving the Retrieve error path in withAWSBedrockOptions (line 297) untested at this layer.
The old TestAWSBedrockCredentialChain tested the "no credential source configured" error path through withAWSBedrockOptions. That test was removed when credential resolution moved to buildBedrockCredentials, which has its own equivalent. But the wrapping error path in withAWSBedrockOptions ("resolve AWS credentials: %w") is now uncovered here. A test case with a Creds provider that returns an error on Retrieve would pin the new error message and ensure future refactors don't silently change the error surface.
(Bisky)
🤖
| return nil, "", xerrors.New("region is required to assume a role, but was not specified") | ||
| } | ||
|
|
||
| // The base identity signs Bedrock requests directly unless a target role is |
There was a problem hiding this comment.
P3 [CRF-17] The comment narrates what the if/else does ("The base identity signs requests directly unless a target RoleARN is configured..."). The code shows this: credsProvider := base.Credentials followed by if cfg.RoleARN != "" { ... }. The one useful observation, that base.Credentials from LoadDefaultConfig is already cache-wrapped so only the AssumeRoleProvider needs explicit wrapping, is buried in mechanism narration.
Suggested: // Only the AssumeRoleProvider needs cache-wrapping; base.Credentials from LoadDefaultConfig is already cached.
(Gon)
🤖
| return nil, xerrors.New("both access key and access key secret must be provided together") | ||
| // Neither set: SDK default credential chain resolves credentials. | ||
| default: | ||
| // Fail fast: ensure credentials can be resolved before signing. Served from |
There was a problem hiding this comment.
P3 [CRF-18] "Served from the shared cache on most requests (no network); on the cold or refresh path this performs the actual STS/IMDS call." This restates the caching behavior already documented on buildBedrockCredentials and on the withAWSBedrockOptions doc comment two lines above. The intent (fail fast before signing) fits in one line.
Suggested: // Fail fast: surface credential errors before signing.
(Gon)
🤖
| // so per-request credential retrieval is served from this cache rather than | ||
| // re-resolving (and re-assuming) on every request. No network call is made here: | ||
| // the base identity and any AssumeRole are resolved lazily on first retrieval. | ||
| func buildBedrockCredentials(ctx context.Context, cfg config.AWSBedrock) (aws.CredentialsProvider, string, error) { |
There was a problem hiding this comment.
Nit [CRF-19] The bare string second return is ambiguous. Callers must read the body to discover it is the SDK-resolved region. Named returns would clarify:
func buildBedrockCredentials(ctx context.Context, cfg config.AWSBedrock) (creds aws.CredentialsProvider, resolvedRegion string, err error) {(Gon)
🤖
Support IAM role assumption for AWS Bedrock in AI Bridge
Summary
Implements https://linear.app/codercom/issue/AIGOV-371/support-dynamic-bedrock-assumerole-across-aws-accounts-for-ai-gateway
A Bedrock provider can now be configured with an IAM role to assume. Before calling Bedrock, the gateway assumes that role via STS and signs requests with the resulting temporary credentials. Whether the role lives in the same account or another one is entirely a matter of the role's trust policy.
Problem
Many organizations prohibit long-lived AWS access keys and expect workloads to authenticate through assumed IAM roles instead. A common case is an organization that runs Bedrock across several AWS accounts, one per business unit, and needs each unit's usage billed to its own account by assuming a role there. AI Bridge previously authenticated a Bedrock provider only with static keys or the gateway's own ambient AWS identity, which is shared by every provider, with no way to assume a role. These deployments had no clean path.
How it works
When a provider is configured with a role ARN, the gateway uses its base identity to assume that role via STS and signs Bedrock requests with the temporary credentials it returns. The base identity is whatever the AWS default credential chain resolves, IRSA, EKS Pod Identity, EC2 Instance Profile, or static keys.
Credentials are resolved once when the provider is set up and are then cached and rotated, so individual requests are served from the cache rather than triggering a new STS call. A deployment that needs several roles configures several providers, each pointing at its own role.
Configuration
The role ARN is part of the Bedrock provider settings and is set through the AI provider API. It is optional: a provider with no role ARN behaves exactly as before.
Scope and trade-offs
Follow-up PR: #26578