mirror of
https://github.com/coder/coder.git
synced 2026-09-24 15:04:27 +08:00
feat: support cross-account Bedrock AssumeRole in AI Bridge (#26527)
# 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 - This PR is backend only. The settings UI for the role ARN ships in a follow-up. - Configuration is not exposed through environment variables. Environment-based provider configuration is being phased out in favor of database-managed providers, so the role ARN is intentionally database and API only. Follow-up PR: https://github.com/coder/coder/pull/26578
This commit is contained in:
+5
-3
@@ -175,7 +175,7 @@ func BuildProviders(ctx context.Context, db database.Store, cfg codersdk.AIBridg
|
||||
if row.Enabled {
|
||||
enabledCount++
|
||||
}
|
||||
prov, err := buildAIProviderFromRow(row, keysByProvider[row.ID], cfg, metrics)
|
||||
prov, err := buildAIProviderFromRow(ctx, row, keysByProvider[row.ID], cfg, metrics)
|
||||
if err != nil {
|
||||
outcome.Status = aibridged.ProviderStatusError
|
||||
outcome.Err = err
|
||||
@@ -210,6 +210,7 @@ func BuildProviders(ctx context.Context, db database.Store, cfg codersdk.AIBridg
|
||||
// Disabled: true; settings decode, key loading, and credential checks
|
||||
// are skipped because the provider will never call upstream.
|
||||
func buildAIProviderFromRow(
|
||||
ctx context.Context,
|
||||
row database.AIProvider,
|
||||
keys []database.AIProviderKey,
|
||||
cfg codersdk.AIBridgeConfig,
|
||||
@@ -284,14 +285,14 @@ func buildAIProviderFromRow(
|
||||
return nil, xerrors.Errorf("anthropic key pool: %w", err)
|
||||
}
|
||||
}
|
||||
return aibridge.NewAnthropicProvider(aibridge.AnthropicConfig{
|
||||
return aibridge.NewAnthropicProvider(ctx, aibridge.AnthropicConfig{
|
||||
Name: row.Name,
|
||||
BaseURL: row.BaseUrl,
|
||||
KeyPool: pool,
|
||||
APIDumpDir: dumpDir,
|
||||
CircuitBreaker: cbCfg,
|
||||
SendActorHeaders: sendActorHeaders,
|
||||
}, bedrock), nil
|
||||
}, bedrock)
|
||||
|
||||
case database.AIProviderTypeCopilot:
|
||||
// Copilot is always BYOK; the per-user token is supplied on each
|
||||
@@ -348,6 +349,7 @@ func bedrockConfigFromRow(row database.AIProvider, settings codersdk.AIProviderS
|
||||
AccessKeySecret: accessKeySecret,
|
||||
Model: bedrockSettings.Model,
|
||||
SmallFastModel: bedrockSettings.SmallFastModel,
|
||||
RoleARN: bedrockSettings.RoleARN,
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -267,6 +267,7 @@ func TestBuildProviders(t *testing.T) {
|
||||
Name: "anthropic-bedrock",
|
||||
BaseUrl: "https://bedrock-runtime.us-west-2.amazonaws.com/",
|
||||
}
|
||||
roleARN := "arn:aws:iam::123456789012:role/BedrockRole"
|
||||
settings := codersdk.AIProviderSettings{
|
||||
Bedrock: &codersdk.AIProviderBedrockSettings{
|
||||
Region: "us-west-2",
|
||||
@@ -274,6 +275,7 @@ func TestBuildProviders(t *testing.T) {
|
||||
AccessKeySecret: &secret,
|
||||
Model: model,
|
||||
SmallFastModel: smallModel,
|
||||
RoleARN: roleARN,
|
||||
},
|
||||
}
|
||||
got := bedrockConfigFromRow(row, settings)
|
||||
@@ -284,6 +286,7 @@ func TestBuildProviders(t *testing.T) {
|
||||
assert.Equal(t, secret, got.AccessKeySecret)
|
||||
assert.Equal(t, model, got.Model)
|
||||
assert.Equal(t, smallModel, got.SmallFastModel)
|
||||
assert.Equal(t, roleARN, got.RoleARN)
|
||||
})
|
||||
|
||||
t.Run("BedrockSettingsEmpty", func(t *testing.T) {
|
||||
|
||||
@@ -741,7 +741,7 @@ func TestBuildAIProviderFromRowSetsAPIDumpDir(t *testing.T) {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
provider, err := buildAIProviderFromRow(tt.row, nil, codersdk.AIBridgeConfig{
|
||||
provider, err := buildAIProviderFromRow(t.Context(), tt.row, nil, codersdk.AIBridgeConfig{
|
||||
AllowBYOK: serpent.Bool(true),
|
||||
APIDumpDir: serpent.String(dumpDir),
|
||||
}, nil)
|
||||
@@ -755,7 +755,7 @@ func TestBuildAIProviderFromRowSetsAPIDumpDir(t *testing.T) {
|
||||
func TestBuildAIProviderFromRowBedrockWithoutSettings(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
_, err := buildAIProviderFromRow(database.AIProvider{
|
||||
_, err := buildAIProviderFromRow(t.Context(), database.AIProvider{
|
||||
Enabled: true,
|
||||
Type: database.AIProviderTypeBedrock,
|
||||
Name: "bedrock-no-settings",
|
||||
|
||||
Reference in New Issue
Block a user