fix: resolve Bedrock application inference profile ARNs #28877
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -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 | ||
|
Contributor
There was a problem hiding this comment. P3 [CRF-17] The
|
||
| // (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,16 +220,26 @@ 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 { | ||
| model = i.bedrock.ResolvedSmallFastModel() | ||
| } | ||
| return model | ||
| } | ||
|
|
||
| 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 { | ||
|
Contributor
There was a problem hiding this comment. There is already
Contributor
Author
There was a problem hiding this comment. @pawbana The idea is that we need two functions:
In other words:
|
||
| 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 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 { | ||
|
Contributor
There was a problem hiding this comment. P3 [CRF-20] The single call site checks
|
||
| return | ||
| } | ||
|
|
||
| model := i.Model() | ||
|
Contributor
There was a problem hiding this comment. Is
Contributor
Author
There was a problem hiding this comment. No, I explained the difference above. |
||
| 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 | ||
|
|
||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
P3 [CRF-3]
BedrockRuntimekeepsCfgandCredsexported while the resolved models are unexported, so a struct literal still compiles and yields an emptyModel(). (Netero)NewBedrockRuntimeis the only way to setresolvedModel, but nothing prevents&BedrockRuntime{Cfg: ..., Creds: ...}. On the InvokeModel pathModel()then returns"", which flows into usage records, pricing lookups, metrics, and the apidump middleware. Today onlyprovider/anthropic.go:95constructs one, so this is latent, not live, but the constructor was introduced to carry an invariant it does not enforce. Fix: unexportCfgandCreds; every in-package read already goes through the receiver, and no other package touches them directly.There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Fair point, but I'd like to keep it as is for now. I think this pattern is pretty common in our codebase. Also, Go generally recommends using exported fields directly instead of getters.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
P3 [CRF-3] Panel disposition: held at P3. The 'exported fields are idiomatic' defense answers a question the finding did not ask: this is not fields-versus-getters, it is a type with two exported and two unexported fields with a dependency between them, whose constructor doc asserts an invariant the type does not hold. New evidence defeats the 'latent, needs a struct literal' framing: no struct literal is required,
base_internal_test.goalready callsNewBedrockRuntime(tt.cfg, ..., "", "")withtt.cfg.Modelset, so the sanctioned constructor already accepts the state the doc says cannot exist. And the round-3 'unreachable because Validate rejects empty Model' argument does not hold:Validatenever inspectsresolvedModel. Two coherent endpoints, either is fine: unexportCfg/Credsso the constructor enforces the invariant (the only cross-package reader isanthropic.go:197, needing one accessor), or drop theConfigured*accessors and readCfgdirectly (then CRF-30 disappears too). What should not ship is the current half-encapsulated shape where the doc claims an invariant the type does not enforce. At minimum, fix the type doc, which still enumerates only 'the static Bedrock config plus the AWS credentials provider.'