diff --git a/aibridge/intercept/messages/base.go b/aibridge/intercept/messages/base.go index 3b77f752a62..67b19b1bdbc 100644 --- a/aibridge/intercept/messages/base.go +++ b/aibridge/intercept/messages/base.go @@ -72,6 +72,49 @@ var bedrockSupportedBetaFlags = map[string]bool{ type BedrockRuntime struct { Cfg aibconfig.AWSBedrock Creds aws.CredentialsProvider + + 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. +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 +} + +// ConfiguredSmallFastModel is [BedrockRuntime.ConfiguredModel] for the +// small/fast model. +func (b *BedrockRuntime) ConfiguredSmallFastModel() string { + return b.Cfg.SmallFastModel +} + +// 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 +} + +// ResolvedSmallFastModel is [BedrockRuntime.ResolvedModel] for the small/fast +// model. +func (b *BedrockRuntime) ResolvedSmallFastModel() string { + return b.resolvedSmallFastModel } type interceptionBase struct { @@ -86,6 +129,14 @@ type interceptionBase struct { // clientHeaders are the original HTTP headers from the client request. clientHeaders http.Header + // 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 @@ -169,9 +220,9 @@ 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 + model := i.bedrock.ResolvedModel() + if i.isSmallFastModel { + model = i.bedrock.ResolvedSmallFastModel() } return model } @@ -179,6 +230,16 @@ 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. +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), @@ -261,12 +322,9 @@ 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 strings.Contains(i.reqPayload.model(), "haiku") +// 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 +473,16 @@ 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 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 { return } model := i.Model() - updated, err := i.reqPayload.withModel(model) + 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 diff --git a/aibridge/intercept/messages/base_internal_test.go b/aibridge/intercept/messages/base_internal_test.go index b5b400dd92c..38c7d185c00 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()) @@ -209,6 +206,143 @@ 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 := 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 + expectConfigured string + }{ + { + name: "primary model", + expectModel: "anthropic.claude-opus-4-8", + expectConfigured: profileARN, + }, + { + name: "small fast model", + smallFast: true, + expectModel: "anthropic.claude-haiku-4-5", + expectConfigured: 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, + isSmallFastModel: tt.smallFast, + logger: slog.Make(), + } + + require.Equal(t, tt.expectModel, i.Model()) + + i.augmentRequestForBedrockInvokeModel() + 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()) + }) + } +} + +// 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) { + t.Parallel() + + i := &interceptionBase{ + reqPayload: mustMessagesPayload(t, `{"model":"claude-opus-4-8","max_tokens":10000}`), + 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(), + } + + require.Equal(t, "eu.anthropic.claude-opus-4-8", i.Model()) +} + func TestAccumulateUsage(t *testing.T) { t.Parallel() @@ -608,6 +742,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 +895,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", @@ -804,14 +957,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", - }, - }, + 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(), } @@ -832,7 +990,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. @@ -1159,14 +1318,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(), } @@ -1217,10 +1373,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/intercept/messages/blocking.go b/aibridge/intercept/messages/blocking.go index ecebb76981e..98d408bc7d8 100644 --- a/aibridge/intercept/messages/blocking.go +++ b/aibridge/intercept/messages/blocking.go @@ -41,13 +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, + 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 70ddc65bdb0..a526015f83c 100644 --- a/aibridge/intercept/messages/streaming.go +++ b/aibridge/intercept/messages/streaming.go @@ -46,13 +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, + id: id, + reqPayload: reqPayload, + cfg: cfg, + cred: cred, + bedrock: bedrock, + clientHeaders: clientHeaders, + tracer: tracer, + isSmallFastModel: isSmallFastModel(reqPayload.model()), }} } @@ -115,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() } diff --git a/aibridge/provider/anthropic.go b/aibridge/provider/anthropic.go index 460b38102cd..e8796fde256 100644 --- a/aibridge/provider/anthropic.go +++ b/aibridge/provider/anthropic.go @@ -69,20 +69,30 @@ 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) } - 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. + 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) } 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 new file mode 100644 index 00000000000..b0ed2f55664 --- /dev/null +++ b/aibridge/provider/bedrock_inference_profile.go @@ -0,0 +1,112 @@ +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" + +// 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, 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 { + return false + } + resourceType, _, ok := strings.Cut(parsed.Resource, "/") + return ok && resourceType == applicationInferenceProfileResourceType +} + +// 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. +func resolveInferenceProfile(ctx context.Context, awsCfg aws.Config, profileARN string) (string, error) { + client := bedrock.NewFromConfig(awsCfg) + + out, err := client.GetInferenceProfile(ctx, &bedrock.GetInferenceProfileInput{ + InferenceProfileIdentifier: aws.String(profileARN), + }) + if err != nil { + 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) + } + + 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, 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 new file mode 100644 index 00000000000..2ac34857e97 --- /dev/null +++ b/aibridge/provider/bedrock_inference_profile_internal_test.go @@ -0,0 +1,230 @@ +package provider + +import ( + "context" + "net/http" + "net/http/httptest" + "testing" + + "github.com/stretchr/testify/require" + + "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) + }) + } +} + +// 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) { + const profileARN = "arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/46u2vhiyo6z5" + + bedrockCfg := func(model, smallFastModel string) *config.AWSBedrock { + return &config.AWSBedrock{ + Region: "us-east-1", + AccessKey: "test-key", + AccessKeySecret: "test-secret", + Model: model, + SmallFastModel: smallFastModel, + } + } + + // 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 + } + + 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, "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. + 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) { + 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, "anthropic.claude-haiku-4-5")) + require.ErrorContains(t, err, "resolve bedrock models") + require.ErrorContains(t, err, "GetInferenceProfile") + }) + + 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) + + _, err := NewAnthropic(context.Background(), config.Anthropic{}, bedrockCfg(profileARN, "anthropic.claude-haiku-4-5")) + 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) + + 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()) + 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") + }) + 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")) + 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) + }) +} 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) } diff --git a/docs/ai-coder/ai-gateway/providers.md b/docs/ai-coder/ai-gateway/providers.md index 47af0656143..7d002271374 100644 --- a/docs/ai-coder/ai-gateway/providers.md +++ b/docs/ai-coder/ai-gateway/providers.md @@ -230,6 +230,22 @@ 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, 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 GitHub Copilot offers three plans: Individual, Business, and Enterprise, 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=