Skip to content

Commit b2ee4ad

Browse files
kw-bivanovkw-bivanov
committed
fix(bedrock): resolve custom ARN capabilities from the underlying model (#1968)
1 parent 2baac5e commit b2ee4ad

34 files changed

Lines changed: 832 additions & 20 deletions

‎packages/types/src/__tests__/bedrock-custom-arn.test.ts‎

Lines changed: 49 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,49 @@
1+
import { resolveBedrockCustomArnBaseModelId } from "../providers/bedrock.js"
2+
3+
const appProfileArn = "arn:aws:bedrock:us-west-2:123456789012:application-inference-profile/abcd1234efgh"
4+
5+
describe("resolveBedrockCustomArnBaseModelId", () => {
6+
it("returns the explicit base model when it is a known Bedrock model", () => {
7+
expect(resolveBedrockCustomArnBaseModelId(appProfileArn, "anthropic.claude-opus-5-5")).toBe(
8+
"anthropic.claude-opus-5-5",
9+
)
10+
})
11+
12+
it("prefers the explicit base model over one named in the ARN", () => {
13+
const arn = "arn:aws:bedrock:us-east-1:123456789012:inference-profile/us.anthropic.claude-opus-4-8"
14+
expect(resolveBedrockCustomArnBaseModelId(arn, "meta.llama3-3-70b-instruct-v1:0")).toBe(
15+
"meta.llama3-3-70b-instruct-v1:0",
16+
)
17+
})
18+
19+
it("detects a model named at the end of the ARN", () => {
20+
expect(
21+
resolveBedrockCustomArnBaseModelId(
22+
"arn:aws:bedrock:us-east-1:123456789012:inference-profile/us.anthropic.claude-opus-4-8",
23+
undefined,
24+
),
25+
).toBe("anthropic.claude-opus-4-8")
26+
expect(
27+
resolveBedrockCustomArnBaseModelId(
28+
"arn:aws:bedrock:us-east-1::foundation-model/anthropic.claude-3-5-sonnet-20241022-v2:0",
29+
"",
30+
),
31+
).toBe("anthropic.claude-3-5-sonnet-20241022-v2:0")
32+
})
33+
34+
it("returns undefined when neither the setting nor the ARN identifies a model", () => {
35+
expect(resolveBedrockCustomArnBaseModelId(appProfileArn, undefined)).toBeUndefined()
36+
expect(resolveBedrockCustomArnBaseModelId(appProfileArn, "")).toBeUndefined()
37+
expect(resolveBedrockCustomArnBaseModelId(appProfileArn, "not-a-bedrock-model")).toBeUndefined()
38+
expect(resolveBedrockCustomArnBaseModelId(undefined, undefined)).toBeUndefined()
39+
})
40+
41+
it("does not confuse models whose IDs share a suffix", () => {
42+
expect(
43+
resolveBedrockCustomArnBaseModelId(
44+
"arn:aws:bedrock:us-east-1:123456789012:inference-profile/us.amazon.nova-2-lite-v1:0",
45+
undefined,
46+
),
47+
).toBe("amazon.nova-2-lite-v1:0")
48+
})
49+
})

‎packages/types/src/provider-settings/bedrock.ts‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@ export const bedrockProviderDefinition = createProviderDefinition({
2626
awsApiKey: z.string().optional(),
2727
awsUseApiKey: z.boolean().optional(),
2828
awsCustomArn: z.string().optional(),
29+
awsCustomArnBaseModelId: z.string().optional(), // Registry model a custom ARN points to; drives capabilities and pricing
2930
awsModelContextWindow: z.number().optional(),
3031
awsBedrockEndpointEnabled: z.boolean().optional(),
3132
awsBedrockEndpoint: z.string().optional(),

‎packages/types/src/providers/bedrock.ts‎

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -739,3 +739,27 @@ export const BEDROCK_SERVICE_TIER_PRICING = {
739739
FLEX: 0.5, // 50% discount from standard
740740
PRIORITY: 1.75, // 75% premium over standard
741741
} as const
742+
743+
const isBedrockModelId = (id: string): id is BedrockModelId => Object.prototype.hasOwnProperty.call(bedrockModels, id)
744+
745+
/**
746+
* Resolves the registry model a custom ARN points to: the user's explicit choice when it is a known
747+
* model, otherwise a model named at the end of the ARN (e.g. `inference-profile/us.anthropic.claude-opus-4-8`).
748+
* Returns undefined when the ARN doesn't identify its model (e.g. application inference profiles).
749+
*/
750+
export const resolveBedrockCustomArnBaseModelId = (
751+
customArn: string | undefined,
752+
baseModelId: string | undefined,
753+
): BedrockModelId | undefined => {
754+
if (baseModelId && isBedrockModelId(baseModelId)) {
755+
return baseModelId
756+
}
757+
758+
if (!customArn) {
759+
return undefined
760+
}
761+
762+
return (Object.keys(bedrockModels) as BedrockModelId[]).find(
763+
(id) => customArn.endsWith(`/${id}`) || customArn.endsWith(`.${id}`),
764+
)
765+
}

‎src/api/providers/__tests__/bedrock.spec.ts‎

Lines changed: 136 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -60,6 +60,9 @@ import {
6060
bedrockModels,
6161
SERVICE_TIER_KEY,
6262
ApiProviderError,
63+
BEDROCK_DEFAULT_CONTEXT,
64+
BEDROCK_MAX_TOKENS,
65+
type ProviderSettings,
6366
} from "@roo-code/types"
6467

6568
import type { Anthropic } from "@anthropic-ai/sdk"
@@ -1948,6 +1951,139 @@ describe("AwsBedrockHandler", () => {
19481951
expect(commandArg.inferenceConfig?.temperature).toBeDefined()
19491952
})
19501953

1954+
describe("custom ARN base model", () => {
1955+
const appProfileArn = "arn:aws:bedrock:us-west-2:123456789012:application-inference-profile/abcd1234efgh"
1956+
const createCustomArnHandler = (options: Partial<ProviderSettings> = {}) =>
1957+
new AwsBedrockHandler({
1958+
apiModelId: "custom-arn",
1959+
awsCustomArn: appProfileArn,
1960+
awsAccessKey: "test-access-key",
1961+
awsSecretKey: "test-secret-key",
1962+
awsRegion: "us-west-2",
1963+
...options,
1964+
})
1965+
const sendFirstRequest = async (handler: AwsBedrockHandler) => {
1966+
await handler.createMessage("System prompt", messages).next()
1967+
return mockConverseStreamCommand.mock.calls[0][0]
1968+
}
1969+
1970+
it("uses Opus 5.5 capabilities: native 1M context, adaptive thinking, no temperature", async () => {
1971+
const handler = createCustomArnHandler({
1972+
awsCustomArnBaseModelId: "anthropic.claude-opus-5-5",
1973+
enableReasoningEffort: true,
1974+
})
1975+
1976+
expect(handler.getModel().info.contextWindow).toBe(1_000_000)
1977+
expect(handler.getModel().info.maxTokens).toBe(bedrockModels["anthropic.claude-opus-5-5"].maxTokens)
1978+
1979+
const input = await sendFirstRequest(handler)
1980+
expect(input).toMatchObject({
1981+
modelId: appProfileArn,
1982+
additionalModelRequestFields: {
1983+
thinking: { type: "adaptive", display: "summarized" },
1984+
output_config: { effort: "xhigh" },
1985+
},
1986+
})
1987+
expect(input.additionalModelRequestFields).not.toMatchObject({
1988+
anthropic_beta: expect.arrayContaining(["context-1m-2025-08-07"]),
1989+
})
1990+
expect(input.inferenceConfig?.temperature).toBeUndefined()
1991+
})
1992+
1993+
it("sends the 1M context beta for an Opus 4.8 base model when enabled", async () => {
1994+
const handler = createCustomArnHandler({
1995+
awsCustomArnBaseModelId: "anthropic.claude-opus-4-8",
1996+
awsBedrock1MContext: true,
1997+
})
1998+
1999+
expect(handler.getModel().info.contextWindow).toBe(1_000_000)
2000+
const input = await sendFirstRequest(handler)
2001+
expect(input).toMatchObject({
2002+
modelId: appProfileArn,
2003+
additionalModelRequestFields: {
2004+
anthropic_beta: expect.arrayContaining(["context-1m-2025-08-07"]),
2005+
},
2006+
})
2007+
})
2008+
2009+
it("keeps budget thinking and temperature for a Sonnet 4.5 base model", async () => {
2010+
const input = await sendFirstRequest(
2011+
createCustomArnHandler({
2012+
awsCustomArnBaseModelId: "anthropic.claude-sonnet-4-5-20250929-v1:0",
2013+
enableReasoningEffort: true,
2014+
}),
2015+
)
2016+
2017+
expect(input).toMatchObject({ additionalModelRequestFields: { thinking: { type: "enabled" } } })
2018+
expect(input.inferenceConfig?.temperature).toBeDefined()
2019+
})
2020+
2021+
it("sends no Anthropic fields for a non-Anthropic base model", async () => {
2022+
const handler = createCustomArnHandler({ awsCustomArnBaseModelId: "meta.llama3-3-70b-instruct-v1:0" })
2023+
2024+
const llama = bedrockModels["meta.llama3-3-70b-instruct-v1:0"]
2025+
expect(handler.getModel().info).toMatchObject({
2026+
contextWindow: llama.contextWindow,
2027+
maxTokens: llama.maxTokens,
2028+
inputPrice: llama.inputPrice,
2029+
})
2030+
const input = await sendFirstRequest(handler)
2031+
expect(input.modelId).toBe(appProfileArn)
2032+
expect(input.additionalModelRequestFields).toBeUndefined()
2033+
expect(input.inferenceConfig?.temperature).toBeDefined()
2034+
})
2035+
2036+
it("uses the generic fallback with user limits when the model is not listed", () => {
2037+
expect(createCustomArnHandler().getModel().info).toMatchObject({
2038+
contextWindow: BEDROCK_DEFAULT_CONTEXT,
2039+
maxTokens: BEDROCK_MAX_TOKENS,
2040+
})
2041+
expect(
2042+
createCustomArnHandler({ awsModelContextWindow: 32_000, modelMaxTokens: 2048 }).getModel().info,
2043+
).toMatchObject({ contextWindow: 32_000, maxTokens: 2048 })
2044+
})
2045+
2046+
it("applies capabilities of a model named in an inference profile ARN", async () => {
2047+
const input = await sendFirstRequest(
2048+
createCustomArnHandler({
2049+
awsCustomArn:
2050+
"arn:aws:bedrock:us-east-1:123456789012:inference-profile/us.anthropic.claude-opus-4-8",
2051+
awsRegion: "us-east-1",
2052+
awsBedrock1MContext: true,
2053+
}),
2054+
)
2055+
2056+
expect(input).toMatchObject({
2057+
additionalModelRequestFields: { anthropic_beta: expect.arrayContaining(["context-1m-2025-08-07"]) },
2058+
})
2059+
expect(input.inferenceConfig?.temperature).toBeUndefined()
2060+
})
2061+
2062+
it("completePrompt omits temperature for an adaptive base model (non-stream path)", async () => {
2063+
const mockConverseCommand = vi.mocked(ConverseCommand)
2064+
mockConverseCommand.mockClear()
2065+
2066+
await createCustomArnHandler({ awsCustomArnBaseModelId: "anthropic.claude-opus-5-5" }).completePrompt(
2067+
"Test prompt",
2068+
)
2069+
2070+
const input = mockConverseCommand.mock.calls[0][0]
2071+
expect(input.modelId).toBe(appProfileArn)
2072+
expect(input.inferenceConfig?.temperature).toBeUndefined()
2073+
})
2074+
2075+
it("ignores the base model for foundation-model ARNs, which already name the invoked model", () => {
2076+
const handler = createCustomArnHandler({
2077+
awsCustomArn:
2078+
"arn:aws:bedrock:us-east-1::foundation-model/anthropic.claude-3-5-sonnet-20241022-v2:0",
2079+
awsRegion: "us-east-1",
2080+
awsCustomArnBaseModelId: "anthropic.claude-opus-5-5",
2081+
})
2082+
2083+
expect(handler.getModel().id).toBe("anthropic.claude-3-5-sonnet-20241022-v2:0")
2084+
})
2085+
})
2086+
19512087
describe("isAdaptiveThinkingModel detection", () => {
19522088
// Unit-cover the private guard directly (same pattern the suite uses for
19532089
// parseBaseModelId / getPrefixForRegion). This exercises all four model

‎src/api/providers/bedrock.ts‎

Lines changed: 26 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,7 @@ import {
3636
BEDROCK_SERVICE_TIER_PRICING,
3737
SERVICE_TIER_KEY,
3838
ApiProviderError,
39+
resolveBedrockCustomArnBaseModelId,
3940
} from "@roo-code/types"
4041
import { TelemetryService } from "@roo-code/telemetry"
4142

@@ -342,7 +343,7 @@ export class AwsBedrockHandler extends BaseProvider implements SingleCompletionH
342343
* the prefix via parseBaseModelId before matching.
343344
*/
344345
private isAdaptiveThinkingModel(modelId: string): boolean {
345-
const baseModelId = this.parseBaseModelId(modelId)
346+
const baseModelId = this.getCapabilityModelId(modelId)
346347
return (
347348
baseModelId.includes("opus-4-7") ||
348349
baseModelId.includes("opus-4-8") ||
@@ -452,7 +453,7 @@ export class AwsBedrockHandler extends BaseProvider implements SingleCompletionH
452453
// Detect models that require the adaptive-thinking API contract (Opus/Sonnet
453454
// 4.7 and 4.8). See isAdaptiveThinkingModel for details. The same guard is
454455
// reused in completePrompt so both request paths stay consistent.
455-
const baseModelId = this.parseBaseModelId(modelConfig.id)
456+
const baseModelId = this.getCapabilityModelId(modelConfig.id)
456457
const isAdaptiveThinkingModel = this.isAdaptiveThinkingModel(modelConfig.id)
457458

458459
// Determine if thinking should be enabled
@@ -1138,6 +1139,22 @@ export class AwsBedrockHandler extends BaseProvider implements SingleCompletionH
11381139
}
11391140
}
11401141

1142+
private getCustomArnBaseModelId(): BedrockModelId | undefined {
1143+
return resolveBedrockCustomArnBaseModelId(this.options.awsCustomArn, this.options.awsCustomArnBaseModelId)
1144+
}
1145+
1146+
/**
1147+
* Model ID used for capability checks (thinking contract, 1M context, service tier, ...).
1148+
* For a custom ARN this is the model it points to rather than the ARN itself.
1149+
*/
1150+
private getCapabilityModelId(modelId: string): string {
1151+
if (this.options.awsCustomArn && modelId === this.options.awsCustomArn) {
1152+
return this.getCustomArnBaseModelId() ?? this.arnInfo?.modelId ?? modelId
1153+
}
1154+
1155+
return this.parseBaseModelId(modelId)
1156+
}
1157+
11411158
//This strips any region prefix that used on cross-region model inference ARNs
11421159
private parseBaseModelId(modelId: string): string {
11431160
if (!modelId) {
@@ -1225,7 +1242,10 @@ export class AwsBedrockHandler extends BaseProvider implements SingleCompletionH
12251242

12261243
// If custom ARN is provided, use it
12271244
if (this.options.awsCustomArn) {
1228-
modelConfig = this.getModelById(this.arnInfo.modelId, this.arnInfo.modelType)
1245+
// A foundation-model ARN already names the model that will be invoked.
1246+
const baseModelId =
1247+
this.arnInfo.modelType === "foundation-model" ? undefined : this.getCustomArnBaseModelId()
1248+
modelConfig = this.getModelById(baseModelId ?? this.arnInfo.modelId, this.arnInfo.modelType)
12291249

12301250
//If the user entered an ARN for a foundation-model they've done the same thing as picking from our list of options.
12311251
//We leave the model data matching the same as if a drop-down input method was used by not overwriting the model ID with the user input ARN
@@ -1253,8 +1273,8 @@ export class AwsBedrockHandler extends BaseProvider implements SingleCompletionH
12531273
}
12541274

12551275
// Check if 1M context is enabled for supported Claude 4 models
1256-
// Use parseBaseModelId to handle cross-region inference prefixes
1257-
const baseModelId = this.parseBaseModelId(modelConfig.id)
1276+
// getCapabilityModelId handles cross-region prefixes and custom ARNs
1277+
const baseModelId = this.getCapabilityModelId(modelConfig.id)
12581278
if (BEDROCK_1M_CONTEXT_MODEL_IDS.includes(baseModelId as any) && this.options.awsBedrock1MContext) {
12591279
// Update context window and pricing to 1M tier when 1M context beta is enabled
12601280
const tier = modelConfig.info.tiers?.[0]
@@ -1278,7 +1298,7 @@ export class AwsBedrockHandler extends BaseProvider implements SingleCompletionH
12781298
})
12791299

12801300
// Apply service tier pricing if specified and model supports it
1281-
const baseModelIdForTier = this.parseBaseModelId(modelConfig.id)
1301+
const baseModelIdForTier = this.getCapabilityModelId(modelConfig.id)
12821302
if (this.options.awsBedrockServiceTier && BEDROCK_SERVICE_TIER_MODEL_IDS.includes(baseModelIdForTier as any)) {
12831303
const pricingMultiplier = BEDROCK_SERVICE_TIER_PRICING[this.options.awsBedrockServiceTier]
12841304
if (pricingMultiplier && pricingMultiplier !== 1.0) {

‎webview-ui/src/components/settings/ApiOptions.tsx‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@ import {
2525
getStaticModelsForProvider,
2626
shouldUseGenericModelPicker,
2727
handleModelChangeSideEffects,
28+
toBedrockModelId,
2829
} from "./utils/providerModelConfig"
2930

3031
import { vscode } from "@src/utils/vscode"
@@ -756,6 +757,11 @@ const ApiOptions = ({
756757
organizationAllowList={organizationAllowList}
757758
errorMessage={modelValidationError}
758759
simplifySettings={fromWelcomeView}
760+
valueTransform={
761+
activeSelectedProvider === providerIdentifiers.bedrock
762+
? toBedrockModelId
763+
: undefined
764+
}
759765
onModelChange={(modelId) =>
760766
handleModelChangeSideEffects(
761767
activeSelectedProvider,

0 commit comments

Comments
 (0)