From a0446bcdbddf99f25639a7513c358bd8b50936ff Mon Sep 17 00:00:00 2001 From: Yevhenii Shcherbina Date: Tue, 1 Sep 2026 14:13:56 +0000 Subject: [PATCH 01/18] fix: resolve Bedrock application inference profile ARNs Application inference profile ARNs are opaque, so Bedrock capability detection matched nothing and adaptive-thinking conversion shaped requests for the wrong model. Resolve the ARN through GetInferenceProfile and use the underlying model ID for capabilities, usage records, pricing, and metrics, while still invoking the profile so AWS attributes spend. Resolution only runs for application inference profile ARNs, so deployments configured with plain model IDs need no extra permission. --- aibridge/intercept/messages/base.go | 64 ++++- .../intercept/messages/base_internal_test.go | 101 ++++++- aibridge/intercept/messages/blocking.go | 1 + aibridge/intercept/messages/streaming.go | 1 + aibridge/provider/anthropic.go | 43 ++- .../provider/bedrock_inference_profile.go | 116 ++++++++ ...bedrock_inference_profile_internal_test.go | 267 ++++++++++++++++++ docs/ai-coder/ai-gateway/providers.md | 14 + go.mod | 1 + go.sum | 2 + 10 files changed, 602 insertions(+), 8 deletions(-) create mode 100644 aibridge/provider/bedrock_inference_profile.go create mode 100644 aibridge/provider/bedrock_inference_profile_internal_test.go diff --git a/aibridge/intercept/messages/base.go b/aibridge/intercept/messages/base.go index 3b77f752a62..5c1a71aa3bc 100644 --- a/aibridge/intercept/messages/base.go +++ b/aibridge/intercept/messages/base.go @@ -72,6 +72,44 @@ var bedrockSupportedBetaFlags = map[string]bool{ type BedrockRuntime struct { Cfg aibconfig.AWSBedrock Creds aws.CredentialsProvider + // ResolvedModel and ResolvedSmallFastModel are the Bedrock model IDs behind + // the configured identifiers. They differ from the configured values only + // when those are application inference profile ARNs, which are opaque and + // must be resolved through AWS. Empty values fall back to the configured + // identifier. + ResolvedModel string + ResolvedSmallFastModel string +} + +// InvocationModel returns the identifier to send upstream as the model. This is +// the configured value, which may be an application inference profile ARN: AWS +// attributes spend to the profile only when the profile itself is invoked. +func (b *BedrockRuntime) InvocationModel() string { + return b.Cfg.Model +} + +// SmallFastInvocationModel is [BedrockRuntime.InvocationModel] for the +// small/fast model. +func (b *BedrockRuntime) SmallFastInvocationModel() string { + return b.Cfg.SmallFastModel +} + +// ModelID returns the Bedrock model ID behind the configured identifier. Model +// capabilities, usage records, pricing, and metrics all key off this rather +// than the invocation target. +func (b *BedrockRuntime) ModelID() string { + if b.ResolvedModel != "" { + return b.ResolvedModel + } + return b.Cfg.Model +} + +// SmallFastModelID is [BedrockRuntime.ModelID] for the small/fast model. +func (b *BedrockRuntime) SmallFastModelID() string { + if b.ResolvedSmallFastModel != "" { + return b.ResolvedSmallFastModel + } + return b.Cfg.SmallFastModel } type interceptionBase struct { @@ -86,6 +124,11 @@ type interceptionBase struct { // clientHeaders are the original HTTP headers from the client request. clientHeaders http.Header + // smallFast classifies the model the client requested. It is captured at + // construction because the Bedrock InvokeModel remap overwrites the model in + // the request payload. + smallFast bool + logger slog.Logger tracer trace.Tracer @@ -169,11 +212,10 @@ func (i *interceptionBase) Model() string { // passthrough, non-Bedrock providers) returns the model the client sent in // the body. if i.isBedrockInvokeModel() { - model := i.bedrock.Cfg.Model if i.isSmallFastModel() { - model = i.bedrock.Cfg.SmallFastModel + return i.bedrock.SmallFastModelID() } - return model + return i.bedrock.ModelID() } return i.reqPayload.model() @@ -266,7 +308,12 @@ func (*interceptionBase) extractModelThoughts(msg *anthropic.Message) []*recorde // See `ANTHROPIC_SMALL_FAST_MODEL`: https://docs.anthropic.com/en/docs/claude-code/settings#environment-variables // https://docs.claude.com/en/docs/claude-code/costs#background-token-usage func (i *interceptionBase) isSmallFastModel() bool { - return strings.Contains(i.reqPayload.model(), "haiku") + return i.smallFast +} + +// isSmallFastModel reports whether the client requested a small/fast model. +func isSmallFastModel(model string) bool { + return strings.Contains(model, "haiku") } // newMessagesService builds the SDK service used for upstream calls. @@ -415,13 +462,20 @@ func (i *interceptionBase) withBedrockMantleOptions(ctx context.Context) ([]opti // Anthropics' model names. It also converts adaptive thinking to enabled with a budget for models that // don't support adaptive thinking natively, or enabled thinking to adaptive for models that only support // adaptive. +// +// The request carries the invocation target, which may be an application +// inference profile ARN, while capability decisions use the model ID behind it. func (i *interceptionBase) augmentRequestForBedrockInvokeModel() { if i.bedrock == nil { return } model := i.Model() - updated, err := i.reqPayload.withModel(model) + invocationModel := i.bedrock.InvocationModel() + if i.isSmallFastModel() { + invocationModel = i.bedrock.SmallFastInvocationModel() + } + updated, err := i.reqPayload.withModel(invocationModel) if err != nil { i.logger.Warn(context.Background(), "failed to set model in request payload for Bedrock", slog.Error(err)) return diff --git a/aibridge/intercept/messages/base_internal_test.go b/aibridge/intercept/messages/base_internal_test.go index b5b400dd92c..2920a003819 100644 --- a/aibridge/intercept/messages/base_internal_test.go +++ b/aibridge/intercept/messages/base_internal_test.go @@ -209,6 +209,84 @@ func TestAWSBedrockOptionsRequireRuntime(t *testing.T) { require.Contains(t, err.Error(), "nil bedrock runtime") } +// TestModelForBedrockInvokeModel covers the split between the invocation target +// and the model identity used for capabilities, usage records, and metrics. +func TestModelForBedrockInvokeModel(t *testing.T) { + t.Parallel() + + const ( + profileARN = "arn:aws:bedrock:eu-west-2:123456789012:application-inference-profile/46u2vhiyo6z5" + smallFastProfileARN = "arn:aws:bedrock:eu-west-2:123456789012:application-inference-profile/8x1qk20fzp3r" + ) + + runtime := &BedrockRuntime{ + Cfg: config.AWSBedrock{ + Model: profileARN, + SmallFastModel: smallFastProfileARN, + }, + ResolvedModel: "anthropic.claude-opus-4-8", + ResolvedSmallFastModel: "anthropic.claude-haiku-4-5", + } + + tests := []struct { + name string + smallFast bool + expectModel string + expectInvocationID string + }{ + { + name: "primary model", + expectModel: "anthropic.claude-opus-4-8", + expectInvocationID: profileARN, + }, + { + name: "small fast model", + smallFast: true, + expectModel: "anthropic.claude-haiku-4-5", + expectInvocationID: smallFastProfileARN, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + i := &interceptionBase{ + reqPayload: mustMessagesPayload(t, `{"model":"claude-opus-4-8","max_tokens":10000}`), + bedrock: runtime, + smallFast: tt.smallFast, + logger: slog.Make(), + } + + require.Equal(t, tt.expectModel, i.Model()) + + i.augmentRequestForBedrockInvokeModel() + require.Equal(t, tt.expectInvocationID, gjson.GetBytes(i.reqPayload, "model").String()) + // The remap must not change how the interception identifies itself. + require.Equal(t, tt.expectModel, i.Model()) + }) + } +} + +// TestModelFallsBackToConfiguredIdentifier covers Bedrock providers configured +// with plain model IDs, which are never resolved. +func TestModelFallsBackToConfiguredIdentifier(t *testing.T) { + t.Parallel() + + i := &interceptionBase{ + reqPayload: mustMessagesPayload(t, `{"model":"claude-opus-4-8","max_tokens":10000}`), + bedrock: &BedrockRuntime{ + Cfg: config.AWSBedrock{ + Model: "eu.anthropic.claude-opus-4-8", + SmallFastModel: "anthropic.claude-haiku-4-5", + }, + }, + logger: slog.Make(), + } + + require.Equal(t, "eu.anthropic.claude-opus-4-8", i.Model()) +} + func TestAccumulateUsage(t *testing.T) { t.Parallel() @@ -608,6 +686,7 @@ func TestAugmentRequestForBedrock_AdaptiveThinking(t *testing.T) { name string bedrockModel string + resolvedModel string // underlying model ID when bedrockModel is an application inference profile ARN requestBody string clientBetaFlags string @@ -760,6 +839,24 @@ func TestAugmentRequestForBedrock_AdaptiveThinking(t *testing.T) { requestBody: `{"max_tokens":10000,"thinking":{"type":"enabled","budget_tokens":8000}}`, expectThinkingType: "adaptive", }, + { + // Application inference profile ARNs are opaque, so capability + // detection runs against the model resolved through AWS. + name: "opaque_application_inference_profile_uses_resolved_model_for_enabled_thinking", + bedrockModel: "arn:aws:bedrock:eu-west-2:123456789012:application-inference-profile/46u2vhiyo6z5", + resolvedModel: "anthropic.claude-opus-4-8", + requestBody: `{"max_tokens":10000,"thinking":{"type":"enabled","budget_tokens":8000}}`, + expectThinkingType: "adaptive", + }, + { + name: "opaque_application_inference_profile_keeps_adaptive_thinking_and_output_config", + bedrockModel: "arn:aws:bedrock:eu-west-2:123456789012:application-inference-profile/46u2vhiyo6z5", + resolvedModel: "anthropic.claude-opus-4-8", + requestBody: `{"max_tokens":10000,"thinking":{"type":"adaptive"},"output_config":{"effort":"high"}}`, + expectThinkingType: "adaptive", + expectEffort: "high", + expectKeptFields: []string{"output_config"}, + }, { name: "opus_4_8_model_with_enabled_thinking_is_converted_to_adaptive_and_drops_budget", bedrockModel: "eu.anthropic.claude-opus-4-8", @@ -811,6 +908,7 @@ func TestAugmentRequestForBedrock_AdaptiveThinking(t *testing.T) { Model: tc.bedrockModel, SmallFastModel: "anthropic.claude-haiku-3-5", }, + ResolvedModel: tc.resolvedModel, }, clientHeaders: clientHeaders, logger: slog.Make(), @@ -832,7 +930,8 @@ func TestAugmentRequestForBedrock_AdaptiveThinking(t *testing.T) { require.Equal(t, tc.expectBudgetTokens, budgetTokens.Int()) } - // Model should always be set to the bedrock model. + // The request always targets the configured identifier, which may be + // an application inference profile ARN. require.Equal(t, tc.bedrockModel, gjson.GetBytes(i.reqPayload, "model").String()) // Verify expected fields are removed. diff --git a/aibridge/intercept/messages/blocking.go b/aibridge/intercept/messages/blocking.go index ecebb76981e..6af28204cdf 100644 --- a/aibridge/intercept/messages/blocking.go +++ b/aibridge/intercept/messages/blocking.go @@ -48,6 +48,7 @@ func NewBlockingInterceptor( bedrock: bedrock, clientHeaders: clientHeaders, tracer: tracer, + smallFast: isSmallFastModel(reqPayload.model()), }} } diff --git a/aibridge/intercept/messages/streaming.go b/aibridge/intercept/messages/streaming.go index 70ddc65bdb0..68c450f87b5 100644 --- a/aibridge/intercept/messages/streaming.go +++ b/aibridge/intercept/messages/streaming.go @@ -53,6 +53,7 @@ func NewStreamingInterceptor( bedrock: bedrock, clientHeaders: clientHeaders, tracer: tracer, + smallFast: isSmallFastModel(reqPayload.model()), }} } diff --git a/aibridge/provider/anthropic.go b/aibridge/provider/anthropic.go index 460b38102cd..8e0aa60d1d6 100644 --- a/aibridge/provider/anthropic.go +++ b/aibridge/provider/anthropic.go @@ -33,6 +33,22 @@ type Anthropic struct { bedrock *messages.BedrockRuntime } +// anthropicOption customizes provider construction. The type is unexported so +// the set of behaviors stays closed to this package. +type anthropicOption func(*anthropicOptions) + +type anthropicOptions struct { + resolveInferenceProfile inferenceProfileResolver +} + +// withInferenceProfileResolver overrides how application inference profile +// ARNs are resolved, so tests do not call AWS. +func withInferenceProfileResolver(resolve inferenceProfileResolver) anthropicOption { + return func(o *anthropicOptions) { + o.resolveInferenceProfile = resolve + } +} + const routeMessages = "/v1/messages" // https://docs.anthropic.com/en/api/messages var anthropicOpenErrorResponse = func() []byte { @@ -51,7 +67,12 @@ var anthropicIsFailure = func(statusCode int) bool { return circuitbreaker.DefaultIsFailure(statusCode) } -func NewAnthropic(ctx context.Context, cfg config.Anthropic, bedrockCfg *config.AWSBedrock) (*Anthropic, error) { +func NewAnthropic(ctx context.Context, cfg config.Anthropic, bedrockCfg *config.AWSBedrock, opts ...anthropicOption) (*Anthropic, error) { + options := anthropicOptions{resolveInferenceProfile: resolveInferenceProfile} + for _, opt := range opts { + opt(&options) + } + if cfg.Name == "" { cfg.Name = config.ProviderAnthropic } @@ -82,7 +103,25 @@ func NewAnthropic(ctx context.Context, cfg config.Anthropic, bedrockCfg *config. if err := runtimeCfg.Validate(); err != nil { return nil, xerrors.Errorf("bedrock config: %w", err) } - bedrock = &messages.BedrockRuntime{Cfg: runtimeCfg, Creds: creds} + + // Resolution only calls AWS for application inference profile ARNs, so + // deployments configured with plain model IDs need no extra permission. + // A failure here fails provider construction: serving the provider with + // an unresolved profile would silently misshape every request, which + // Bedrock rejects outright on models that only accept adaptive thinking. + resolveCtx, cancel := context.WithTimeout(ctx, inferenceProfileResolutionTimeout) + defer cancel() + model, smallFastModel, err := resolveBedrockModels(resolveCtx, runtimeCfg, creds, options.resolveInferenceProfile) + if err != nil { + return nil, xerrors.Errorf("resolve bedrock models: %w", err) + } + + bedrock = &messages.BedrockRuntime{ + Cfg: runtimeCfg, + Creds: creds, + ResolvedModel: model, + ResolvedSmallFastModel: smallFastModel, + } } return &Anthropic{ diff --git a/aibridge/provider/bedrock_inference_profile.go b/aibridge/provider/bedrock_inference_profile.go new file mode 100644 index 00000000000..4e409a7076b --- /dev/null +++ b/aibridge/provider/bedrock_inference_profile.go @@ -0,0 +1,116 @@ +package provider + +import ( + "context" + "strings" + "time" + + "github.com/aws/aws-sdk-go-v2/aws" + "github.com/aws/aws-sdk-go-v2/aws/arn" + "github.com/aws/aws-sdk-go-v2/service/bedrock" + "golang.org/x/xerrors" + + "github.com/coder/coder/v2/aibridge/config" +) + +// bedrockService is the ARN service namespace for Amazon Bedrock resources. +const bedrockService = "bedrock" + +// applicationInferenceProfileResourceType is the ARN resource type of an +// application inference profile, the AWS-native mechanism for attributing +// Bedrock spend to a team or workload via cost allocation tags. +const applicationInferenceProfileResourceType = "application-inference-profile" + +// inferenceProfileResolver resolves an application inference profile ARN to the +// Bedrock model ID it wraps. +type inferenceProfileResolver func(ctx context.Context, cfg config.AWSBedrock, creds aws.CredentialsProvider, profileARN string) (string, error) + +// inferenceProfileResolutionTimeout bounds the Bedrock control-plane calls made +// while constructing a provider, which also cover the first credential +// resolution (STS/IRSA). +const inferenceProfileResolutionTimeout = 30 * time.Second + +// isApplicationInferenceProfileARN reports whether model is an application +// inference profile ARN. +// +// Plain model IDs and system-defined inference profile ARNs both carry the +// model ID in the identifier itself, so only application inference profiles, +// whose identifier is opaque, need resolving through AWS. +func isApplicationInferenceProfileARN(model string) bool { + parsed, err := arn.Parse(model) + if err != nil || parsed.Service != bedrockService { + return false + } + resourceType, _, ok := strings.Cut(parsed.Resource, "/") + return ok && resourceType == applicationInferenceProfileResourceType +} + +// resolveInferenceProfile returns the Bedrock model ID behind an application +// inference profile ARN. +// +// The caller's credentials sign the call, so the required +// bedrock:GetInferenceProfile permission belongs to the identity that already +// invokes Bedrock, including any role assumed via config.AWSBedrock.RoleARN. +func resolveInferenceProfile(ctx context.Context, cfg config.AWSBedrock, creds aws.CredentialsProvider, profileARN string) (string, error) { + client := bedrock.NewFromConfig(aws.Config{ + Region: cfg.Region, + Credentials: creds, + }) + + out, err := client.GetInferenceProfile(ctx, &bedrock.GetInferenceProfileInput{ + InferenceProfileIdentifier: aws.String(profileARN), + }) + if err != nil { + return "", xerrors.Errorf("get inference profile %q (requires the %s:GetInferenceProfile permission): %w", profileARN, bedrockService, err) + } + if len(out.Models) == 0 || out.Models[0].ModelArn == nil { + return "", xerrors.Errorf("inference profile %q references no model", profileARN) + } + + modelARN := *out.Models[0].ModelArn + model, err := modelIDFromARN(modelARN) + if err != nil { + return "", xerrors.Errorf("inference profile %q: %w", profileARN, err) + } + return model, nil +} + +// modelIDFromARN extracts the model ID from the ARN an inference profile +// points at. The ARN is either a foundation model +// (arn:aws:bedrock:{region}::foundation-model/{model}) or a system-defined +// inference profile (arn:aws:bedrock:{region}:{account}:inference-profile/{model}), +// and both carry the model ID as the resource identifier. +func modelIDFromARN(modelARN string) (string, error) { + parsed, err := arn.Parse(modelARN) + if err != nil { + return "", xerrors.Errorf("parse model arn %q: %w", modelARN, err) + } + _, model, ok := strings.Cut(parsed.Resource, "/") + if !ok || model == "" { + return "", xerrors.Errorf("model arn %q has no model identifier", modelARN) + } + return model, nil +} + +// resolveBedrockModels resolves the configured model identifiers to the model +// IDs used for capability detection, usage recording, and pricing. Identifiers +// that are not application inference profile ARNs are returned unchanged and +// cost no AWS call. +func resolveBedrockModels(ctx context.Context, cfg config.AWSBedrock, creds aws.CredentialsProvider, resolve inferenceProfileResolver) (model, smallFastModel string, err error) { + resolveOne := func(configured string) (string, error) { + if !isApplicationInferenceProfileARN(configured) { + return configured, nil + } + return resolve(ctx, cfg, creds, configured) + } + + model, err = resolveOne(cfg.Model) + if err != nil { + return "", "", xerrors.Errorf("resolve model: %w", err) + } + smallFastModel, err = resolveOne(cfg.SmallFastModel) + if err != nil { + return "", "", xerrors.Errorf("resolve small fast model: %w", err) + } + return model, smallFastModel, nil +} diff --git a/aibridge/provider/bedrock_inference_profile_internal_test.go b/aibridge/provider/bedrock_inference_profile_internal_test.go new file mode 100644 index 00000000000..145cf01088d --- /dev/null +++ b/aibridge/provider/bedrock_inference_profile_internal_test.go @@ -0,0 +1,267 @@ +package provider + +import ( + "context" + "testing" + + "github.com/aws/aws-sdk-go-v2/aws" + "github.com/stretchr/testify/require" + "golang.org/x/xerrors" + + "github.com/coder/coder/v2/aibridge/config" +) + +func TestIsApplicationInferenceProfileARN(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + model string + want bool + }{ + { + name: "application inference profile arn", + model: "arn:aws:bedrock:eu-west-2:123456789012:application-inference-profile/46u2vhiyo6z5", + want: true, + }, + { + name: "plain model id", + model: "anthropic.claude-opus-4-8", + want: false, + }, + { + name: "regional model id", + model: "eu.anthropic.claude-opus-4-8", + want: false, + }, + { + // System-defined inference profiles carry the model ID, so they need + // no resolution and must not require the extra AWS permission. + name: "system defined inference profile arn", + model: "arn:aws:bedrock:us-east-1:123456789012:inference-profile/us.anthropic.claude-opus-4-8", + want: false, + }, + { + name: "foundation model arn", + model: "arn:aws:bedrock:us-east-1::foundation-model/anthropic.claude-opus-4-8", + want: false, + }, + { + name: "non bedrock arn", + model: "arn:aws:iam::123456789012:role/BedrockRole", + want: false, + }, + { + name: "malformed arn", + model: "arn:aws:bedrock:broken", + want: false, + }, + { + name: "resource without type separator", + model: "arn:aws:bedrock:us-east-1:123456789012:application-inference-profile", + want: false, + }, + { + name: "empty", + model: "", + want: false, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + require.Equal(t, tt.want, isApplicationInferenceProfileARN(tt.model)) + }) + } +} + +func TestModelIDFromARN(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + modelARN string + want string + errorMsg string + }{ + { + name: "foundation model", + modelARN: "arn:aws:bedrock:us-east-1::foundation-model/anthropic.claude-opus-4-8", + want: "anthropic.claude-opus-4-8", + }, + { + name: "system defined inference profile", + modelARN: "arn:aws:bedrock:us-east-1:123456789012:inference-profile/us.anthropic.claude-opus-4-8", + want: "us.anthropic.claude-opus-4-8", + }, + { + name: "not an arn", + modelARN: "anthropic.claude-opus-4-8", + errorMsg: "parse model arn", + }, + { + name: "no model identifier", + modelARN: "arn:aws:bedrock:us-east-1::foundation-model", + errorMsg: "has no model identifier", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + got, err := modelIDFromARN(tt.modelARN) + if tt.errorMsg != "" { + require.Error(t, err) + require.Contains(t, err.Error(), tt.errorMsg) + return + } + require.NoError(t, err) + require.Equal(t, tt.want, got) + }) + } +} + +func TestResolveBedrockModels(t *testing.T) { + t.Parallel() + + const ( + profileARN = "arn:aws:bedrock:eu-west-2:123456789012:application-inference-profile/46u2vhiyo6z5" + smallFastProfileARN = "arn:aws:bedrock:eu-west-2:123456789012:application-inference-profile/8x1qk20fzp3r" + ) + + t.Run("plain model ids are not resolved", func(t *testing.T) { + t.Parallel() + + cfg := config.AWSBedrock{ + Model: "eu.anthropic.claude-opus-4-8", + SmallFastModel: "anthropic.claude-haiku-4-5", + } + resolve := func(context.Context, config.AWSBedrock, aws.CredentialsProvider, string) (string, error) { + t.Error("resolver called for a plain model id") + return "", nil + } + + model, smallFastModel, err := resolveBedrockModels(context.Background(), cfg, nil, resolve) + require.NoError(t, err) + require.Equal(t, cfg.Model, model) + require.Equal(t, cfg.SmallFastModel, smallFastModel) + }) + + t.Run("profile arns are resolved", func(t *testing.T) { + t.Parallel() + + cfg := config.AWSBedrock{ + Model: profileARN, + SmallFastModel: smallFastProfileARN, + } + resolved := map[string]string{ + profileARN: "anthropic.claude-opus-4-8", + smallFastProfileARN: "anthropic.claude-haiku-4-5", + } + resolve := func(_ context.Context, _ config.AWSBedrock, _ aws.CredentialsProvider, arn string) (string, error) { + return resolved[arn], nil + } + + model, smallFastModel, err := resolveBedrockModels(context.Background(), cfg, nil, resolve) + require.NoError(t, err) + require.Equal(t, "anthropic.claude-opus-4-8", model) + require.Equal(t, "anthropic.claude-haiku-4-5", smallFastModel) + }) + + t.Run("small fast model resolves independently", func(t *testing.T) { + t.Parallel() + + cfg := config.AWSBedrock{ + Model: "eu.anthropic.claude-opus-4-8", + SmallFastModel: smallFastProfileARN, + } + var calls []string + resolve := func(_ context.Context, _ config.AWSBedrock, _ aws.CredentialsProvider, arn string) (string, error) { + calls = append(calls, arn) + return "anthropic.claude-haiku-4-5", nil + } + + model, smallFastModel, err := resolveBedrockModels(context.Background(), cfg, nil, resolve) + require.NoError(t, err) + require.Equal(t, []string{smallFastProfileARN}, calls) + require.Equal(t, cfg.Model, model) + require.Equal(t, "anthropic.claude-haiku-4-5", smallFastModel) + }) + + t.Run("resolution failure is returned", func(t *testing.T) { + t.Parallel() + + cfg := config.AWSBedrock{ + Model: profileARN, + SmallFastModel: "anthropic.claude-haiku-4-5", + } + resolve := func(context.Context, config.AWSBedrock, aws.CredentialsProvider, string) (string, error) { + return "", xerrors.New("AccessDeniedException") + } + + _, _, err := resolveBedrockModels(context.Background(), cfg, nil, resolve) + require.Error(t, err) + require.Contains(t, err.Error(), "resolve model") + require.Contains(t, err.Error(), "AccessDeniedException") + }) +} + +func TestNewAnthropic_InferenceProfileResolution(t *testing.T) { + t.Parallel() + + const profileARN = "arn:aws:bedrock:eu-west-2:123456789012:application-inference-profile/46u2vhiyo6z5" + + bedrockCfg := func(model string) *config.AWSBedrock { + return &config.AWSBedrock{ + Region: "eu-west-2", + AccessKey: "test-key", + AccessKeySecret: "test-secret", + Model: model, + SmallFastModel: "anthropic.claude-haiku-4-5", + } + } + + t.Run("resolved profile drives the model id", func(t *testing.T) { + t.Parallel() + + resolve := func(_ context.Context, _ config.AWSBedrock, _ aws.CredentialsProvider, arn string) (string, error) { + require.Equal(t, profileARN, arn) + return "anthropic.claude-opus-4-8", nil + } + + p, err := NewAnthropic(context.Background(), config.Anthropic{}, bedrockCfg(profileARN), withInferenceProfileResolver(resolve)) + require.NoError(t, err) + require.Equal(t, "anthropic.claude-opus-4-8", p.bedrock.ModelID()) + // The profile stays the invocation target so AWS attributes spend to it. + require.Equal(t, profileARN, p.bedrock.InvocationModel()) + require.Equal(t, "anthropic.claude-haiku-4-5", p.bedrock.SmallFastModelID()) + }) + + t.Run("failed resolution fails construction", func(t *testing.T) { + t.Parallel() + + resolve := func(context.Context, config.AWSBedrock, aws.CredentialsProvider, string) (string, error) { + return "", xerrors.New("AccessDeniedException") + } + + _, err := NewAnthropic(context.Background(), config.Anthropic{}, bedrockCfg(profileARN), withInferenceProfileResolver(resolve)) + require.ErrorContains(t, err, "resolve bedrock models") + }) + + t.Run("plain model id needs no resolution", func(t *testing.T) { + t.Parallel() + + resolve := func(context.Context, config.AWSBedrock, aws.CredentialsProvider, string) (string, error) { + t.Error("resolver called for a plain model id") + return "", nil + } + + p, err := NewAnthropic(context.Background(), config.Anthropic{}, bedrockCfg("eu.anthropic.claude-opus-4-8"), withInferenceProfileResolver(resolve)) + require.NoError(t, err) + require.Equal(t, "eu.anthropic.claude-opus-4-8", p.bedrock.ModelID()) + require.Equal(t, "eu.anthropic.claude-opus-4-8", p.bedrock.InvocationModel()) + }) +} diff --git a/docs/ai-coder/ai-gateway/providers.md b/docs/ai-coder/ai-gateway/providers.md index 47af0656143..b89de1554a4 100644 --- a/docs/ai-coder/ai-gateway/providers.md +++ b/docs/ai-coder/ai-gateway/providers.md @@ -127,6 +127,20 @@ export ANTHROPIC_BEDROCK_MANTLE_BASE_URL="/api/v2/ai-gatewa export ANTHROPIC_AUTH_TOKEN="" ``` +#### Application inference profiles + +For InvokeModel, the **model** and **small fast model** identifiers can be +[application inference profile](https://docs.aws.amazon.com/bedrock/latest/userguide/cost-mgmt-application-inference-profiles.html) +ARNs, which attribute Bedrock spend to a team or workload through AWS cost +allocation tags. AI Gateway invokes the profile so AWS records the +attribution, and resolves the model behind it to shape requests correctly and +to price usage. + +Resolution calls `GetInferenceProfile`, so the identity the gateway uses must +also permit `bedrock:GetInferenceProfile` for the profile. Providers +configured with plain model identifiers do not need this permission. When +resolution fails, the provider is not served and the failure is logged. + #### AWS credentials Do not attach API keys to a Bedrock provider. diff --git a/go.mod b/go.mod index fd9f838814f..efd1965a629 100644 --- a/go.mod +++ b/go.mod @@ -528,6 +528,7 @@ require ( charm.land/fantasy v0.8.1 github.com/anthropics/anthropic-sdk-go v1.19.0 github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream v1.7.15 + github.com/aws/aws-sdk-go-v2/service/bedrock v1.70.0 github.com/aymanbagabas/go-udiff v0.4.1 github.com/brianvoe/gofakeit/v7 v7.16.0 github.com/coder/agentapi-sdk-go v0.0.0-20250505131810-560d1d88d225 diff --git a/go.sum b/go.sum index fae4d0d9854..98f649231bf 100644 --- a/go.sum +++ b/go.sum @@ -188,6 +188,8 @@ github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.8.1 h1:K0JsbZQj+1h208Ro1zH github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.8.1/go.mod h1:W3/vL6EtCIatICGy9ab29QhMuae+cOKPWcMxv02CO+Q= github.com/aws/aws-sdk-go-v2/internal/v4a v1.5.1 h1:yhw5KD1phVyP9vijxOUzDfEtJx+bt+L63k+VfuiYFAA= github.com/aws/aws-sdk-go-v2/internal/v4a v1.5.1/go.mod h1:ZW2e0d7DYlRxlS9hEiMXE47gTdX5KRN4byUiNbUpG+Q= +github.com/aws/aws-sdk-go-v2/service/bedrock v1.70.0 h1:aZqG2s7EoLfCi50CT/lcW8S7Yt3xSVfB1du9rk/e5yg= +github.com/aws/aws-sdk-go-v2/service/bedrock v1.70.0/go.mod h1:nKmZ+J5ZhnK908kqblT/LTOk2VSW1MVxg2Qnr1KxQ78= github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.13.19 h1:bAdDl/HkGCcGPoe25ToSHEw23VIxt6CT5fLcg111BKg= github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.13.19/go.mod h1:KaUzbLxv4CeSxh6ZCl9B4m7CuFenS8kUEaDs+f/DQr4= github.com/aws/aws-sdk-go-v2/service/internal/checksum v1.9.26 h1:Eflerh7atY6HN0yz60peNLOkJA2ZKUyYjZexMbqwMCE= From 57f583aaae940460ce040b089e87d48bc87253d4 Mon Sep 17 00:00:00 2001 From: Yevhenii Shcherbina Date: Tue, 1 Sep 2026 15:30:11 +0000 Subject: [PATCH 02/18] refactor(aibridge): name Bedrock model accessors by configured and resolved Drop the empty-value fallback so the runtime is always constructed with resolved model IDs, and name the accessors after what they return rather than how they are used. --- aibridge/intercept/messages/base.go | 75 ++++++++------- .../intercept/messages/base_internal_test.go | 93 ++++++++----------- aibridge/provider/anthropic.go | 7 +- ...bedrock_inference_profile_internal_test.go | 12 +-- 4 files changed, 88 insertions(+), 99 deletions(-) diff --git a/aibridge/intercept/messages/base.go b/aibridge/intercept/messages/base.go index 5c1a71aa3bc..8d481657b97 100644 --- a/aibridge/intercept/messages/base.go +++ b/aibridge/intercept/messages/base.go @@ -72,44 +72,49 @@ var bedrockSupportedBetaFlags = map[string]bool{ type BedrockRuntime struct { Cfg aibconfig.AWSBedrock Creds aws.CredentialsProvider - // ResolvedModel and ResolvedSmallFastModel are the Bedrock model IDs behind - // the configured identifiers. They differ from the configured values only - // when those are application inference profile ARNs, which are opaque and - // must be resolved through AWS. Empty values fall back to the configured - // identifier. - ResolvedModel string - ResolvedSmallFastModel string -} - -// InvocationModel returns the identifier to send upstream as the model. This is -// the configured value, which may be an application inference profile ARN: AWS -// attributes spend to the profile only when the profile itself is invoked. -func (b *BedrockRuntime) InvocationModel() string { + + resolvedModel string + resolvedSmallFastModel string +} + +// NewBedrockRuntime bundles the Bedrock config and credentials with the model +// IDs behind the configured identifiers. The resolved IDs differ from the +// configured ones only when those are application inference profile ARNs, which +// are opaque and must be resolved through AWS; every other identifier resolves +// to itself. Taking them as arguments keeps the runtime always resolved. +func NewBedrockRuntime(cfg aibconfig.AWSBedrock, creds aws.CredentialsProvider, resolvedModel, resolvedSmallFastModel string) *BedrockRuntime { + return &BedrockRuntime{ + Cfg: cfg, + Creds: creds, + resolvedModel: resolvedModel, + resolvedSmallFastModel: resolvedSmallFastModel, + } +} + +// ConfiguredModel returns the identifier the operator configured, which may be +// an application inference profile ARN. Requests carry it as the model because +// AWS attributes spend to a profile only when the profile itself is invoked. +func (b *BedrockRuntime) ConfiguredModel() string { return b.Cfg.Model } -// SmallFastInvocationModel is [BedrockRuntime.InvocationModel] for the +// ConfiguredSmallFastModel is [BedrockRuntime.ConfiguredModel] for the // small/fast model. -func (b *BedrockRuntime) SmallFastInvocationModel() string { +func (b *BedrockRuntime) ConfiguredSmallFastModel() string { return b.Cfg.SmallFastModel } -// ModelID returns the Bedrock model ID behind the configured identifier. Model -// capabilities, usage records, pricing, and metrics all key off this rather -// than the invocation target. -func (b *BedrockRuntime) ModelID() string { - if b.ResolvedModel != "" { - return b.ResolvedModel - } - return b.Cfg.Model +// ResolvedModel returns the Bedrock model ID behind the configured identifier. +// Model capabilities, usage records, pricing, and metrics all key off this +// rather than the configured identifier. +func (b *BedrockRuntime) ResolvedModel() string { + return b.resolvedModel } -// SmallFastModelID is [BedrockRuntime.ModelID] for the small/fast model. -func (b *BedrockRuntime) SmallFastModelID() string { - if b.ResolvedSmallFastModel != "" { - return b.ResolvedSmallFastModel - } - return b.Cfg.SmallFastModel +// ResolvedSmallFastModel is [BedrockRuntime.ResolvedModel] for the small/fast +// model. +func (b *BedrockRuntime) ResolvedSmallFastModel() string { + return b.resolvedSmallFastModel } type interceptionBase struct { @@ -213,9 +218,9 @@ func (i *interceptionBase) Model() string { // the body. if i.isBedrockInvokeModel() { if i.isSmallFastModel() { - return i.bedrock.SmallFastModelID() + return i.bedrock.ResolvedSmallFastModel() } - return i.bedrock.ModelID() + return i.bedrock.ResolvedModel() } return i.reqPayload.model() @@ -463,7 +468,7 @@ func (i *interceptionBase) withBedrockMantleOptions(ctx context.Context) ([]opti // don't support adaptive thinking natively, or enabled thinking to adaptive for models that only support // adaptive. // -// The request carries the invocation target, which may be an application +// The request carries the configured identifier, which may be an application // inference profile ARN, while capability decisions use the model ID behind it. func (i *interceptionBase) augmentRequestForBedrockInvokeModel() { if i.bedrock == nil { @@ -471,11 +476,11 @@ func (i *interceptionBase) augmentRequestForBedrockInvokeModel() { } model := i.Model() - invocationModel := i.bedrock.InvocationModel() + configuredModel := i.bedrock.ConfiguredModel() if i.isSmallFastModel() { - invocationModel = i.bedrock.SmallFastInvocationModel() + configuredModel = i.bedrock.ConfiguredSmallFastModel() } - updated, err := i.reqPayload.withModel(invocationModel) + updated, err := i.reqPayload.withModel(configuredModel) if err != nil { i.logger.Warn(context.Background(), "failed to set model in request payload for Bedrock", slog.Error(err)) return diff --git a/aibridge/intercept/messages/base_internal_test.go b/aibridge/intercept/messages/base_internal_test.go index 2920a003819..616fead2091 100644 --- a/aibridge/intercept/messages/base_internal_test.go +++ b/aibridge/intercept/messages/base_internal_test.go @@ -179,10 +179,7 @@ func TestAWSBedrockValidation(t *testing.T) { t.Parallel() base := &interceptionBase{ - bedrock: &BedrockRuntime{ - Cfg: tt.cfg, - Creds: credentials.NewStaticCredentialsProvider("test-key", "test-secret", ""), - }, + bedrock: NewBedrockRuntime(tt.cfg, credentials.NewStaticCredentialsProvider("test-key", "test-secret", ""), "", ""), } opts, err := base.withBedrockInvokeModelOptions(context.Background()) @@ -219,31 +216,27 @@ func TestModelForBedrockInvokeModel(t *testing.T) { smallFastProfileARN = "arn:aws:bedrock:eu-west-2:123456789012:application-inference-profile/8x1qk20fzp3r" ) - runtime := &BedrockRuntime{ - Cfg: config.AWSBedrock{ - Model: profileARN, - SmallFastModel: smallFastProfileARN, - }, - ResolvedModel: "anthropic.claude-opus-4-8", - ResolvedSmallFastModel: "anthropic.claude-haiku-4-5", - } + runtime := NewBedrockRuntime(config.AWSBedrock{ + Model: profileARN, + SmallFastModel: smallFastProfileARN, + }, nil, "anthropic.claude-opus-4-8", "anthropic.claude-haiku-4-5") tests := []struct { - name string - smallFast bool - expectModel string - expectInvocationID string + name string + smallFast bool + expectModel string + expectConfigured string }{ { - name: "primary model", - expectModel: "anthropic.claude-opus-4-8", - expectInvocationID: profileARN, + name: "primary model", + expectModel: "anthropic.claude-opus-4-8", + expectConfigured: profileARN, }, { - name: "small fast model", - smallFast: true, - expectModel: "anthropic.claude-haiku-4-5", - expectInvocationID: smallFastProfileARN, + name: "small fast model", + smallFast: true, + expectModel: "anthropic.claude-haiku-4-5", + expectConfigured: smallFastProfileARN, }, } @@ -261,26 +254,24 @@ func TestModelForBedrockInvokeModel(t *testing.T) { require.Equal(t, tt.expectModel, i.Model()) i.augmentRequestForBedrockInvokeModel() - require.Equal(t, tt.expectInvocationID, gjson.GetBytes(i.reqPayload, "model").String()) + require.Equal(t, tt.expectConfigured, gjson.GetBytes(i.reqPayload, "model").String()) // The remap must not change how the interception identifies itself. require.Equal(t, tt.expectModel, i.Model()) }) } } -// TestModelFallsBackToConfiguredIdentifier covers Bedrock providers configured -// with plain model IDs, which are never resolved. -func TestModelFallsBackToConfiguredIdentifier(t *testing.T) { +// TestModelForPlainBedrockModelID covers Bedrock providers configured with +// plain model IDs, which resolve to themselves. +func TestModelForPlainBedrockModelID(t *testing.T) { t.Parallel() i := &interceptionBase{ reqPayload: mustMessagesPayload(t, `{"model":"claude-opus-4-8","max_tokens":10000}`), - bedrock: &BedrockRuntime{ - Cfg: config.AWSBedrock{ - Model: "eu.anthropic.claude-opus-4-8", - SmallFastModel: "anthropic.claude-haiku-4-5", - }, - }, + bedrock: NewBedrockRuntime(config.AWSBedrock{ + Model: "eu.anthropic.claude-opus-4-8", + SmallFastModel: "anthropic.claude-haiku-4-5", + }, nil, "eu.anthropic.claude-opus-4-8", "anthropic.claude-haiku-4-5"), logger: slog.Make(), } @@ -901,15 +892,19 @@ func TestAugmentRequestForBedrock_AdaptiveThinking(t *testing.T) { } } + // Plain model IDs resolve to themselves; an application inference + // profile ARN resolves to the model behind it. + resolvedModel := tc.resolvedModel + if resolvedModel == "" { + resolvedModel = tc.bedrockModel + } + i := &interceptionBase{ reqPayload: mustMessagesPayload(t, tc.requestBody), - bedrock: &BedrockRuntime{ - Cfg: config.AWSBedrock{ - Model: tc.bedrockModel, - SmallFastModel: "anthropic.claude-haiku-3-5", - }, - ResolvedModel: tc.resolvedModel, - }, + bedrock: NewBedrockRuntime(config.AWSBedrock{ + Model: tc.bedrockModel, + SmallFastModel: "anthropic.claude-haiku-3-5", + }, nil, resolvedModel, "anthropic.claude-haiku-3-5"), clientHeaders: clientHeaders, logger: slog.Make(), } @@ -1258,14 +1253,11 @@ func TestBedrockMantleIsPassthrough(t *testing.T) { i := &interceptionBase{ reqPayload: mustMessagesPayload(t, `{"model":"anthropic.claude-opus-4-8","max_tokens":10000,"thinking":{"type":"adaptive"},"metadata":{"user_id":"u123"},"context_management":{"type":"auto"}}`), - bedrock: &BedrockRuntime{ - Cfg: config.AWSBedrock{ - Region: "us-east-1", - BaseURL: "https://bedrock-mantle.us-east-1.api.aws/anthropic", - Protocol: config.BedrockProtocolMantle, - }, - Creds: credentials.NewStaticCredentialsProvider("test-key", "test-secret", ""), - }, + bedrock: NewBedrockRuntime(config.AWSBedrock{ + Region: "us-east-1", + BaseURL: "https://bedrock-mantle.us-east-1.api.aws/anthropic", + Protocol: config.BedrockProtocolMantle, + }, credentials.NewStaticCredentialsProvider("test-key", "test-secret", ""), "", ""), logger: slog.Make(), } @@ -1316,10 +1308,7 @@ func TestAWSMantleOptionsValidation(t *testing.T) { t.Parallel() base := &interceptionBase{ - bedrock: &BedrockRuntime{ - Cfg: tt.cfg, - Creds: credentials.NewStaticCredentialsProvider("test-key", "test-secret", ""), - }, + bedrock: NewBedrockRuntime(tt.cfg, credentials.NewStaticCredentialsProvider("test-key", "test-secret", ""), "", ""), } opts, err := base.withBedrockMantleOptions(t.Context()) if tt.errorMsg != "" { diff --git a/aibridge/provider/anthropic.go b/aibridge/provider/anthropic.go index 8e0aa60d1d6..1de0890b878 100644 --- a/aibridge/provider/anthropic.go +++ b/aibridge/provider/anthropic.go @@ -116,12 +116,7 @@ func NewAnthropic(ctx context.Context, cfg config.Anthropic, bedrockCfg *config. return nil, xerrors.Errorf("resolve bedrock models: %w", err) } - bedrock = &messages.BedrockRuntime{ - Cfg: runtimeCfg, - Creds: creds, - ResolvedModel: model, - ResolvedSmallFastModel: smallFastModel, - } + bedrock = messages.NewBedrockRuntime(runtimeCfg, creds, model, smallFastModel) } return &Anthropic{ diff --git a/aibridge/provider/bedrock_inference_profile_internal_test.go b/aibridge/provider/bedrock_inference_profile_internal_test.go index 145cf01088d..40b07f5641e 100644 --- a/aibridge/provider/bedrock_inference_profile_internal_test.go +++ b/aibridge/provider/bedrock_inference_profile_internal_test.go @@ -234,10 +234,10 @@ func TestNewAnthropic_InferenceProfileResolution(t *testing.T) { p, err := NewAnthropic(context.Background(), config.Anthropic{}, bedrockCfg(profileARN), withInferenceProfileResolver(resolve)) require.NoError(t, err) - require.Equal(t, "anthropic.claude-opus-4-8", p.bedrock.ModelID()) - // The profile stays the invocation target so AWS attributes spend to it. - require.Equal(t, profileARN, p.bedrock.InvocationModel()) - require.Equal(t, "anthropic.claude-haiku-4-5", p.bedrock.SmallFastModelID()) + require.Equal(t, "anthropic.claude-opus-4-8", p.bedrock.ResolvedModel()) + // The profile stays the configured identifier so AWS attributes spend to it. + require.Equal(t, profileARN, p.bedrock.ConfiguredModel()) + require.Equal(t, "anthropic.claude-haiku-4-5", p.bedrock.ResolvedSmallFastModel()) }) t.Run("failed resolution fails construction", func(t *testing.T) { @@ -261,7 +261,7 @@ func TestNewAnthropic_InferenceProfileResolution(t *testing.T) { p, err := NewAnthropic(context.Background(), config.Anthropic{}, bedrockCfg("eu.anthropic.claude-opus-4-8"), withInferenceProfileResolver(resolve)) require.NoError(t, err) - require.Equal(t, "eu.anthropic.claude-opus-4-8", p.bedrock.ModelID()) - require.Equal(t, "eu.anthropic.claude-opus-4-8", p.bedrock.InvocationModel()) + require.Equal(t, "eu.anthropic.claude-opus-4-8", p.bedrock.ResolvedModel()) + require.Equal(t, "eu.anthropic.claude-opus-4-8", p.bedrock.ConfiguredModel()) }) } From 0e031e7599accf69e14a3f269305ffe3f80f3e18 Mon Sep 17 00:00:00 2001 From: Yevhenii Shcherbina Date: Tue, 1 Sep 2026 15:34:46 +0000 Subject: [PATCH 03/18] refactor(aibridge/intercept/messages): keep Model() structure unchanged --- aibridge/intercept/messages/base.go | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/aibridge/intercept/messages/base.go b/aibridge/intercept/messages/base.go index 8d481657b97..097d9c4125d 100644 --- a/aibridge/intercept/messages/base.go +++ b/aibridge/intercept/messages/base.go @@ -217,10 +217,11 @@ func (i *interceptionBase) Model() string { // passthrough, non-Bedrock providers) returns the model the client sent in // the body. if i.isBedrockInvokeModel() { + model := i.bedrock.ResolvedModel() if i.isSmallFastModel() { - return i.bedrock.ResolvedSmallFastModel() + model = i.bedrock.ResolvedSmallFastModel() } - return i.bedrock.ResolvedModel() + return model } return i.reqPayload.model() From 1fd6ec5f9eb69945573acd31658acf6307393ca3 Mon Sep 17 00:00:00 2001 From: Yevhenii Shcherbina Date: Tue, 1 Sep 2026 16:33:49 +0000 Subject: [PATCH 04/18] refactor(aibridge/intercept/messages): extract upstreamModel --- aibridge/intercept/messages/base.go | 18 +++++++++++++----- 1 file changed, 13 insertions(+), 5 deletions(-) diff --git a/aibridge/intercept/messages/base.go b/aibridge/intercept/messages/base.go index 097d9c4125d..eac11a90ac8 100644 --- a/aibridge/intercept/messages/base.go +++ b/aibridge/intercept/messages/base.go @@ -227,6 +227,18 @@ func (i *interceptionBase) Model() string { return i.reqPayload.model() } +// upstreamModel returns the identifier sent to Bedrock as the invocation +// target, which may be an application inference profile ARN. It is the +// counterpart to [interceptionBase.Model]: AWS bills and attributes what is +// invoked, while everything internal keys off the model behind it. +func (i *interceptionBase) upstreamModel() string { + model := i.bedrock.ConfiguredModel() + if i.isSmallFastModel() { + model = i.bedrock.ConfiguredSmallFastModel() + } + return model +} + func (i *interceptionBase) baseTraceAttributes(r *http.Request, streaming bool) []attribute.KeyValue { attrs := []attribute.KeyValue{ attribute.String(tracing.RequestPath, r.URL.Path), @@ -477,11 +489,7 @@ func (i *interceptionBase) augmentRequestForBedrockInvokeModel() { } model := i.Model() - configuredModel := i.bedrock.ConfiguredModel() - if i.isSmallFastModel() { - configuredModel = i.bedrock.ConfiguredSmallFastModel() - } - updated, err := i.reqPayload.withModel(configuredModel) + updated, err := i.reqPayload.withModel(i.upstreamModel()) if err != nil { i.logger.Warn(context.Background(), "failed to set model in request payload for Bedrock", slog.Error(err)) return From 7742db144714ee975b8f9eee842e81600900714e Mon Sep 17 00:00:00 2001 From: Yevhenii Shcherbina Date: Tue, 1 Sep 2026 16:48:35 +0000 Subject: [PATCH 05/18] refactor(aibridge/intercept/messages): read small/fast classification from the field --- aibridge/intercept/messages/base.go | 23 ++++++++----------- .../intercept/messages/base_internal_test.go | 8 +++---- aibridge/intercept/messages/blocking.go | 16 ++++++------- aibridge/intercept/messages/streaming.go | 18 +++++++-------- 4 files changed, 30 insertions(+), 35 deletions(-) diff --git a/aibridge/intercept/messages/base.go b/aibridge/intercept/messages/base.go index eac11a90ac8..70449c9a08f 100644 --- a/aibridge/intercept/messages/base.go +++ b/aibridge/intercept/messages/base.go @@ -129,10 +129,13 @@ type interceptionBase struct { // clientHeaders are the original HTTP headers from the client request. clientHeaders http.Header - // smallFast classifies the model the client requested. It is captured at - // construction because the Bedrock InvokeModel remap overwrites the model in - // the request payload. - smallFast bool + // isSmallFastModel reports whether the client requested a small/fast model + // (Haiku 3.5), which is optimized for tasks like code autocomplete and other + // small, quick operations. It is captured at construction because the Bedrock + // InvokeModel remap overwrites the model in the request payload. + // See `ANTHROPIC_SMALL_FAST_MODEL`: https://docs.anthropic.com/en/docs/claude-code/settings#environment-variables + // https://docs.claude.com/en/docs/claude-code/costs#background-token-usage + isSmallFastModel bool logger slog.Logger tracer trace.Tracer @@ -218,7 +221,7 @@ func (i *interceptionBase) Model() string { // the body. if i.isBedrockInvokeModel() { model := i.bedrock.ResolvedModel() - if i.isSmallFastModel() { + if i.isSmallFastModel { model = i.bedrock.ResolvedSmallFastModel() } return model @@ -233,7 +236,7 @@ func (i *interceptionBase) Model() string { // invoked, while everything internal keys off the model behind it. func (i *interceptionBase) upstreamModel() string { model := i.bedrock.ConfiguredModel() - if i.isSmallFastModel() { + if i.isSmallFastModel { model = i.bedrock.ConfiguredSmallFastModel() } return model @@ -321,14 +324,6 @@ func (*interceptionBase) extractModelThoughts(msg *anthropic.Message) []*recorde return thoughtRecords } -// IsSmallFastModel checks if the model is a small/fast model (Haiku 3.5). -// These models are optimized for tasks like code autocomplete and other small, quick operations. -// See `ANTHROPIC_SMALL_FAST_MODEL`: https://docs.anthropic.com/en/docs/claude-code/settings#environment-variables -// https://docs.claude.com/en/docs/claude-code/costs#background-token-usage -func (i *interceptionBase) isSmallFastModel() bool { - return i.smallFast -} - // isSmallFastModel reports whether the client requested a small/fast model. func isSmallFastModel(model string) bool { return strings.Contains(model, "haiku") diff --git a/aibridge/intercept/messages/base_internal_test.go b/aibridge/intercept/messages/base_internal_test.go index 616fead2091..9b553a17490 100644 --- a/aibridge/intercept/messages/base_internal_test.go +++ b/aibridge/intercept/messages/base_internal_test.go @@ -245,10 +245,10 @@ func TestModelForBedrockInvokeModel(t *testing.T) { t.Parallel() i := &interceptionBase{ - reqPayload: mustMessagesPayload(t, `{"model":"claude-opus-4-8","max_tokens":10000}`), - bedrock: runtime, - smallFast: tt.smallFast, - logger: slog.Make(), + reqPayload: mustMessagesPayload(t, `{"model":"claude-opus-4-8","max_tokens":10000}`), + bedrock: runtime, + isSmallFastModel: tt.smallFast, + logger: slog.Make(), } require.Equal(t, tt.expectModel, i.Model()) diff --git a/aibridge/intercept/messages/blocking.go b/aibridge/intercept/messages/blocking.go index 6af28204cdf..98d408bc7d8 100644 --- a/aibridge/intercept/messages/blocking.go +++ b/aibridge/intercept/messages/blocking.go @@ -41,14 +41,14 @@ func NewBlockingInterceptor( tracer trace.Tracer, ) *BlockingInterception { return &BlockingInterception{interceptionBase: interceptionBase{ - id: id, - reqPayload: reqPayload, - cfg: cfg, - cred: cred, - bedrock: bedrock, - clientHeaders: clientHeaders, - tracer: tracer, - smallFast: isSmallFastModel(reqPayload.model()), + id: id, + reqPayload: reqPayload, + cfg: cfg, + cred: cred, + bedrock: bedrock, + clientHeaders: clientHeaders, + tracer: tracer, + isSmallFastModel: isSmallFastModel(reqPayload.model()), }} } diff --git a/aibridge/intercept/messages/streaming.go b/aibridge/intercept/messages/streaming.go index 68c450f87b5..a526015f83c 100644 --- a/aibridge/intercept/messages/streaming.go +++ b/aibridge/intercept/messages/streaming.go @@ -46,14 +46,14 @@ func NewStreamingInterceptor( tracer trace.Tracer, ) *StreamingInterception { return &StreamingInterception{interceptionBase: interceptionBase{ - id: id, - reqPayload: reqPayload, - cfg: cfg, - cred: cred, - bedrock: bedrock, - clientHeaders: clientHeaders, - tracer: tracer, - smallFast: isSmallFastModel(reqPayload.model()), + id: id, + reqPayload: reqPayload, + cfg: cfg, + cred: cred, + bedrock: bedrock, + clientHeaders: clientHeaders, + tracer: tracer, + isSmallFastModel: isSmallFastModel(reqPayload.model()), }} } @@ -116,7 +116,7 @@ func (i *StreamingInterception) ProcessRequest(w http.ResponseWriter, r *http.Re } // Claude Code uses a "small/fast model" for certain tasks. - if !i.isSmallFastModel() { + if !i.isSmallFastModel { // Only inject tools into "actual" request. i.injectTools() } From 431d2e7a69f54a83da5c424a7298211ca4dc0e07 Mon Sep 17 00:00:00 2001 From: evgeniy-scherbina Date: Tue, 1 Sep 2026 13:12:04 -0400 Subject: [PATCH 06/18] docs: minor changes --- aibridge/intercept/messages/base.go | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/aibridge/intercept/messages/base.go b/aibridge/intercept/messages/base.go index 70449c9a08f..67b19b1bdbc 100644 --- a/aibridge/intercept/messages/base.go +++ b/aibridge/intercept/messages/base.go @@ -81,7 +81,7 @@ type BedrockRuntime struct { // IDs behind the configured identifiers. The resolved IDs differ from the // configured ones only when those are application inference profile ARNs, which // are opaque and must be resolved through AWS; every other identifier resolves -// to itself. Taking them as arguments keeps the runtime always resolved. +// to itself. func NewBedrockRuntime(cfg aibconfig.AWSBedrock, creds aws.CredentialsProvider, resolvedModel, resolvedSmallFastModel string) *BedrockRuntime { return &BedrockRuntime{ Cfg: cfg, @@ -231,9 +231,7 @@ func (i *interceptionBase) Model() string { } // upstreamModel returns the identifier sent to Bedrock as the invocation -// target, which may be an application inference profile ARN. It is the -// counterpart to [interceptionBase.Model]: AWS bills and attributes what is -// invoked, while everything internal keys off the model behind it. +// target, which may be an application inference profile ARN. func (i *interceptionBase) upstreamModel() string { model := i.bedrock.ConfiguredModel() if i.isSmallFastModel { From 54508c4825f4405daf3b4951563fb16e6019c49e Mon Sep 17 00:00:00 2001 From: Yevhenii Shcherbina Date: Tue, 1 Sep 2026 18:00:53 +0000 Subject: [PATCH 07/18] docs(aibridge/provider): clarify which Bedrock identifiers are resolved --- aibridge/provider/bedrock_inference_profile.go | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) diff --git a/aibridge/provider/bedrock_inference_profile.go b/aibridge/provider/bedrock_inference_profile.go index 4e409a7076b..cafa67425e5 100644 --- a/aibridge/provider/bedrock_inference_profile.go +++ b/aibridge/provider/bedrock_inference_profile.go @@ -31,11 +31,10 @@ type inferenceProfileResolver func(ctx context.Context, cfg config.AWSBedrock, c const inferenceProfileResolutionTimeout = 30 * time.Second // isApplicationInferenceProfileARN reports whether model is an application -// inference profile ARN. -// -// Plain model IDs and system-defined inference profile ARNs both carry the -// model ID in the identifier itself, so only application inference profiles, -// whose identifier is opaque, need resolving through AWS. +// inference profile ARN, whose identifier is opaque and must be resolved +// through AWS. Plain model IDs and system-defined inference profile ARNs, which +// AWS documents as {geoRegion}.{modelId}, embed the model ID and need no +// lookup. func isApplicationInferenceProfileARN(model string) bool { parsed, err := arn.Parse(model) if err != nil || parsed.Service != bedrockService { From e9a19164cbceafce4a4b54c0b6ff422851c5e58b Mon Sep 17 00:00:00 2001 From: Yevhenii Shcherbina Date: Wed, 2 Sep 2026 13:36:15 +0000 Subject: [PATCH 08/18] test(aibridge/provider): resolve inference profiles against a mock endpoint Drop the injected resolver in favor of the mock-endpoint pattern the Bedrock credential tests already use, so the AWS client, response decoding, and error wrapping are exercised. Building the client from the loaded AWS config also honors custom control-plane endpoints. --- aibridge/provider/anthropic.go | 25 +----- .../provider/bedrock_inference_profile.go | 19 ++--- ...bedrock_inference_profile_internal_test.go | 80 +++++++++++++------ 3 files changed, 68 insertions(+), 56 deletions(-) diff --git a/aibridge/provider/anthropic.go b/aibridge/provider/anthropic.go index 1de0890b878..bb5446ef494 100644 --- a/aibridge/provider/anthropic.go +++ b/aibridge/provider/anthropic.go @@ -33,22 +33,6 @@ type Anthropic struct { bedrock *messages.BedrockRuntime } -// anthropicOption customizes provider construction. The type is unexported so -// the set of behaviors stays closed to this package. -type anthropicOption func(*anthropicOptions) - -type anthropicOptions struct { - resolveInferenceProfile inferenceProfileResolver -} - -// withInferenceProfileResolver overrides how application inference profile -// ARNs are resolved, so tests do not call AWS. -func withInferenceProfileResolver(resolve inferenceProfileResolver) anthropicOption { - return func(o *anthropicOptions) { - o.resolveInferenceProfile = resolve - } -} - const routeMessages = "/v1/messages" // https://docs.anthropic.com/en/api/messages var anthropicOpenErrorResponse = func() []byte { @@ -67,12 +51,7 @@ var anthropicIsFailure = func(statusCode int) bool { return circuitbreaker.DefaultIsFailure(statusCode) } -func NewAnthropic(ctx context.Context, cfg config.Anthropic, bedrockCfg *config.AWSBedrock, opts ...anthropicOption) (*Anthropic, error) { - options := anthropicOptions{resolveInferenceProfile: resolveInferenceProfile} - for _, opt := range opts { - opt(&options) - } - +func NewAnthropic(ctx context.Context, cfg config.Anthropic, bedrockCfg *config.AWSBedrock) (*Anthropic, error) { if cfg.Name == "" { cfg.Name = config.ProviderAnthropic } @@ -111,7 +90,7 @@ func NewAnthropic(ctx context.Context, cfg config.Anthropic, bedrockCfg *config. // Bedrock rejects outright on models that only accept adaptive thinking. resolveCtx, cancel := context.WithTimeout(ctx, inferenceProfileResolutionTimeout) defer cancel() - model, smallFastModel, err := resolveBedrockModels(resolveCtx, runtimeCfg, creds, options.resolveInferenceProfile) + model, smallFastModel, err := resolveBedrockModels(resolveCtx, runtimeCfg, creds, resolveInferenceProfile) if err != nil { return nil, xerrors.Errorf("resolve bedrock models: %w", err) } diff --git a/aibridge/provider/bedrock_inference_profile.go b/aibridge/provider/bedrock_inference_profile.go index cafa67425e5..13ce1c56334 100644 --- a/aibridge/provider/bedrock_inference_profile.go +++ b/aibridge/provider/bedrock_inference_profile.go @@ -7,6 +7,7 @@ import ( "github.com/aws/aws-sdk-go-v2/aws" "github.com/aws/aws-sdk-go-v2/aws/arn" + awsconfig "github.com/aws/aws-sdk-go-v2/config" "github.com/aws/aws-sdk-go-v2/service/bedrock" "golang.org/x/xerrors" @@ -21,10 +22,6 @@ const bedrockService = "bedrock" // Bedrock spend to a team or workload via cost allocation tags. const applicationInferenceProfileResourceType = "application-inference-profile" -// inferenceProfileResolver resolves an application inference profile ARN to the -// Bedrock model ID it wraps. -type inferenceProfileResolver func(ctx context.Context, cfg config.AWSBedrock, creds aws.CredentialsProvider, profileARN string) (string, error) - // inferenceProfileResolutionTimeout bounds the Bedrock control-plane calls made // while constructing a provider, which also cover the first credential // resolution (STS/IRSA). @@ -50,11 +47,15 @@ func isApplicationInferenceProfileARN(model string) bool { // The caller's credentials sign the call, so the required // bedrock:GetInferenceProfile permission belongs to the identity that already // invokes Bedrock, including any role assumed via config.AWSBedrock.RoleARN. +// The rest of the client configuration comes from the AWS environment, so +// custom control-plane endpoints are honored. func resolveInferenceProfile(ctx context.Context, cfg config.AWSBedrock, creds aws.CredentialsProvider, profileARN string) (string, error) { - client := bedrock.NewFromConfig(aws.Config{ - Region: cfg.Region, - Credentials: creds, - }) + awsCfg, err := awsconfig.LoadDefaultConfig(ctx, awsconfig.WithRegion(cfg.Region)) + if err != nil { + return "", xerrors.Errorf("load AWS config: %w", err) + } + awsCfg.Credentials = creds + client := bedrock.NewFromConfig(awsCfg) out, err := client.GetInferenceProfile(ctx, &bedrock.GetInferenceProfileInput{ InferenceProfileIdentifier: aws.String(profileARN), @@ -95,7 +96,7 @@ func modelIDFromARN(modelARN string) (string, error) { // IDs used for capability detection, usage recording, and pricing. Identifiers // that are not application inference profile ARNs are returned unchanged and // cost no AWS call. -func resolveBedrockModels(ctx context.Context, cfg config.AWSBedrock, creds aws.CredentialsProvider, resolve inferenceProfileResolver) (model, smallFastModel string, err error) { +func resolveBedrockModels(ctx context.Context, cfg config.AWSBedrock, creds aws.CredentialsProvider, resolve func(ctx context.Context, cfg config.AWSBedrock, creds aws.CredentialsProvider, profileARN string) (string, error)) (model, smallFastModel string, err error) { resolveOne := func(configured string) (string, error) { if !isApplicationInferenceProfileARN(configured) { return configured, nil diff --git a/aibridge/provider/bedrock_inference_profile_internal_test.go b/aibridge/provider/bedrock_inference_profile_internal_test.go index 40b07f5641e..2745327ffd6 100644 --- a/aibridge/provider/bedrock_inference_profile_internal_test.go +++ b/aibridge/provider/bedrock_inference_profile_internal_test.go @@ -2,6 +2,8 @@ package provider import ( "context" + "net/http" + "net/http/httptest" "testing" "github.com/aws/aws-sdk-go-v2/aws" @@ -209,14 +211,16 @@ func TestResolveBedrockModels(t *testing.T) { }) } +// TestNewAnthropic_InferenceProfileResolution drives the Bedrock +// GetInferenceProfile path against a mock endpoint. +// https://docs.aws.amazon.com/bedrock/latest/APIReference/API_GetInferenceProfile.html +// NOTE: no t.Parallel() because the subtests use t.Setenv. func TestNewAnthropic_InferenceProfileResolution(t *testing.T) { - t.Parallel() - - const profileARN = "arn:aws:bedrock:eu-west-2:123456789012:application-inference-profile/46u2vhiyo6z5" + const profileARN = "arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/46u2vhiyo6z5" bedrockCfg := func(model string) *config.AWSBedrock { return &config.AWSBedrock{ - Region: "eu-west-2", + Region: "us-east-1", AccessKey: "test-key", AccessKeySecret: "test-secret", Model: model, @@ -224,44 +228,72 @@ func TestNewAnthropic_InferenceProfileResolution(t *testing.T) { } } - t.Run("resolved profile drives the model id", func(t *testing.T) { - t.Parallel() + // mockBedrock serves the Bedrock control-plane API and records the paths it + // receives. Callers point the SDK at the returned URL. + mockBedrock := func(t *testing.T, handler http.HandlerFunc) (url string, paths *[]string) { + t.Helper() + + var got []string + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + got = append(got, r.URL.Path) + handler(w, r) + })) + t.Cleanup(srv.Close) + return srv.URL, &got + } - resolve := func(_ context.Context, _ config.AWSBedrock, _ aws.CredentialsProvider, arn string) (string, error) { - require.Equal(t, profileARN, arn) - return "anthropic.claude-opus-4-8", nil - } + t.Run("resolved profile drives the model id", func(t *testing.T) { + url, paths := mockBedrock(t, func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"models":[{"modelArn":"arn:aws:bedrock:us-east-1::foundation-model/anthropic.claude-opus-4-8"}]}`)) + }) + t.Setenv("AWS_ENDPOINT_URL_BEDROCK", url) - p, err := NewAnthropic(context.Background(), config.Anthropic{}, bedrockCfg(profileARN), withInferenceProfileResolver(resolve)) + p, err := NewAnthropic(context.Background(), config.Anthropic{}, bedrockCfg(profileARN)) require.NoError(t, err) require.Equal(t, "anthropic.claude-opus-4-8", p.bedrock.ResolvedModel()) // The profile stays the configured identifier so AWS attributes spend to it. require.Equal(t, profileARN, p.bedrock.ConfiguredModel()) require.Equal(t, "anthropic.claude-haiku-4-5", p.bedrock.ResolvedSmallFastModel()) + require.Len(t, *paths, 1, "only the profile ARN is resolved") + require.Contains(t, (*paths)[0], profileARN) }) t.Run("failed resolution fails construction", func(t *testing.T) { - t.Parallel() - - resolve := func(context.Context, config.AWSBedrock, aws.CredentialsProvider, string) (string, error) { - return "", xerrors.New("AccessDeniedException") - } + url, _ := mockBedrock(t, func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.Header().Set("X-Amzn-Errortype", "AccessDeniedException") + w.WriteHeader(http.StatusForbidden) + _, _ = w.Write([]byte(`{"message":"not authorized to perform bedrock:GetInferenceProfile"}`)) + }) + t.Setenv("AWS_ENDPOINT_URL_BEDROCK", url) - _, err := NewAnthropic(context.Background(), config.Anthropic{}, bedrockCfg(profileARN), withInferenceProfileResolver(resolve)) + _, err := NewAnthropic(context.Background(), config.Anthropic{}, bedrockCfg(profileARN)) require.ErrorContains(t, err, "resolve bedrock models") + require.ErrorContains(t, err, "GetInferenceProfile") }) - t.Run("plain model id needs no resolution", func(t *testing.T) { - t.Parallel() + t.Run("profile without a model fails construction", func(t *testing.T) { + url, _ := mockBedrock(t, func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"models":[]}`)) + }) + t.Setenv("AWS_ENDPOINT_URL_BEDROCK", url) - resolve := func(context.Context, config.AWSBedrock, aws.CredentialsProvider, string) (string, error) { - t.Error("resolver called for a plain model id") - return "", nil - } + _, err := NewAnthropic(context.Background(), config.Anthropic{}, bedrockCfg(profileARN)) + require.ErrorContains(t, err, "references no model") + }) + + t.Run("plain model id needs no resolution", func(t *testing.T) { + url, paths := mockBedrock(t, func(http.ResponseWriter, *http.Request) { + t.Error("Bedrock called for a plain model id") + }) + t.Setenv("AWS_ENDPOINT_URL_BEDROCK", url) - p, err := NewAnthropic(context.Background(), config.Anthropic{}, bedrockCfg("eu.anthropic.claude-opus-4-8"), withInferenceProfileResolver(resolve)) + p, err := NewAnthropic(context.Background(), config.Anthropic{}, bedrockCfg("eu.anthropic.claude-opus-4-8")) require.NoError(t, err) require.Equal(t, "eu.anthropic.claude-opus-4-8", p.bedrock.ResolvedModel()) require.Equal(t, "eu.anthropic.claude-opus-4-8", p.bedrock.ConfiguredModel()) + require.Empty(t, *paths) }) } From 1ee5a5be8484f9fa7d6b2b0357564d4b5ec4a0d2 Mon Sep 17 00:00:00 2001 From: Yevhenii Shcherbina Date: Wed, 2 Sep 2026 14:33:11 +0000 Subject: [PATCH 09/18] refactor(aibridge/provider): return the loaded AWS config from buildBedrockCredentials Inference profile resolution reused the credentials but reloaded the AWS config and overwrote its credentials. Return the config that was already loaded so the control-plane client shares one identity and one set of environment-derived settings. --- aibridge/provider/anthropic.go | 10 ++--- aibridge/provider/bedrock.go | 32 +++++++++------ .../provider/bedrock_inference_profile.go | 20 +++------- ...bedrock_inference_profile_internal_test.go | 16 ++++---- aibridge/provider/bedrock_internal_test.go | 40 +++++++++---------- 5 files changed, 58 insertions(+), 60 deletions(-) diff --git a/aibridge/provider/anthropic.go b/aibridge/provider/anthropic.go index bb5446ef494..70a0c1c596c 100644 --- a/aibridge/provider/anthropic.go +++ b/aibridge/provider/anthropic.go @@ -69,15 +69,15 @@ func NewAnthropic(ctx context.Context, cfg config.Anthropic, bedrockCfg *config. // so it is cheap to run at construction. var bedrock *messages.BedrockRuntime if bedrockCfg != nil { - creds, resolvedRegion, err := buildBedrockCredentials(ctx, *bedrockCfg) + awsCfg, err := buildBedrockCredentials(ctx, *bedrockCfg) if err != nil { return nil, xerrors.Errorf("build bedrock credentials: %w", err) } runtimeCfg := *bedrockCfg - // resolvedRegion is bedrockCfg.Region if provided; + // awsCfg.Region is bedrockCfg.Region if provided; // otherwise, it is resolved from the environment via awsconfig.LoadDefaultConfig if runtimeCfg.Region == "" { - runtimeCfg.Region = resolvedRegion + runtimeCfg.Region = awsCfg.Region } if err := runtimeCfg.Validate(); err != nil { return nil, xerrors.Errorf("bedrock config: %w", err) @@ -90,12 +90,12 @@ func NewAnthropic(ctx context.Context, cfg config.Anthropic, bedrockCfg *config. // Bedrock rejects outright on models that only accept adaptive thinking. resolveCtx, cancel := context.WithTimeout(ctx, inferenceProfileResolutionTimeout) defer cancel() - model, smallFastModel, err := resolveBedrockModels(resolveCtx, runtimeCfg, creds, resolveInferenceProfile) + model, smallFastModel, err := resolveBedrockModels(resolveCtx, runtimeCfg, awsCfg, resolveInferenceProfile) if err != nil { return nil, xerrors.Errorf("resolve bedrock models: %w", err) } - bedrock = messages.NewBedrockRuntime(runtimeCfg, creds, model, smallFastModel) + bedrock = messages.NewBedrockRuntime(runtimeCfg, awsCfg.Credentials, model, smallFastModel) } return &Anthropic{ diff --git a/aibridge/provider/bedrock.go b/aibridge/provider/bedrock.go index 23f5b1db6aa..3ef73ae522c 100644 --- a/aibridge/provider/bedrock.go +++ b/aibridge/provider/bedrock.go @@ -24,16 +24,21 @@ const bedrockSessionName = "coder-aigateway" // static keys or the AWS SDK default credential chain, which covers IRSA, // EKS Pod Identity, EC2 Instance Profile, and more. // -// 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 -// Provider should be shared across all LLM requests to the Bedrock Provider, -// 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) { +// It returns the loaded AWS config with the resolved credentials attached, so +// callers that need an AWS client reuse the same environment-derived settings +// and identity rather than assembling their own. +// +// The credentials are 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 Provider should be shared across all LLM requests to the Bedrock +// Provider, 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.Config, error) { if cfg.Region == "" && cfg.BaseURL == "" { - return nil, "", xerrors.New("region or base url required") + return aws.Config{}, xerrors.New("region or base url required") } var loadOpts []func(*awsconfig.LoadOptions) error @@ -55,21 +60,21 @@ func buildBedrockCredentials(ctx context.Context, cfg config.AWSBedrock) (aws.Cr )) // Only one set: misconfiguration. case cfg.AccessKey != "" || cfg.AccessKeySecret != "": - return nil, "", xerrors.New("both access key and access key secret must be provided together") + return aws.Config{}, xerrors.New("both access key and access key secret must be provided together") // Neither set: SDK default credential chain resolves the base identity. default: } base, err := awsconfig.LoadDefaultConfig(ctx, loadOpts...) if err != nil { - return nil, "", xerrors.Errorf("failed to load AWS Bedrock config: %w", err) + return aws.Config{}, xerrors.Errorf("failed to load AWS Bedrock config: %w", err) } // Assuming a role calls STS, which needs a region to resolve its endpoint. // The region may come from the config or the AWS environment; if neither // supplies one, fail here. if cfg.RoleARN != "" && base.Region == "" { - return nil, "", xerrors.New("region is required to assume a role, but was not specified") + return aws.Config{}, 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 @@ -107,5 +112,6 @@ func buildBedrockCredentials(ctx context.Context, cfg config.AWSBedrock) (aws.Cr // base.Region is the region the SDK resolved (explicit config, AWS_REGION / // AWS_DEFAULT_REGION, shared config, or IMDS). - return credsProvider, base.Region, nil + base.Credentials = credsProvider + return base, nil } diff --git a/aibridge/provider/bedrock_inference_profile.go b/aibridge/provider/bedrock_inference_profile.go index 13ce1c56334..019af0cd448 100644 --- a/aibridge/provider/bedrock_inference_profile.go +++ b/aibridge/provider/bedrock_inference_profile.go @@ -7,7 +7,6 @@ import ( "github.com/aws/aws-sdk-go-v2/aws" "github.com/aws/aws-sdk-go-v2/aws/arn" - awsconfig "github.com/aws/aws-sdk-go-v2/config" "github.com/aws/aws-sdk-go-v2/service/bedrock" "golang.org/x/xerrors" @@ -44,17 +43,10 @@ func isApplicationInferenceProfileARN(model string) bool { // resolveInferenceProfile returns the Bedrock model ID behind an application // inference profile ARN. // -// The caller's credentials sign the call, so the required -// bedrock:GetInferenceProfile permission belongs to the identity that already -// invokes Bedrock, including any role assumed via config.AWSBedrock.RoleARN. -// The rest of the client configuration comes from the AWS environment, so -// custom control-plane endpoints are honored. -func resolveInferenceProfile(ctx context.Context, cfg config.AWSBedrock, creds aws.CredentialsProvider, profileARN string) (string, error) { - awsCfg, err := awsconfig.LoadDefaultConfig(ctx, awsconfig.WithRegion(cfg.Region)) - if err != nil { - return "", xerrors.Errorf("load AWS config: %w", err) - } - awsCfg.Credentials = creds +// awsCfg carries the identity that invokes Bedrock, including any role assumed +// via config.AWSBedrock.RoleARN, so the required bedrock:GetInferenceProfile +// permission belongs to that identity. +func resolveInferenceProfile(ctx context.Context, awsCfg aws.Config, profileARN string) (string, error) { client := bedrock.NewFromConfig(awsCfg) out, err := client.GetInferenceProfile(ctx, &bedrock.GetInferenceProfileInput{ @@ -96,12 +88,12 @@ func modelIDFromARN(modelARN string) (string, error) { // IDs used for capability detection, usage recording, and pricing. Identifiers // that are not application inference profile ARNs are returned unchanged and // cost no AWS call. -func resolveBedrockModels(ctx context.Context, cfg config.AWSBedrock, creds aws.CredentialsProvider, resolve func(ctx context.Context, cfg config.AWSBedrock, creds aws.CredentialsProvider, profileARN string) (string, error)) (model, smallFastModel string, err error) { +func resolveBedrockModels(ctx context.Context, cfg config.AWSBedrock, awsCfg aws.Config, resolve func(ctx context.Context, awsCfg aws.Config, profileARN string) (string, error)) (model, smallFastModel string, err error) { resolveOne := func(configured string) (string, error) { if !isApplicationInferenceProfileARN(configured) { return configured, nil } - return resolve(ctx, cfg, creds, configured) + return resolve(ctx, awsCfg, configured) } model, err = resolveOne(cfg.Model) diff --git a/aibridge/provider/bedrock_inference_profile_internal_test.go b/aibridge/provider/bedrock_inference_profile_internal_test.go index 2745327ffd6..10d4316d1ca 100644 --- a/aibridge/provider/bedrock_inference_profile_internal_test.go +++ b/aibridge/provider/bedrock_inference_profile_internal_test.go @@ -141,12 +141,12 @@ func TestResolveBedrockModels(t *testing.T) { Model: "eu.anthropic.claude-opus-4-8", SmallFastModel: "anthropic.claude-haiku-4-5", } - resolve := func(context.Context, config.AWSBedrock, aws.CredentialsProvider, string) (string, error) { + resolve := func(context.Context, aws.Config, string) (string, error) { t.Error("resolver called for a plain model id") return "", nil } - model, smallFastModel, err := resolveBedrockModels(context.Background(), cfg, nil, resolve) + model, smallFastModel, err := resolveBedrockModels(context.Background(), cfg, aws.Config{}, resolve) require.NoError(t, err) require.Equal(t, cfg.Model, model) require.Equal(t, cfg.SmallFastModel, smallFastModel) @@ -163,11 +163,11 @@ func TestResolveBedrockModels(t *testing.T) { profileARN: "anthropic.claude-opus-4-8", smallFastProfileARN: "anthropic.claude-haiku-4-5", } - resolve := func(_ context.Context, _ config.AWSBedrock, _ aws.CredentialsProvider, arn string) (string, error) { + resolve := func(_ context.Context, _ aws.Config, arn string) (string, error) { return resolved[arn], nil } - model, smallFastModel, err := resolveBedrockModels(context.Background(), cfg, nil, resolve) + model, smallFastModel, err := resolveBedrockModels(context.Background(), cfg, aws.Config{}, resolve) require.NoError(t, err) require.Equal(t, "anthropic.claude-opus-4-8", model) require.Equal(t, "anthropic.claude-haiku-4-5", smallFastModel) @@ -181,12 +181,12 @@ func TestResolveBedrockModels(t *testing.T) { SmallFastModel: smallFastProfileARN, } var calls []string - resolve := func(_ context.Context, _ config.AWSBedrock, _ aws.CredentialsProvider, arn string) (string, error) { + resolve := func(_ context.Context, _ aws.Config, arn string) (string, error) { calls = append(calls, arn) return "anthropic.claude-haiku-4-5", nil } - model, smallFastModel, err := resolveBedrockModels(context.Background(), cfg, nil, resolve) + model, smallFastModel, err := resolveBedrockModels(context.Background(), cfg, aws.Config{}, resolve) require.NoError(t, err) require.Equal(t, []string{smallFastProfileARN}, calls) require.Equal(t, cfg.Model, model) @@ -200,11 +200,11 @@ func TestResolveBedrockModels(t *testing.T) { Model: profileARN, SmallFastModel: "anthropic.claude-haiku-4-5", } - resolve := func(context.Context, config.AWSBedrock, aws.CredentialsProvider, string) (string, error) { + resolve := func(context.Context, aws.Config, string) (string, error) { return "", xerrors.New("AccessDeniedException") } - _, _, err := resolveBedrockModels(context.Background(), cfg, nil, resolve) + _, _, err := resolveBedrockModels(context.Background(), cfg, aws.Config{}, resolve) require.Error(t, err) require.Contains(t, err.Error(), "resolve model") require.Contains(t, err.Error(), "AccessDeniedException") diff --git a/aibridge/provider/bedrock_internal_test.go b/aibridge/provider/bedrock_internal_test.go index e9827b7895f..1577f3b5262 100644 --- a/aibridge/provider/bedrock_internal_test.go +++ b/aibridge/provider/bedrock_internal_test.go @@ -49,7 +49,7 @@ func TestBuildBedrockCredentialsValidation(t *testing.T) { t.Run(tt.name, func(t *testing.T) { t.Parallel() - _, _, err := buildBedrockCredentials(context.Background(), tt.cfg) + _, err := buildBedrockCredentials(context.Background(), tt.cfg) require.Error(t, err) require.Contains(t, err.Error(), tt.errorMsg) }) @@ -60,14 +60,14 @@ func TestBuildBedrockCredentialsValidation(t *testing.T) { func TestBuildBedrockCredentialsStatic(t *testing.T) { t.Parallel() - creds, _, err := buildBedrockCredentials(context.Background(), config.AWSBedrock{ + awsCfg, err := buildBedrockCredentials(context.Background(), config.AWSBedrock{ Region: "us-east-1", AccessKey: "test-key", AccessKeySecret: "test-secret", }) require.NoError(t, err) - got, err := creds.Retrieve(context.Background()) + got, err := awsCfg.Credentials.Retrieve(context.Background()) require.NoError(t, err) require.Equal(t, "test-key", got.AccessKeyID) require.Equal(t, "test-secret", got.SecretAccessKey) @@ -135,13 +135,13 @@ func TestBuildBedrockCredentialsDefaultChain(t *testing.T) { // buildBedrockCredentials only wires up the provider chain; it // does not resolve credentials, so it succeeds regardless of // credential availability. Resolution failures surface on Retrieve. - creds, _, err := buildBedrockCredentials(context.Background(), config.AWSBedrock{ + awsCfg, err := buildBedrockCredentials(context.Background(), config.AWSBedrock{ Region: "us-east-1", }) require.NoError(t, err) - require.NotNil(t, creds) + require.NotNil(t, awsCfg.Credentials) - got, err := creds.Retrieve(context.Background()) + got, err := awsCfg.Credentials.Retrieve(context.Background()) if tt.expectError { require.Error(t, err) return @@ -193,13 +193,13 @@ func TestBuildBedrockCredentialsAssumeRole(t *testing.T) { t.Setenv("AWS_ACCESS_KEY_ID", "base-key") t.Setenv("AWS_SECRET_ACCESS_KEY", "base-secret") - creds, _, err := buildBedrockCredentials(context.Background(), config.AWSBedrock{ + awsCfg, err := buildBedrockCredentials(context.Background(), config.AWSBedrock{ Region: "us-east-1", RoleARN: "arn:aws:iam::123456789012:role/target", }) require.NoError(t, err) - got, err := creds.Retrieve(context.Background()) + got, err := awsCfg.Credentials.Retrieve(context.Background()) require.NoError(t, err) require.Equal(t, "ASIAASSUMED", got.AccessKeyID) require.Equal(t, "assumed-secret", got.SecretAccessKey) @@ -254,14 +254,14 @@ func TestBuildBedrockCredentialsAssumeRoleExternalID(t *testing.T) { t.Setenv("AWS_ACCESS_KEY_ID", "base-key") t.Setenv("AWS_SECRET_ACCESS_KEY", "base-secret") - creds, _, err := buildBedrockCredentials(context.Background(), config.AWSBedrock{ + awsCfg, err := buildBedrockCredentials(context.Background(), config.AWSBedrock{ Region: "us-east-1", RoleARN: "arn:aws:iam::123456789012:role/target", ExternalID: tt.externalID, }) require.NoError(t, err) - _, err = creds.Retrieve(context.Background()) + _, err = awsCfg.Credentials.Retrieve(context.Background()) require.NoError(t, err) require.Equal(t, tt.wantExternalID, gotExternalID) }) @@ -293,13 +293,13 @@ func TestBuildBedrockCredentialsAssumeRoleError(t *testing.T) { t.Setenv("AWS_ACCESS_KEY_ID", "base-key") t.Setenv("AWS_SECRET_ACCESS_KEY", "base-secret") - creds, _, err := buildBedrockCredentials(context.Background(), config.AWSBedrock{ + awsCfg, err := buildBedrockCredentials(context.Background(), config.AWSBedrock{ Region: "us-east-1", RoleARN: "arn:aws:iam::123456789012:role/target", }) require.NoError(t, err) // Build is lazy; the STS call happens on Retrieve. - _, err = creds.Retrieve(context.Background()) + _, err = awsCfg.Credentials.Retrieve(context.Background()) require.Error(t, err) // The error must carry the STS operation and failure code so operators can // tell this is an AssumeRole authorization problem, not missing credentials. @@ -337,7 +337,7 @@ func TestBuildBedrockCredentialsAssumeRoleCaches(t *testing.T) { t.Setenv("AWS_ACCESS_KEY_ID", "base-key") t.Setenv("AWS_SECRET_ACCESS_KEY", "base-secret") - creds, _, err := buildBedrockCredentials(context.Background(), config.AWSBedrock{ + awsCfg, err := buildBedrockCredentials(context.Background(), config.AWSBedrock{ Region: "us-east-1", RoleARN: "arn:aws:iam::123456789012:role/target", }) @@ -346,7 +346,7 @@ func TestBuildBedrockCredentialsAssumeRoleCaches(t *testing.T) { // Each retrieval stands in for an LLM request resolving credentials from the // shared provider. Only the first should reach STS. for range 5 { - got, err := creds.Retrieve(context.Background()) + got, err := awsCfg.Credentials.Retrieve(context.Background()) require.NoError(t, err) require.Equal(t, "ASIAASSUMED", got.AccessKeyID) } @@ -386,15 +386,15 @@ func TestBuildBedrockCredentialsAssumeRoleRefreshesOnExpiry(t *testing.T) { t.Setenv("AWS_ACCESS_KEY_ID", "base-key") t.Setenv("AWS_SECRET_ACCESS_KEY", "base-secret") - creds, _, err := buildBedrockCredentials(context.Background(), config.AWSBedrock{ + awsCfg, err := buildBedrockCredentials(context.Background(), config.AWSBedrock{ Region: "us-east-1", RoleARN: "arn:aws:iam::123456789012:role/target", }) require.NoError(t, err) - _, err = creds.Retrieve(context.Background()) + _, err = awsCfg.Credentials.Retrieve(context.Background()) require.NoError(t, err) - _, err = creds.Retrieve(context.Background()) + _, err = awsCfg.Credentials.Retrieve(context.Background()) require.NoError(t, err) require.Equal(t, int64(2), stsCalls.Load(), @@ -415,7 +415,7 @@ func TestBuildBedrockCredentialsAssumeRoleRequiresRegion(t *testing.T) { t.Setenv("AWS_SHARED_CREDENTIALS_FILE", "/dev/null") t.Setenv("AWS_EC2_METADATA_DISABLED", "true") - _, _, err := buildBedrockCredentials(context.Background(), config.AWSBedrock{ + _, err := buildBedrockCredentials(context.Background(), config.AWSBedrock{ BaseURL: "https://bedrock-runtime.example.com", RoleARN: "arn:aws:iam::123456789012:role/target", }) @@ -430,10 +430,10 @@ func TestBuildBedrockCredentialsAssumeRoleRegionFromEnv(t *testing.T) { t.Setenv("AWS_REGION", "us-west-2") // BaseURL set with no explicit region: the region comes from AWS_REGION. - _, region, err := buildBedrockCredentials(context.Background(), config.AWSBedrock{ + awsCfg, err := buildBedrockCredentials(context.Background(), config.AWSBedrock{ BaseURL: "https://bedrock-runtime.example.com", RoleARN: "arn:aws:iam::123456789012:role/target", }) require.NoError(t, err) - require.Equal(t, "us-west-2", region) + require.Equal(t, "us-west-2", awsCfg.Region) } From a928bb5ed0bac850e645d2e8bd2ad289a4889a00 Mon Sep 17 00:00:00 2001 From: Yevhenii Shcherbina Date: Wed, 2 Sep 2026 14:50:26 +0000 Subject: [PATCH 10/18] refactor(aibridge/provider): call the inference profile resolver directly The injected resolver existed only for tests, which now drive resolution through a mock Bedrock endpoint. Asserting that no request reaches Bedrock is also a stronger statement than asserting a stub was unused. --- aibridge/provider/anthropic.go | 2 +- .../provider/bedrock_inference_profile.go | 4 +- ...bedrock_inference_profile_internal_test.go | 108 ++++-------------- 3 files changed, 24 insertions(+), 90 deletions(-) diff --git a/aibridge/provider/anthropic.go b/aibridge/provider/anthropic.go index 70a0c1c596c..7876adc06b9 100644 --- a/aibridge/provider/anthropic.go +++ b/aibridge/provider/anthropic.go @@ -90,7 +90,7 @@ func NewAnthropic(ctx context.Context, cfg config.Anthropic, bedrockCfg *config. // Bedrock rejects outright on models that only accept adaptive thinking. resolveCtx, cancel := context.WithTimeout(ctx, inferenceProfileResolutionTimeout) defer cancel() - model, smallFastModel, err := resolveBedrockModels(resolveCtx, runtimeCfg, awsCfg, resolveInferenceProfile) + model, smallFastModel, err := resolveBedrockModels(resolveCtx, runtimeCfg, awsCfg) if err != nil { return nil, xerrors.Errorf("resolve bedrock models: %w", err) } diff --git a/aibridge/provider/bedrock_inference_profile.go b/aibridge/provider/bedrock_inference_profile.go index 019af0cd448..092ede806e7 100644 --- a/aibridge/provider/bedrock_inference_profile.go +++ b/aibridge/provider/bedrock_inference_profile.go @@ -88,12 +88,12 @@ func modelIDFromARN(modelARN string) (string, error) { // IDs used for capability detection, usage recording, and pricing. Identifiers // that are not application inference profile ARNs are returned unchanged and // cost no AWS call. -func resolveBedrockModels(ctx context.Context, cfg config.AWSBedrock, awsCfg aws.Config, resolve func(ctx context.Context, awsCfg aws.Config, profileARN string) (string, error)) (model, smallFastModel string, err error) { +func resolveBedrockModels(ctx context.Context, cfg config.AWSBedrock, awsCfg aws.Config) (model, smallFastModel string, err error) { resolveOne := func(configured string) (string, error) { if !isApplicationInferenceProfileARN(configured) { return configured, nil } - return resolve(ctx, awsCfg, configured) + return resolveInferenceProfile(ctx, awsCfg, configured) } model, err = resolveOne(cfg.Model) diff --git a/aibridge/provider/bedrock_inference_profile_internal_test.go b/aibridge/provider/bedrock_inference_profile_internal_test.go index 10d4316d1ca..df7458b1f23 100644 --- a/aibridge/provider/bedrock_inference_profile_internal_test.go +++ b/aibridge/provider/bedrock_inference_profile_internal_test.go @@ -6,9 +6,7 @@ import ( "net/http/httptest" "testing" - "github.com/aws/aws-sdk-go-v2/aws" "github.com/stretchr/testify/require" - "golang.org/x/xerrors" "github.com/coder/coder/v2/aibridge/config" ) @@ -126,91 +124,6 @@ func TestModelIDFromARN(t *testing.T) { } } -func TestResolveBedrockModels(t *testing.T) { - t.Parallel() - - const ( - profileARN = "arn:aws:bedrock:eu-west-2:123456789012:application-inference-profile/46u2vhiyo6z5" - smallFastProfileARN = "arn:aws:bedrock:eu-west-2:123456789012:application-inference-profile/8x1qk20fzp3r" - ) - - t.Run("plain model ids are not resolved", func(t *testing.T) { - t.Parallel() - - cfg := config.AWSBedrock{ - Model: "eu.anthropic.claude-opus-4-8", - SmallFastModel: "anthropic.claude-haiku-4-5", - } - resolve := func(context.Context, aws.Config, string) (string, error) { - t.Error("resolver called for a plain model id") - return "", nil - } - - model, smallFastModel, err := resolveBedrockModels(context.Background(), cfg, aws.Config{}, resolve) - require.NoError(t, err) - require.Equal(t, cfg.Model, model) - require.Equal(t, cfg.SmallFastModel, smallFastModel) - }) - - t.Run("profile arns are resolved", func(t *testing.T) { - t.Parallel() - - cfg := config.AWSBedrock{ - Model: profileARN, - SmallFastModel: smallFastProfileARN, - } - resolved := map[string]string{ - profileARN: "anthropic.claude-opus-4-8", - smallFastProfileARN: "anthropic.claude-haiku-4-5", - } - resolve := func(_ context.Context, _ aws.Config, arn string) (string, error) { - return resolved[arn], nil - } - - model, smallFastModel, err := resolveBedrockModels(context.Background(), cfg, aws.Config{}, resolve) - require.NoError(t, err) - require.Equal(t, "anthropic.claude-opus-4-8", model) - require.Equal(t, "anthropic.claude-haiku-4-5", smallFastModel) - }) - - t.Run("small fast model resolves independently", func(t *testing.T) { - t.Parallel() - - cfg := config.AWSBedrock{ - Model: "eu.anthropic.claude-opus-4-8", - SmallFastModel: smallFastProfileARN, - } - var calls []string - resolve := func(_ context.Context, _ aws.Config, arn string) (string, error) { - calls = append(calls, arn) - return "anthropic.claude-haiku-4-5", nil - } - - model, smallFastModel, err := resolveBedrockModels(context.Background(), cfg, aws.Config{}, resolve) - require.NoError(t, err) - require.Equal(t, []string{smallFastProfileARN}, calls) - require.Equal(t, cfg.Model, model) - require.Equal(t, "anthropic.claude-haiku-4-5", smallFastModel) - }) - - t.Run("resolution failure is returned", func(t *testing.T) { - t.Parallel() - - cfg := config.AWSBedrock{ - Model: profileARN, - SmallFastModel: "anthropic.claude-haiku-4-5", - } - resolve := func(context.Context, aws.Config, string) (string, error) { - return "", xerrors.New("AccessDeniedException") - } - - _, _, err := resolveBedrockModels(context.Background(), cfg, aws.Config{}, resolve) - require.Error(t, err) - require.Contains(t, err.Error(), "resolve model") - require.Contains(t, err.Error(), "AccessDeniedException") - }) -} - // TestNewAnthropic_InferenceProfileResolution drives the Bedrock // GetInferenceProfile path against a mock endpoint. // https://docs.aws.amazon.com/bedrock/latest/APIReference/API_GetInferenceProfile.html @@ -284,6 +197,27 @@ func TestNewAnthropic_InferenceProfileResolution(t *testing.T) { require.ErrorContains(t, err, "references no model") }) + t.Run("small fast profile resolves independently", func(t *testing.T) { + const smallFastProfileARN = "arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/8x1qk20fzp3r" + + url, paths := mockBedrock(t, func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"models":[{"modelArn":"arn:aws:bedrock:us-east-1::foundation-model/anthropic.claude-haiku-4-5"}]}`)) + }) + t.Setenv("AWS_ENDPOINT_URL_BEDROCK", url) + + cfg := bedrockCfg("eu.anthropic.claude-opus-4-8") + cfg.SmallFastModel = smallFastProfileARN + + p, err := NewAnthropic(context.Background(), config.Anthropic{}, cfg) + require.NoError(t, err) + require.Equal(t, "eu.anthropic.claude-opus-4-8", p.bedrock.ResolvedModel()) + require.Equal(t, "anthropic.claude-haiku-4-5", p.bedrock.ResolvedSmallFastModel()) + require.Equal(t, smallFastProfileARN, p.bedrock.ConfiguredSmallFastModel()) + require.Len(t, *paths, 1, "only the small fast profile ARN is resolved") + require.Contains(t, (*paths)[0], smallFastProfileARN) + }) + t.Run("plain model id needs no resolution", func(t *testing.T) { url, paths := mockBedrock(t, func(http.ResponseWriter, *http.Request) { t.Error("Bedrock called for a plain model id") From 6e837d52f77a938f946749dc157c3f26d5aebc36 Mon Sep 17 00:00:00 2001 From: evgeniy-scherbina Date: Wed, 2 Sep 2026 11:51:16 -0400 Subject: [PATCH 11/18] docs: minor changes --- aibridge/provider/anthropic.go | 3 --- 1 file changed, 3 deletions(-) diff --git a/aibridge/provider/anthropic.go b/aibridge/provider/anthropic.go index 7876adc06b9..e8796fde256 100644 --- a/aibridge/provider/anthropic.go +++ b/aibridge/provider/anthropic.go @@ -85,9 +85,6 @@ func NewAnthropic(ctx context.Context, cfg config.Anthropic, bedrockCfg *config. // Resolution only calls AWS for application inference profile ARNs, so // deployments configured with plain model IDs need no extra permission. - // A failure here fails provider construction: serving the provider with - // an unresolved profile would silently misshape every request, which - // Bedrock rejects outright on models that only accept adaptive thinking. resolveCtx, cancel := context.WithTimeout(ctx, inferenceProfileResolutionTimeout) defer cancel() model, smallFastModel, err := resolveBedrockModels(resolveCtx, runtimeCfg, awsCfg) From 457b81a5dbf669153756b2b0d2539892223ba26f Mon Sep 17 00:00:00 2001 From: Yevhenii Shcherbina Date: Wed, 2 Sep 2026 16:40:39 +0000 Subject: [PATCH 12/18] test(aibridge/provider): make Bedrock model setup explicit --- .../bedrock_inference_profile_internal_test.go | 17 +++++++---------- 1 file changed, 7 insertions(+), 10 deletions(-) diff --git a/aibridge/provider/bedrock_inference_profile_internal_test.go b/aibridge/provider/bedrock_inference_profile_internal_test.go index df7458b1f23..2ac34857e97 100644 --- a/aibridge/provider/bedrock_inference_profile_internal_test.go +++ b/aibridge/provider/bedrock_inference_profile_internal_test.go @@ -131,13 +131,13 @@ func TestModelIDFromARN(t *testing.T) { func TestNewAnthropic_InferenceProfileResolution(t *testing.T) { const profileARN = "arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/46u2vhiyo6z5" - bedrockCfg := func(model string) *config.AWSBedrock { + bedrockCfg := func(model, smallFastModel string) *config.AWSBedrock { return &config.AWSBedrock{ Region: "us-east-1", AccessKey: "test-key", AccessKeySecret: "test-secret", Model: model, - SmallFastModel: "anthropic.claude-haiku-4-5", + SmallFastModel: smallFastModel, } } @@ -162,7 +162,7 @@ func TestNewAnthropic_InferenceProfileResolution(t *testing.T) { }) t.Setenv("AWS_ENDPOINT_URL_BEDROCK", url) - p, err := NewAnthropic(context.Background(), config.Anthropic{}, bedrockCfg(profileARN)) + p, err := NewAnthropic(context.Background(), config.Anthropic{}, bedrockCfg(profileARN, "anthropic.claude-haiku-4-5")) require.NoError(t, err) require.Equal(t, "anthropic.claude-opus-4-8", p.bedrock.ResolvedModel()) // The profile stays the configured identifier so AWS attributes spend to it. @@ -181,7 +181,7 @@ func TestNewAnthropic_InferenceProfileResolution(t *testing.T) { }) t.Setenv("AWS_ENDPOINT_URL_BEDROCK", url) - _, err := NewAnthropic(context.Background(), config.Anthropic{}, bedrockCfg(profileARN)) + _, err := NewAnthropic(context.Background(), config.Anthropic{}, bedrockCfg(profileARN, "anthropic.claude-haiku-4-5")) require.ErrorContains(t, err, "resolve bedrock models") require.ErrorContains(t, err, "GetInferenceProfile") }) @@ -193,7 +193,7 @@ func TestNewAnthropic_InferenceProfileResolution(t *testing.T) { }) t.Setenv("AWS_ENDPOINT_URL_BEDROCK", url) - _, err := NewAnthropic(context.Background(), config.Anthropic{}, bedrockCfg(profileARN)) + _, err := NewAnthropic(context.Background(), config.Anthropic{}, bedrockCfg(profileARN, "anthropic.claude-haiku-4-5")) require.ErrorContains(t, err, "references no model") }) @@ -206,10 +206,7 @@ func TestNewAnthropic_InferenceProfileResolution(t *testing.T) { }) t.Setenv("AWS_ENDPOINT_URL_BEDROCK", url) - cfg := bedrockCfg("eu.anthropic.claude-opus-4-8") - cfg.SmallFastModel = smallFastProfileARN - - p, err := NewAnthropic(context.Background(), config.Anthropic{}, cfg) + p, err := NewAnthropic(context.Background(), config.Anthropic{}, bedrockCfg("eu.anthropic.claude-opus-4-8", smallFastProfileARN)) require.NoError(t, err) require.Equal(t, "eu.anthropic.claude-opus-4-8", p.bedrock.ResolvedModel()) require.Equal(t, "anthropic.claude-haiku-4-5", p.bedrock.ResolvedSmallFastModel()) @@ -224,7 +221,7 @@ func TestNewAnthropic_InferenceProfileResolution(t *testing.T) { }) t.Setenv("AWS_ENDPOINT_URL_BEDROCK", url) - p, err := NewAnthropic(context.Background(), config.Anthropic{}, bedrockCfg("eu.anthropic.claude-opus-4-8")) + p, err := NewAnthropic(context.Background(), config.Anthropic{}, bedrockCfg("eu.anthropic.claude-opus-4-8", "anthropic.claude-haiku-4-5")) require.NoError(t, err) require.Equal(t, "eu.anthropic.claude-opus-4-8", p.bedrock.ResolvedModel()) require.Equal(t, "eu.anthropic.claude-opus-4-8", p.bedrock.ConfiguredModel()) From 3d65242b1a40322fc1e1f3abc0beb335726b5528 Mon Sep 17 00:00:00 2001 From: Yevhenii Shcherbina Date: Thu, 3 Sep 2026 13:19:12 +0000 Subject: [PATCH 13/18] docs(docs/ai-coder/ai-gateway): move application inference profiles after external ID --- docs/ai-coder/ai-gateway/providers.md | 28 +++++++++++++-------------- 1 file changed, 14 insertions(+), 14 deletions(-) diff --git a/docs/ai-coder/ai-gateway/providers.md b/docs/ai-coder/ai-gateway/providers.md index b89de1554a4..f5ba0d223d1 100644 --- a/docs/ai-coder/ai-gateway/providers.md +++ b/docs/ai-coder/ai-gateway/providers.md @@ -127,20 +127,6 @@ export ANTHROPIC_BEDROCK_MANTLE_BASE_URL="/api/v2/ai-gatewa export ANTHROPIC_AUTH_TOKEN="" ``` -#### Application inference profiles - -For InvokeModel, the **model** and **small fast model** identifiers can be -[application inference profile](https://docs.aws.amazon.com/bedrock/latest/userguide/cost-mgmt-application-inference-profiles.html) -ARNs, which attribute Bedrock spend to a team or workload through AWS cost -allocation tags. AI Gateway invokes the profile so AWS records the -attribution, and resolves the model behind it to shape requests correctly and -to price usage. - -Resolution calls `GetInferenceProfile`, so the identity the gateway uses must -also permit `bedrock:GetInferenceProfile` for the profile. Providers -configured with plain model identifiers do not need this permission. When -resolution fails, the provider is not served and the failure is logged. - #### AWS credentials Do not attach API keys to a Bedrock provider. @@ -244,6 +230,20 @@ To enforce it, add the external ID to the target role's trust policy as an > not enforced, and the role can still be assumed without it. To rotate the > external ID, recreate the provider. +#### Application inference profiles + +For InvokeModel, the **model** and **small fast model** identifiers can be +[application inference profile](https://docs.aws.amazon.com/bedrock/latest/userguide/cost-mgmt-application-inference-profiles.html) +ARNs, which attribute Bedrock spend to a team or workload through AWS cost +allocation tags. AI Gateway passes the profile upstream so AWS records the +attribution, and resolves the model behind it. The underlying model is the +identity the gateway uses internally, including to price usage. + +Resolution calls `GetInferenceProfile`, so the identity the gateway uses must +also permit `bedrock:GetInferenceProfile` for the profile. Providers +configured with plain model identifiers do not need this permission. When +resolution fails, the provider is skipped. + ### GitHub Copilot GitHub Copilot offers three plans: Individual, Business, and Enterprise, From a56799503aff5ec10f15e606950450c1654f28fa Mon Sep 17 00:00:00 2001 From: Yevhenii Shcherbina Date: Thu, 3 Sep 2026 13:25:36 +0000 Subject: [PATCH 14/18] docs(docs/ai-coder/ai-gateway): clarify application inference profile resolution --- docs/ai-coder/ai-gateway/providers.md | 18 ++++++++++-------- 1 file changed, 10 insertions(+), 8 deletions(-) diff --git a/docs/ai-coder/ai-gateway/providers.md b/docs/ai-coder/ai-gateway/providers.md index f5ba0d223d1..7d002271374 100644 --- a/docs/ai-coder/ai-gateway/providers.md +++ b/docs/ai-coder/ai-gateway/providers.md @@ -235,14 +235,16 @@ To enforce it, add the external ID to the target role's trust policy as an For InvokeModel, the **model** and **small fast model** identifiers can be [application inference profile](https://docs.aws.amazon.com/bedrock/latest/userguide/cost-mgmt-application-inference-profiles.html) ARNs, which attribute Bedrock spend to a team or workload through AWS cost -allocation tags. AI Gateway passes the profile upstream so AWS records the -attribution, and resolves the model behind it. The underlying model is the -identity the gateway uses internally, including to price usage. - -Resolution calls `GetInferenceProfile`, so the identity the gateway uses must -also permit `bedrock:GetInferenceProfile` for the profile. Providers -configured with plain model identifiers do not need this permission. When -resolution fails, the provider is skipped. +allocation tags. + +AI Gateway passes the profile upstream so AWS records the attribution, while +internally resolving and using the underlying model identity, including for +usage pricing. + +Resolution requires a `GetInferenceProfile` call, so the AWS identity used by +the gateway must have `bedrock:GetInferenceProfile` permission for the +profile. Providers configured with plain model identifiers do not need this +permission. If resolution fails, the provider is skipped. ### GitHub Copilot From 0ba11afade5cc94508e706fb6055dc6f331f5599 Mon Sep 17 00:00:00 2001 From: Yevhenii Shcherbina Date: Fri, 4 Sep 2026 14:45:34 +0000 Subject: [PATCH 15/18] test(aibridge/intercept/messages): cover small fast capture in the interceptor constructors --- .../intercept/messages/base_internal_test.go | 65 +++++++++++++++++++ 1 file changed, 65 insertions(+) diff --git a/aibridge/intercept/messages/base_internal_test.go b/aibridge/intercept/messages/base_internal_test.go index 9b553a17490..38c7d185c00 100644 --- a/aibridge/intercept/messages/base_internal_test.go +++ b/aibridge/intercept/messages/base_internal_test.go @@ -261,6 +261,71 @@ func TestModelForBedrockInvokeModel(t *testing.T) { } } +// TestSmallFastModelCapturedAtConstruction covers the classification captured +// by the interceptor constructors. The configured identifiers are opaque +// profile ARNs, so the classification must come from the client payload as it +// arrives. +func TestSmallFastModelCapturedAtConstruction(t *testing.T) { + t.Parallel() + + const ( + profileARN = "arn:aws:bedrock:eu-west-2:123456789012:application-inference-profile/46u2vhiyo6z5" + smallFastProfileARN = "arn:aws:bedrock:eu-west-2:123456789012:application-inference-profile/8x1qk20fzp3r" + ) + + runtime := NewBedrockRuntime(config.AWSBedrock{ + Model: profileARN, + SmallFastModel: smallFastProfileARN, + }, nil, "anthropic.claude-opus-4-8", "anthropic.claude-haiku-4-5") + + const haikuPayload = `{"model":"claude-haiku-4-5","max_tokens":10000}` + const opusPayload = `{"model":"claude-opus-4-8","max_tokens":10000}` + + constructors := []struct { + name string + newInterception func(payload RequestPayload) *interceptionBase + }{ + {name: "blocking", newInterception: func(payload RequestPayload) *interceptionBase { + return &NewBlockingInterceptor(uuid.New(), payload, intercept.Config{}, nil, runtime, http.Header{}, nil).interceptionBase + }}, + {name: "streaming", newInterception: func(payload RequestPayload) *interceptionBase { + return &NewStreamingInterceptor(uuid.New(), payload, intercept.Config{}, nil, runtime, http.Header{}, nil).interceptionBase + }}, + } + + tests := []struct { + name string + payload string + expectModel string + expectConfigured string + }{ + { + name: "small fast model", + payload: haikuPayload, + expectModel: "anthropic.claude-haiku-4-5", + expectConfigured: smallFastProfileARN, + }, + { + name: "primary model", + payload: opusPayload, + expectModel: "anthropic.claude-opus-4-8", + expectConfigured: profileARN, + }, + } + + for _, c := range constructors { + for _, tt := range tests { + t.Run(c.name+" "+tt.name, func(t *testing.T) { + t.Parallel() + + i := c.newInterception(mustMessagesPayload(t, tt.payload)) + require.Equal(t, tt.expectModel, i.Model()) + require.Equal(t, tt.expectConfigured, i.upstreamModel()) + }) + } + } +} + // TestModelForPlainBedrockModelID covers Bedrock providers configured with // plain model IDs, which resolve to themselves. func TestModelForPlainBedrockModelID(t *testing.T) { From 958c981f786c9614eb5f43c7d3daf9b86b782775 Mon Sep 17 00:00:00 2001 From: Yevhenii Shcherbina Date: Fri, 4 Sep 2026 16:05:23 +0000 Subject: [PATCH 16/18] docs(aibridge/provider): record why the first inference profile model is used --- aibridge/provider/bedrock_inference_profile.go | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/aibridge/provider/bedrock_inference_profile.go b/aibridge/provider/bedrock_inference_profile.go index 092ede806e7..7b7411aa70a 100644 --- a/aibridge/provider/bedrock_inference_profile.go +++ b/aibridge/provider/bedrock_inference_profile.go @@ -46,6 +46,10 @@ func isApplicationInferenceProfileARN(model string) bool { // awsCfg carries the identity that invokes Bedrock, including any role assumed // via config.AWSBedrock.RoleARN, so the required bedrock:GetInferenceProfile // permission belongs to that identity. +// +// A profile that wraps a cross-region system-defined profile lists one model +// per region. Those entries differ only in the ARN region, which the model ID +// does not carry, so any entry resolves to the same model. func resolveInferenceProfile(ctx context.Context, awsCfg aws.Config, profileARN string) (string, error) { client := bedrock.NewFromConfig(awsCfg) From abda63c695d1a96a7d4dc5366fc9a92a07305f3c Mon Sep 17 00:00:00 2001 From: Yevhenii Shcherbina Date: Fri, 4 Sep 2026 16:28:21 +0000 Subject: [PATCH 17/18] fix(aibridge/provider): stop attributing every profile lookup failure to permissions --- aibridge/provider/bedrock_inference_profile.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/aibridge/provider/bedrock_inference_profile.go b/aibridge/provider/bedrock_inference_profile.go index 7b7411aa70a..b0ed2f55664 100644 --- a/aibridge/provider/bedrock_inference_profile.go +++ b/aibridge/provider/bedrock_inference_profile.go @@ -57,7 +57,7 @@ func resolveInferenceProfile(ctx context.Context, awsCfg aws.Config, profileARN InferenceProfileIdentifier: aws.String(profileARN), }) if err != nil { - return "", xerrors.Errorf("get inference profile %q (requires the %s:GetInferenceProfile permission): %w", profileARN, bedrockService, err) + return "", xerrors.Errorf("get inference profile %q: %w", profileARN, err) } if len(out.Models) == 0 || out.Models[0].ModelArn == nil { return "", xerrors.Errorf("inference profile %q references no model", profileARN) From 070d51edab31d840d8193f899b0503f84639deb6 Mon Sep 17 00:00:00 2001 From: Yevhenii Shcherbina Date: Fri, 4 Sep 2026 18:22:38 +0000 Subject: [PATCH 18/18] feat(aibridge): resolve bedrock inference profiles on first request --- aibridge/aibridgetest/aibridgetest.go | 5 +- aibridge/api.go | 14 +- .../integrationtest/bridge_internal_test.go | 2 +- aibridge/passthrough_internal_test.go | 2 +- aibridge/provider/anthropic.go | 104 ++++++---- aibridge/provider/anthropic_internal_test.go | 10 +- .../provider/bedrock_inference_profile.go | 101 +++++++--- ...bedrock_inference_profile_internal_test.go | 182 ++++++++++++++---- cli/aibridged.go | 14 +- cli/aibridged_internal_test.go | 2 +- cli/server_aibridge_internal_test.go | 4 +- docs/ai-coder/ai-gateway/providers.md | 4 +- 12 files changed, 321 insertions(+), 123 deletions(-) diff --git a/aibridge/aibridgetest/aibridgetest.go b/aibridge/aibridgetest/aibridgetest.go index c86bee683ab..1c32e2428ff 100644 --- a/aibridge/aibridgetest/aibridgetest.go +++ b/aibridge/aibridgetest/aibridgetest.go @@ -10,10 +10,11 @@ import ( ) // NewAnthropicProvider builds an Anthropic provider for tests, failing the test -// if credential resolution fails. +// if credential resolution fails. Each call gets its own inference profile +// cache so tests do not share resolutions. func NewAnthropicProvider(t testing.TB, cfg aibridge.AnthropicConfig, bedrockCfg *aibridge.AWSBedrockConfig) aibridge.Provider { t.Helper() - p, err := aibridge.NewAnthropicProvider(context.Background(), cfg, bedrockCfg) + p, err := aibridge.NewAnthropicProvider(context.Background(), cfg, bedrockCfg, aibridge.NewInferenceProfileCache()) require.NoError(t, err) return p } diff --git a/aibridge/api.go b/aibridge/api.go index 99c58fc4f94..27b243d90c0 100644 --- a/aibridge/api.go +++ b/aibridge/api.go @@ -40,14 +40,24 @@ type ( AWSBedrockConfig = config.AWSBedrock OpenAIConfig = config.OpenAI CopilotConfig = config.Copilot + + // InferenceProfileCache caches Bedrock application inference profile + // resolutions across provider reloads. Create one per process. + InferenceProfileCache = provider.InferenceProfileCache ) +// NewInferenceProfileCache returns a cache shared by every Bedrock provider in +// the process. +func NewInferenceProfileCache() *InferenceProfileCache { + return provider.NewInferenceProfileCache() +} + func AsActor(ctx context.Context, actorID string, metadata recorder.Metadata) context.Context { return aibcontext.AsActor(ctx, actorID, metadata) } -func NewAnthropicProvider(ctx context.Context, cfg config.Anthropic, bedrockCfg *config.AWSBedrock) (provider.Provider, error) { - return provider.NewAnthropic(ctx, cfg, bedrockCfg) +func NewAnthropicProvider(ctx context.Context, cfg config.Anthropic, bedrockCfg *config.AWSBedrock, profiles *InferenceProfileCache) (provider.Provider, error) { + return provider.NewAnthropic(ctx, cfg, bedrockCfg, profiles) } func NewOpenAIProvider(cfg config.OpenAI) provider.Provider { diff --git a/aibridge/internal/integrationtest/bridge_internal_test.go b/aibridge/internal/integrationtest/bridge_internal_test.go index d5f348a8930..fcef8847818 100644 --- a/aibridge/internal/integrationtest/bridge_internal_test.go +++ b/aibridge/internal/integrationtest/bridge_internal_test.go @@ -317,7 +317,7 @@ func TestAWSBedrockIntegration(t *testing.T) { SmallFastModel: "test-haiku", } - _, err := provider.NewAnthropic(ctx, anthropicCfg("http://unused", apiKey), bedrockCfg) + _, err := provider.NewAnthropic(ctx, anthropicCfg("http://unused", apiKey), bedrockCfg, provider.NewInferenceProfileCache()) require.ErrorContains(t, err, "region or base url required") }) diff --git a/aibridge/passthrough_internal_test.go b/aibridge/passthrough_internal_test.go index 7bd519f040b..472dc4725c2 100644 --- a/aibridge/passthrough_internal_test.go +++ b/aibridge/passthrough_internal_test.go @@ -320,7 +320,7 @@ func TestPassthrough_KeyFailover(t *testing.T) { p, err := provider.NewAnthropic(context.Background(), config.Anthropic{ BaseURL: baseURL, KeyPool: pool, - }, nil) + }, nil, provider.NewInferenceProfileCache()) require.NoError(t, err) return p }, diff --git a/aibridge/provider/anthropic.go b/aibridge/provider/anthropic.go index e8796fde256..0f0ac06ce77 100644 --- a/aibridge/provider/anthropic.go +++ b/aibridge/provider/anthropic.go @@ -8,6 +8,7 @@ import ( "net/http" "strings" + "github.com/aws/aws-sdk-go-v2/aws" "github.com/google/uuid" "go.opentelemetry.io/otel/codes" "go.opentelemetry.io/otel/trace" @@ -29,8 +30,14 @@ var _ Provider = &Anthropic{} // Anthropic allows for interactions with the Anthropic API. type Anthropic struct { cfg config.Anthropic - // bedrock is nil for non-Bedrock providers. - bedrock *messages.BedrockRuntime + // bedrockCfg is nil for non-Bedrock providers. + bedrockCfg *config.AWSBedrock + // awsCfg carries the region and the credentials provider that sign Bedrock + // requests. It is meaningful only alongside bedrockCfg. + awsCfg aws.Config + // profiles resolves configured model identifiers on first use. It is shared + // across providers so a reload does not discard resolutions. + profiles *InferenceProfileCache } const routeMessages = "/v1/messages" // https://docs.anthropic.com/en/api/messages @@ -51,7 +58,12 @@ var anthropicIsFailure = func(statusCode int) bool { return circuitbreaker.DefaultIsFailure(statusCode) } -func NewAnthropic(ctx context.Context, cfg config.Anthropic, bedrockCfg *config.AWSBedrock) (*Anthropic, error) { +// NewAnthropic constructs a provider. It makes no network call: Bedrock +// credentials resolve lazily on first retrieval, and application inference +// profile ARNs resolve on the first request that needs them. Construction +// therefore cannot fail because AWS is slow or briefly unreachable, which would +// otherwise drop the provider from the reloaded snapshot. +func NewAnthropic(ctx context.Context, cfg config.Anthropic, bedrockCfg *config.AWSBedrock, profiles *InferenceProfileCache) (*Anthropic, error) { if cfg.Name == "" { cfg.Name = config.ProviderAnthropic } @@ -63,42 +75,35 @@ func NewAnthropic(ctx context.Context, cfg config.Anthropic, bedrockCfg *config. cfg.CircuitBreaker.OpenErrorResponse = anthropicOpenErrorResponse } + p := &Anthropic{cfg: cfg, profiles: profiles} + if bedrockCfg == nil { + return p, nil + } + if profiles == nil { + return nil, xerrors.New("developer error: bedrock provider requires an inference profile cache") + } + // Resolve the AWS credentials provider once and bundle it with the config. // This performs no network call (the base identity and any AssumeRole // resolve lazily on first retrieval); it only wires up the provider chain, // so it is cheap to run at construction. - var bedrock *messages.BedrockRuntime - if bedrockCfg != nil { - awsCfg, err := buildBedrockCredentials(ctx, *bedrockCfg) - if err != nil { - return nil, xerrors.Errorf("build bedrock credentials: %w", err) - } - runtimeCfg := *bedrockCfg - // awsCfg.Region is bedrockCfg.Region if provided; - // otherwise, it is resolved from the environment via awsconfig.LoadDefaultConfig - if runtimeCfg.Region == "" { - runtimeCfg.Region = awsCfg.Region - } - if err := runtimeCfg.Validate(); err != nil { - return nil, xerrors.Errorf("bedrock config: %w", err) - } - - // Resolution only calls AWS for application inference profile ARNs, so - // deployments configured with plain model IDs need no extra permission. - resolveCtx, cancel := context.WithTimeout(ctx, inferenceProfileResolutionTimeout) - defer cancel() - model, smallFastModel, err := resolveBedrockModels(resolveCtx, runtimeCfg, awsCfg) - if err != nil { - return nil, xerrors.Errorf("resolve bedrock models: %w", err) - } - - bedrock = messages.NewBedrockRuntime(runtimeCfg, awsCfg.Credentials, model, smallFastModel) + awsCfg, err := buildBedrockCredentials(ctx, *bedrockCfg) + if err != nil { + return nil, xerrors.Errorf("build bedrock credentials: %w", err) + } + runtimeCfg := *bedrockCfg + // awsCfg.Region is bedrockCfg.Region if provided; + // otherwise, it is resolved from the environment via awsconfig.LoadDefaultConfig + if runtimeCfg.Region == "" { + runtimeCfg.Region = awsCfg.Region + } + if err := runtimeCfg.Validate(); err != nil { + return nil, xerrors.Errorf("bedrock config: %w", err) } - return &Anthropic{ - cfg: cfg, - bedrock: bedrock, - }, nil + p.bedrockCfg = &runtimeCfg + p.awsCfg = awsCfg + return p, nil } func (*Anthropic) Type() string { @@ -161,16 +166,41 @@ func (p *Anthropic) CreateInterceptor(_ http.ResponseWriter, r *http.Request, tr return nil, xerrors.Errorf("resolve credential: %w", err) } + var bedrock *messages.BedrockRuntime + if p.bedrockCfg != nil { + bedrock, err = p.bedrockRuntime(r.Context()) + if err != nil { + span.SetStatus(codes.Error, err.Error()) + return nil, err + } + } + var interceptor intercept.Interceptor if reqPayload.Stream() { - interceptor = messages.NewStreamingInterceptor(id, reqPayload, cfg, cred, p.bedrock, r.Header, tracer) + interceptor = messages.NewStreamingInterceptor(id, reqPayload, cfg, cred, bedrock, r.Header, tracer) } else { - interceptor = messages.NewBlockingInterceptor(id, reqPayload, cfg, cred, p.bedrock, r.Header, tracer) + interceptor = messages.NewBlockingInterceptor(id, reqPayload, cfg, cred, bedrock, r.Header, tracer) } span.SetAttributes(interceptor.TraceAttributes(r)...) return interceptor, nil } +// bedrockRuntime pairs the configured Bedrock settings with the resolved model +// identities. Resolution only calls AWS for application inference profile ARNs, +// and only until the cache holds them, so deployments configured with plain +// model IDs never call AWS here and need no extra permission. +func (p *Anthropic) bedrockRuntime(ctx context.Context) (*messages.BedrockRuntime, error) { + model, err := p.profiles.Resolve(ctx, p.awsCfg, p.bedrockCfg.Model) + if err != nil { + return nil, xerrors.Errorf("resolve model: %w", err) + } + smallFastModel, err := p.profiles.Resolve(ctx, p.awsCfg, p.bedrockCfg.SmallFastModel) + if err != nil { + return nil, xerrors.Errorf("resolve small fast model: %w", err) + } + return messages.NewBedrockRuntime(*p.bedrockCfg, p.awsCfg.Credentials, model, smallFastModel), nil +} + // resolveCredential determines the upstream credential for a request. At this // point the request contains only LLM provider headers. Any Coder-specific // authentication has already been stripped. @@ -193,8 +223,8 @@ func (p *Anthropic) resolveCredential(r *http.Request) (intercept.Credential, er if p.cfg.KeyPool != nil { return &intercept.CentralizedPool{Pool: p.cfg.KeyPool, Header: p.AuthHeader()}, nil } - if p.bedrock != nil { - return intercept.Bedrock{AccessKey: p.bedrock.Cfg.AccessKey}, nil + if p.bedrockCfg != nil { + return intercept.Bedrock{AccessKey: p.bedrockCfg.AccessKey}, nil } return nil, ErrNoCredential } diff --git a/aibridge/provider/anthropic_internal_test.go b/aibridge/provider/anthropic_internal_test.go index cdc8afe9148..d6965c5db16 100644 --- a/aibridge/provider/anthropic_internal_test.go +++ b/aibridge/provider/anthropic_internal_test.go @@ -24,7 +24,7 @@ import ( // would create an import cycle. func newTestAnthropic(t testing.TB, cfg config.Anthropic, bedrockCfg *config.AWSBedrock) *Anthropic { t.Helper() - p, err := NewAnthropic(context.Background(), cfg, bedrockCfg) + p, err := NewAnthropic(context.Background(), cfg, bedrockCfg, NewInferenceProfileCache()) require.NoError(t, err) return p } @@ -124,10 +124,10 @@ func TestNewAnthropic_BedrockRegionResolution(t *testing.T) { Protocol: config.BedrockProtocolMantle, AccessKey: "test-key", AccessKeySecret: "test-secret", - }) + }, NewInferenceProfileCache()) require.NoError(t, err) - require.NotNil(t, p.bedrock) - require.Equal(t, "us-west-2", p.bedrock.Cfg.Region) + require.NotNil(t, p.bedrockCfg) + require.Equal(t, "us-west-2", p.bedrockCfg.Region) }) t.Run("mantle_no_region_anywhere", func(t *testing.T) { @@ -145,7 +145,7 @@ func TestNewAnthropic_BedrockRegionResolution(t *testing.T) { Protocol: config.BedrockProtocolMantle, AccessKey: "test-key", AccessKeySecret: "test-secret", - }) + }, NewInferenceProfileCache()) require.ErrorContains(t, err, "region required") }) } diff --git a/aibridge/provider/bedrock_inference_profile.go b/aibridge/provider/bedrock_inference_profile.go index b0ed2f55664..03395613c90 100644 --- a/aibridge/provider/bedrock_inference_profile.go +++ b/aibridge/provider/bedrock_inference_profile.go @@ -3,14 +3,14 @@ package provider import ( "context" "strings" + "sync" "time" "github.com/aws/aws-sdk-go-v2/aws" "github.com/aws/aws-sdk-go-v2/aws/arn" "github.com/aws/aws-sdk-go-v2/service/bedrock" + "golang.org/x/sync/singleflight" "golang.org/x/xerrors" - - "github.com/coder/coder/v2/aibridge/config" ) // bedrockService is the ARN service namespace for Amazon Bedrock resources. @@ -21,11 +21,75 @@ const bedrockService = "bedrock" // Bedrock spend to a team or workload via cost allocation tags. const applicationInferenceProfileResourceType = "application-inference-profile" -// inferenceProfileResolutionTimeout bounds the Bedrock control-plane calls made -// while constructing a provider, which also cover the first credential -// resolution (STS/IRSA). +// inferenceProfileResolutionTimeout bounds a single Bedrock control-plane +// lookup, which also covers the first credential resolution (STS/IRSA). The +// request context still applies, so a client that gives up earlier cancels the +// lookup. const inferenceProfileResolutionTimeout = 30 * time.Second +// InferenceProfileCache resolves configured Bedrock model identifiers to the +// model IDs used for capability detection, usage recording, and pricing. +// +// Application inference profile ARNs are opaque and cost one AWS lookup each; +// every other identifier is returned unchanged without calling AWS. Successful +// lookups are cached for the process lifetime: Bedrock has no +// UpdateInferenceProfile, so the model behind a profile is fixed when the +// profile is created, and pointing at a different model means a new ARN. +// Failures are not cached, so a transient one is retried on the next request. +// +// A cache outlives the providers that use it, so a provider reload does not +// discard resolutions. +type InferenceProfileCache struct { + resolutions singleflight.Group + + mu sync.RWMutex + models map[string]string +} + +// NewInferenceProfileCache returns a cache ready for use. +func NewInferenceProfileCache() *InferenceProfileCache { + return &InferenceProfileCache{models: make(map[string]string)} +} + +// Resolve returns the model ID behind configured. Concurrent resolutions of the +// same identifier share a single AWS lookup. +// +// awsCfg carries the identity that invokes Bedrock, including any role assumed +// via config.AWSBedrock.RoleARN, so the required bedrock:GetInferenceProfile +// permission belongs to that identity. +func (c *InferenceProfileCache) Resolve(ctx context.Context, awsCfg aws.Config, configured string) (string, error) { + if !isApplicationInferenceProfileARN(configured) { + return configured, nil + } + + c.mu.RLock() + model, ok := c.models[configured] + c.mu.RUnlock() + if ok { + return model, nil + } + + resolved, err, _ := c.resolutions.Do(configured, func() (any, error) { + resolveCtx, cancel := context.WithTimeout(ctx, inferenceProfileResolutionTimeout) + defer cancel() + + model, err := resolveInferenceProfile(resolveCtx, awsCfg, configured) + if err != nil { + return "", err + } + + c.mu.Lock() + c.models[configured] = model + c.mu.Unlock() + return model, nil + }) + if err != nil { + return "", err + } + model, _ = resolved.(string) + return model, nil +} + // isApplicationInferenceProfileARN reports whether model is an application // inference profile ARN, whose identifier is opaque and must be resolved // through AWS. Plain model IDs and system-defined inference profile ARNs, which @@ -43,10 +107,6 @@ func isApplicationInferenceProfileARN(model string) bool { // resolveInferenceProfile returns the Bedrock model ID behind an application // inference profile ARN. // -// awsCfg carries the identity that invokes Bedrock, including any role assumed -// via config.AWSBedrock.RoleARN, so the required bedrock:GetInferenceProfile -// permission belongs to that identity. -// // A profile that wraps a cross-region system-defined profile lists one model // per region. Those entries differ only in the ARN region, which the model ID // does not carry, so any entry resolves to the same model. @@ -87,26 +147,3 @@ func modelIDFromARN(modelARN string) (string, error) { } return model, nil } - -// resolveBedrockModels resolves the configured model identifiers to the model -// IDs used for capability detection, usage recording, and pricing. Identifiers -// that are not application inference profile ARNs are returned unchanged and -// cost no AWS call. -func resolveBedrockModels(ctx context.Context, cfg config.AWSBedrock, awsCfg aws.Config) (model, smallFastModel string, err error) { - resolveOne := func(configured string) (string, error) { - if !isApplicationInferenceProfileARN(configured) { - return configured, nil - } - return resolveInferenceProfile(ctx, awsCfg, configured) - } - - model, err = resolveOne(cfg.Model) - if err != nil { - return "", "", xerrors.Errorf("resolve model: %w", err) - } - smallFastModel, err = resolveOne(cfg.SmallFastModel) - if err != nil { - return "", "", xerrors.Errorf("resolve small fast model: %w", err) - } - return model, smallFastModel, nil -} diff --git a/aibridge/provider/bedrock_inference_profile_internal_test.go b/aibridge/provider/bedrock_inference_profile_internal_test.go index 2ac34857e97..b4e3fa79b64 100644 --- a/aibridge/provider/bedrock_inference_profile_internal_test.go +++ b/aibridge/provider/bedrock_inference_profile_internal_test.go @@ -2,13 +2,20 @@ package provider import ( "context" + "fmt" "net/http" "net/http/httptest" + "strings" + "sync" + "sync/atomic" "testing" "github.com/stretchr/testify/require" + "golang.org/x/xerrors" "github.com/coder/coder/v2/aibridge/config" + "github.com/coder/coder/v2/aibridge/intercept" + "github.com/coder/coder/v2/testutil" ) func TestIsApplicationInferenceProfileARN(t *testing.T) { @@ -124,12 +131,15 @@ func TestModelIDFromARN(t *testing.T) { } } -// TestNewAnthropic_InferenceProfileResolution drives the Bedrock +// TestInferenceProfileResolutionOnFirstRequest drives the Bedrock // GetInferenceProfile path against a mock endpoint. // https://docs.aws.amazon.com/bedrock/latest/APIReference/API_GetInferenceProfile.html // NOTE: no t.Parallel() because the subtests use t.Setenv. -func TestNewAnthropic_InferenceProfileResolution(t *testing.T) { - const profileARN = "arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/46u2vhiyo6z5" +func TestInferenceProfileResolutionOnFirstRequest(t *testing.T) { + const ( + profileARN = "arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/46u2vhiyo6z5" + smallFastProfileARN = "arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/8x1qk20fzp3r" + ) bedrockCfg := func(model, smallFastModel string) *config.AWSBedrock { return &config.AWSBedrock{ @@ -146,71 +156,127 @@ func TestNewAnthropic_InferenceProfileResolution(t *testing.T) { mockBedrock := func(t *testing.T, handler http.HandlerFunc) (url string, paths *[]string) { t.Helper() - var got []string + var ( + mu sync.Mutex + got []string + ) srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + mu.Lock() got = append(got, r.URL.Path) + mu.Unlock() handler(w, r) })) t.Cleanup(srv.Close) return srv.URL, &got } - t.Run("resolved profile drives the model id", func(t *testing.T) { - url, paths := mockBedrock(t, func(w http.ResponseWriter, _ *http.Request) { + // interceptFor runs a request through the provider the way the bridge does. + interceptFor := func(t *testing.T, p *Anthropic, model string) (intercept.Interceptor, error) { + t.Helper() + + body := fmt.Sprintf(`{"model":%q,"max_tokens":10000}`, model) + req := httptest.NewRequest(http.MethodPost, p.RoutePrefix()+routeMessages, strings.NewReader(body)) + return p.CreateInterceptor(httptest.NewRecorder(), req, testTracer) + } + + respondWithModel := func(modelARN string) http.HandlerFunc { + return func(w http.ResponseWriter, _ *http.Request) { w.Header().Set("Content-Type", "application/json") - _, _ = w.Write([]byte(`{"models":[{"modelArn":"arn:aws:bedrock:us-east-1::foundation-model/anthropic.claude-opus-4-8"}]}`)) + _, _ = fmt.Fprintf(w, `{"models":[{"modelArn":%q}]}`, modelARN) + } + } + + t.Run("construction makes no aws call", func(t *testing.T) { + url, paths := mockBedrock(t, func(http.ResponseWriter, *http.Request) { + t.Error("Bedrock called during construction") }) t.Setenv("AWS_ENDPOINT_URL_BEDROCK", url) - p, err := NewAnthropic(context.Background(), config.Anthropic{}, bedrockCfg(profileARN, "anthropic.claude-haiku-4-5")) + _, err := NewAnthropic(context.Background(), config.Anthropic{}, bedrockCfg(profileARN, "anthropic.claude-haiku-4-5"), NewInferenceProfileCache()) + require.NoError(t, err) + require.Empty(t, *paths) + }) + + t.Run("resolved profile drives the model id", func(t *testing.T) { + url, paths := mockBedrock(t, respondWithModel("arn:aws:bedrock:us-east-1::foundation-model/anthropic.claude-opus-4-8")) + t.Setenv("AWS_ENDPOINT_URL_BEDROCK", url) + + p, err := NewAnthropic(context.Background(), config.Anthropic{}, bedrockCfg(profileARN, "anthropic.claude-haiku-4-5"), NewInferenceProfileCache()) require.NoError(t, err) - require.Equal(t, "anthropic.claude-opus-4-8", p.bedrock.ResolvedModel()) - // The profile stays the configured identifier so AWS attributes spend to it. - require.Equal(t, profileARN, p.bedrock.ConfiguredModel()) - require.Equal(t, "anthropic.claude-haiku-4-5", p.bedrock.ResolvedSmallFastModel()) + + interceptor, err := interceptFor(t, p, "claude-opus-4-8") + require.NoError(t, err) + require.Equal(t, "anthropic.claude-opus-4-8", interceptor.Model()) require.Len(t, *paths, 1, "only the profile ARN is resolved") require.Contains(t, (*paths)[0], profileARN) }) - t.Run("failed resolution fails construction", func(t *testing.T) { + t.Run("resolution is cached across requests", func(t *testing.T) { + url, paths := mockBedrock(t, respondWithModel("arn:aws:bedrock:us-east-1::foundation-model/anthropic.claude-opus-4-8")) + t.Setenv("AWS_ENDPOINT_URL_BEDROCK", url) + + p, err := NewAnthropic(context.Background(), config.Anthropic{}, bedrockCfg(profileARN, "anthropic.claude-haiku-4-5"), NewInferenceProfileCache()) + require.NoError(t, err) + + for range 3 { + interceptor, err := interceptFor(t, p, "claude-opus-4-8") + require.NoError(t, err) + require.Equal(t, "anthropic.claude-opus-4-8", interceptor.Model()) + } + require.Len(t, *paths, 1, "the profile is resolved once") + }) + + t.Run("failed resolution fails the request and is retried", func(t *testing.T) { + var attempts atomic.Int64 url, _ := mockBedrock(t, func(w http.ResponseWriter, _ *http.Request) { - w.Header().Set("Content-Type", "application/json") - w.Header().Set("X-Amzn-Errortype", "AccessDeniedException") - w.WriteHeader(http.StatusForbidden) - _, _ = w.Write([]byte(`{"message":"not authorized to perform bedrock:GetInferenceProfile"}`)) + if attempts.Add(1) == 1 { + w.Header().Set("Content-Type", "application/json") + w.Header().Set("X-Amzn-Errortype", "AccessDeniedException") + w.WriteHeader(http.StatusForbidden) + _, _ = w.Write([]byte(`{"message":"not authorized to perform bedrock:GetInferenceProfile"}`)) + return + } + respondWithModel("arn:aws:bedrock:us-east-1::foundation-model/anthropic.claude-opus-4-8")(w, nil) }) t.Setenv("AWS_ENDPOINT_URL_BEDROCK", url) - _, err := NewAnthropic(context.Background(), config.Anthropic{}, bedrockCfg(profileARN, "anthropic.claude-haiku-4-5")) - require.ErrorContains(t, err, "resolve bedrock models") + p, err := NewAnthropic(context.Background(), config.Anthropic{}, bedrockCfg(profileARN, "anthropic.claude-haiku-4-5"), NewInferenceProfileCache()) + require.NoError(t, err) + + _, err = interceptFor(t, p, "claude-opus-4-8") + require.ErrorContains(t, err, "resolve model") require.ErrorContains(t, err, "GetInferenceProfile") + + // The provider stays usable: the failure was not cached. + interceptor, err := interceptFor(t, p, "claude-opus-4-8") + require.NoError(t, err) + require.Equal(t, "anthropic.claude-opus-4-8", interceptor.Model()) }) - t.Run("profile without a model fails construction", func(t *testing.T) { + t.Run("profile without a model fails the request", func(t *testing.T) { url, _ := mockBedrock(t, func(w http.ResponseWriter, _ *http.Request) { w.Header().Set("Content-Type", "application/json") _, _ = w.Write([]byte(`{"models":[]}`)) }) t.Setenv("AWS_ENDPOINT_URL_BEDROCK", url) - _, err := NewAnthropic(context.Background(), config.Anthropic{}, bedrockCfg(profileARN, "anthropic.claude-haiku-4-5")) + p, err := NewAnthropic(context.Background(), config.Anthropic{}, bedrockCfg(profileARN, "anthropic.claude-haiku-4-5"), NewInferenceProfileCache()) + require.NoError(t, err) + + _, err = interceptFor(t, p, "claude-opus-4-8") require.ErrorContains(t, err, "references no model") }) t.Run("small fast profile resolves independently", func(t *testing.T) { - const smallFastProfileARN = "arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/8x1qk20fzp3r" - - url, paths := mockBedrock(t, func(w http.ResponseWriter, _ *http.Request) { - w.Header().Set("Content-Type", "application/json") - _, _ = w.Write([]byte(`{"models":[{"modelArn":"arn:aws:bedrock:us-east-1::foundation-model/anthropic.claude-haiku-4-5"}]}`)) - }) + url, paths := mockBedrock(t, respondWithModel("arn:aws:bedrock:us-east-1::foundation-model/anthropic.claude-haiku-4-5")) t.Setenv("AWS_ENDPOINT_URL_BEDROCK", url) - p, err := NewAnthropic(context.Background(), config.Anthropic{}, bedrockCfg("eu.anthropic.claude-opus-4-8", smallFastProfileARN)) + p, err := NewAnthropic(context.Background(), config.Anthropic{}, bedrockCfg("eu.anthropic.claude-opus-4-8", smallFastProfileARN), NewInferenceProfileCache()) require.NoError(t, err) - require.Equal(t, "eu.anthropic.claude-opus-4-8", p.bedrock.ResolvedModel()) - require.Equal(t, "anthropic.claude-haiku-4-5", p.bedrock.ResolvedSmallFastModel()) - require.Equal(t, smallFastProfileARN, p.bedrock.ConfiguredSmallFastModel()) + + interceptor, err := interceptFor(t, p, "claude-haiku-4-5") + require.NoError(t, err) + require.Equal(t, "anthropic.claude-haiku-4-5", interceptor.Model()) require.Len(t, *paths, 1, "only the small fast profile ARN is resolved") require.Contains(t, (*paths)[0], smallFastProfileARN) }) @@ -221,10 +287,58 @@ func TestNewAnthropic_InferenceProfileResolution(t *testing.T) { }) t.Setenv("AWS_ENDPOINT_URL_BEDROCK", url) - p, err := NewAnthropic(context.Background(), config.Anthropic{}, bedrockCfg("eu.anthropic.claude-opus-4-8", "anthropic.claude-haiku-4-5")) + p, err := NewAnthropic(context.Background(), config.Anthropic{}, bedrockCfg("eu.anthropic.claude-opus-4-8", "anthropic.claude-haiku-4-5"), NewInferenceProfileCache()) require.NoError(t, err) - require.Equal(t, "eu.anthropic.claude-opus-4-8", p.bedrock.ResolvedModel()) - require.Equal(t, "eu.anthropic.claude-opus-4-8", p.bedrock.ConfiguredModel()) + + interceptor, err := interceptFor(t, p, "claude-opus-4-8") + require.NoError(t, err) + require.Equal(t, "eu.anthropic.claude-opus-4-8", interceptor.Model()) require.Empty(t, *paths) }) } + +// TestInferenceProfileCacheSharesConcurrentResolutions verifies that a burst of +// requests for an unresolved profile issues a single AWS lookup. +func TestInferenceProfileCacheSharesConcurrentResolutions(t *testing.T) { + const profileARN = "arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/46u2vhiyo6z5" + + var calls atomic.Int64 + release := make(chan struct{}) + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + calls.Add(1) + <-release + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"models":[{"modelArn":"arn:aws:bedrock:us-east-1::foundation-model/anthropic.claude-opus-4-8"}]}`)) + })) + t.Cleanup(srv.Close) + t.Setenv("AWS_ENDPOINT_URL_BEDROCK", srv.URL) + + awsCfg, err := buildBedrockCredentials(context.Background(), config.AWSBedrock{ + Region: "us-east-1", + AccessKey: "test-key", + AccessKeySecret: "test-secret", + }) + require.NoError(t, err) + + cache := NewInferenceProfileCache() + const callers = 8 + results := make(chan error, callers) + for range callers { + go func() { + model, err := cache.Resolve(context.Background(), awsCfg, profileARN) + if err == nil && model != "anthropic.claude-opus-4-8" { + err = xerrors.Errorf("unexpected model %q", model) + } + results <- err + }() + } + + // Hold the handler until every caller has joined the in-flight resolution. + require.Eventually(t, func() bool { return calls.Load() == 1 }, testutil.WaitShort, testutil.IntervalFast) + close(release) + + for range callers { + require.NoError(t, <-results) + } + require.Equal(t, int64(1), calls.Load()) +} diff --git a/cli/aibridged.go b/cli/aibridged.go index 1dbc443c426..1e125aecebf 100644 --- a/cli/aibridged.go +++ b/cli/aibridged.go @@ -97,6 +97,9 @@ type poolRPCReloader struct { logger slog.Logger aibridgeMetrics *aibridge.Metrics providerMetrics *aibridged.Metrics + // profiles is created once per reloader so Bedrock application inference + // profile resolutions survive the provider rebuild each reload performs. + profiles *aibridge.InferenceProfileCache } // NewPoolRPCReloader builds an [aibridged.ProviderReloader] that fetches the @@ -119,6 +122,7 @@ func NewPoolRPCReloader( logger: logger, aibridgeMetrics: aibridgeMetrics, providerMetrics: providerMetrics, + profiles: aibridge.NewInferenceProfileCache(), } } @@ -136,7 +140,7 @@ func (r *poolRPCReloader) Reload(ctx context.Context) error { // beyond the operator's actual misconfiguration. return xerrors.Errorf("fetch ai providers: %w", err) } - providers, outcomes := BuildProvidersFromProto(ctx, resp.GetProviders(), r.cfg, r.logger, r.aibridgeMetrics) + providers, outcomes := BuildProvidersFromProto(ctx, resp.GetProviders(), r.cfg, r.logger, r.aibridgeMetrics, r.profiles) r.pool.ReplaceProviders(providers) r.providerMetrics.RecordReloadSuccess(outcomes) return nil @@ -152,7 +156,7 @@ func (r *poolRPCReloader) Reload(ctx context.Context) error { // excluded from the returned snapshot; this keeps a single misconfigured // provider from taking the whole daemon down. The returned outcomes mirror the // per-provider status for metrics reporting. -func BuildProvidersFromProto(ctx context.Context, protoProviders []*proto.AIProvider, cfg codersdk.AIBridgeConfig, logger slog.Logger, metrics *aibridge.Metrics) ([]aibridge.Provider, []aibridged.ProviderOutcome) { +func BuildProvidersFromProto(ctx context.Context, protoProviders []*proto.AIProvider, cfg codersdk.AIBridgeConfig, logger slog.Logger, metrics *aibridge.Metrics, profiles *aibridge.InferenceProfileCache) ([]aibridge.Provider, []aibridged.ProviderOutcome) { providers := make([]aibridge.Provider, 0, len(protoProviders)) outcomes := make([]aibridged.ProviderOutcome, 0, len(protoProviders)) enabledCount := 0 @@ -165,7 +169,7 @@ func BuildProvidersFromProto(ctx context.Context, protoProviders []*proto.AIProv if spec.Enabled { enabledCount++ } - prov, err := buildProvider(ctx, spec, cfg, metrics) + prov, err := buildProvider(ctx, spec, cfg, metrics, profiles) if err != nil { outcome.Status = aibridged.ProviderStatusError outcome.Err = err @@ -239,7 +243,7 @@ type aiProviderSpec struct { // buildProvider constructs the appropriate [aibridge.Provider] for a // single provider spec, independent of where the spec was sourced from. -func buildProvider(ctx context.Context, spec aiProviderSpec, cfg codersdk.AIBridgeConfig, metrics *aibridge.Metrics) (aibridge.Provider, error) { +func buildProvider(ctx context.Context, spec aiProviderSpec, cfg codersdk.AIBridgeConfig, metrics *aibridge.Metrics, profiles *aibridge.InferenceProfileCache) (aibridge.Provider, error) { if !spec.Enabled { return aibridge.NewDisabledProviderStub(spec.Name, string(spec.Type)), nil } @@ -311,7 +315,7 @@ func buildProvider(ctx context.Context, spec aiProviderSpec, cfg codersdk.AIBrid APIDumpDir: dumpDir, CircuitBreaker: cbCfg, SendActorHeaders: sendActorHeaders, - }, bedrock) + }, bedrock, profiles) case database.AIProviderTypeCopilot: // Copilot is always BYOK; the per-user token is supplied on each diff --git a/cli/aibridged_internal_test.go b/cli/aibridged_internal_test.go index 7cd4f64d742..3a2e0b7f12a 100644 --- a/cli/aibridged_internal_test.go +++ b/cli/aibridged_internal_test.go @@ -69,7 +69,7 @@ func buildFromDB(ctx context.Context, t *testing.T, db database.Store, cfg coder if err != nil { return nil, nil, err } - providers, outcomes := BuildProvidersFromProto(ctx, resp.GetProviders(), cfg, logger, nil) + providers, outcomes := BuildProvidersFromProto(ctx, resp.GetProviders(), cfg, logger, nil, aibridge.NewInferenceProfileCache()) return providers, outcomes, nil } diff --git a/cli/server_aibridge_internal_test.go b/cli/server_aibridge_internal_test.go index 781b642f938..f3ce6961ff9 100644 --- a/cli/server_aibridge_internal_test.go +++ b/cli/server_aibridge_internal_test.go @@ -742,7 +742,7 @@ func TestBuildProviderFromProtoSetsAPIDumpDir(t *testing.T) { provider, err := buildProvider(t.Context(), protoToProviderSpec(tt.provider), codersdk.AIBridgeConfig{ AllowBYOK: serpent.Bool(true), APIDumpDir: serpent.String(dumpDir), - }, nil) + }, nil, aibridge.NewInferenceProfileCache()) require.NoError(t, err) assert.Equal(t, dumpDir, provider.APIDumpDir()) assert.Equal(t, tt.expectedType, provider.Type()) @@ -760,7 +760,7 @@ func TestBuildProviderFromProtoBedrockWithoutSettings(t *testing.T) { BaseUrl: "https://bedrock-runtime.us-east-1.amazonaws.com/", }), codersdk.AIBridgeConfig{ AllowBYOK: serpent.Bool(true), - }, nil) + }, nil, aibridge.NewInferenceProfileCache()) require.Error(t, err) assert.Contains(t, err.Error(), "bedrock provider has no bedrock credentials configured") } diff --git a/docs/ai-coder/ai-gateway/providers.md b/docs/ai-coder/ai-gateway/providers.md index 7d002271374..c7baddbf60a 100644 --- a/docs/ai-coder/ai-gateway/providers.md +++ b/docs/ai-coder/ai-gateway/providers.md @@ -244,7 +244,9 @@ usage pricing. Resolution requires a `GetInferenceProfile` call, so the AWS identity used by the gateway must have `bedrock:GetInferenceProfile` permission for the profile. Providers configured with plain model identifiers do not need this -permission. If resolution fails, the provider is skipped. +permission. The gateway resolves a profile on the first request that uses it +and caches the result, so a failed resolution fails that request and the next +request tries again. ### GitHub Copilot