Skip to content

Commit 9b4c326

Browse files
committed
feat: add global inference support for AWS Bedrock models
- Add awsUseGlobalInference setting to provider settings schema - Add support for global. prefix in Bedrock provider for supported models - Add UI checkbox for global inference in Bedrock settings - Prioritize global inference over cross-region when both are enabled - Add comprehensive tests for global inference functionality - Support global inference in custom ARN parsing Supported models: - Claude Sonnet 4, 4.5 - Claude Opus 4, 4.1 - Claude 3.7 Sonnet - Claude Haiku 4.5 Fixes #8750
1 parent f3a505f commit 9b4c326

File tree

6 files changed

+362
-11
lines changed

6 files changed

+362
-11
lines changed

packages/types/src/provider-settings.ts

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -219,6 +219,7 @@ const bedrockSchema = apiModelIdProviderModelSchema.extend({
219219
awsSessionToken: z.string().optional(),
220220
awsRegion: z.string().optional(),
221221
awsUseCrossRegionInference: z.boolean().optional(),
222+
awsUseGlobalInference: z.boolean().optional(),
222223
awsUsePromptCache: z.boolean().optional(),
223224
awsProfile: z.string().optional(),
224225
awsUseProfile: z.boolean().optional(),

packages/types/src/providers/bedrock.ts

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -442,6 +442,21 @@ export const AWS_INFERENCE_PROFILE_MAPPING: Array<[string, string]> = [
442442
["sa-", "sa."],
443443
]
444444

445+
// Global Inference Profile prefix
446+
// https://docs.aws.amazon.com/bedrock/latest/userguide/global-inference.html
447+
export const AWS_GLOBAL_INFERENCE_PREFIX = "global."
448+
449+
// Models that support Global Inference
450+
// Based on AWS documentation, these models can use the global. prefix
451+
export const BEDROCK_GLOBAL_INFERENCE_MODEL_IDS = [
452+
"anthropic.claude-sonnet-4-20250514-v1:0",
453+
"anthropic.claude-sonnet-4-5-20250929-v1:0",
454+
"anthropic.claude-opus-4-20250514-v1:0",
455+
"anthropic.claude-opus-4-1-20250805-v1:0",
456+
"anthropic.claude-3-7-sonnet-20250219-v1:0",
457+
"anthropic.claude-haiku-4-5-20251001-v1:0",
458+
] as const
459+
445460
// Amazon Bedrock supported regions for the regions dropdown
446461
// Based on official AWS documentation
447462
export const BEDROCK_REGIONS = [
Lines changed: 272 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,272 @@
1+
// npx vitest run src/api/providers/__tests__/bedrock-global-inference.spec.ts
2+
3+
import { AWS_GLOBAL_INFERENCE_PREFIX, BEDROCK_GLOBAL_INFERENCE_MODEL_IDS } from "@roo-code/types"
4+
import { AwsBedrockHandler } from "../bedrock"
5+
import { ApiHandlerOptions } from "../../../shared/api"
6+
7+
// Mock AWS SDK
8+
vitest.mock("@aws-sdk/client-bedrock-runtime", () => {
9+
return {
10+
BedrockRuntimeClient: vitest.fn().mockImplementation(() => ({
11+
send: vitest.fn(),
12+
config: { region: "us-east-1" },
13+
})),
14+
ConverseCommand: vitest.fn(),
15+
ConverseStreamCommand: vitest.fn(),
16+
}
17+
})
18+
19+
describe("Amazon Bedrock Global Inference", () => {
20+
// Helper function to create a handler with specific options
21+
const createHandler = (options: Partial<ApiHandlerOptions> = {}) => {
22+
const defaultOptions: ApiHandlerOptions = {
23+
apiModelId: "anthropic.claude-sonnet-4-20250514-v1:0",
24+
awsRegion: "us-east-1",
25+
...options,
26+
}
27+
return new AwsBedrockHandler(defaultOptions)
28+
}
29+
30+
describe("AWS_GLOBAL_INFERENCE_PREFIX constant", () => {
31+
it("should have the correct global inference prefix", () => {
32+
expect(AWS_GLOBAL_INFERENCE_PREFIX).toBe("global.")
33+
})
34+
})
35+
36+
describe("BEDROCK_GLOBAL_INFERENCE_MODEL_IDS constant", () => {
37+
it("should contain the expected models that support global inference", () => {
38+
const expectedModels = [
39+
"anthropic.claude-sonnet-4-20250514-v1:0",
40+
"anthropic.claude-sonnet-4-5-20250929-v1:0",
41+
"anthropic.claude-opus-4-20250514-v1:0",
42+
"anthropic.claude-opus-4-1-20250805-v1:0",
43+
"anthropic.claude-3-7-sonnet-20250219-v1:0",
44+
"anthropic.claude-haiku-4-5-20251001-v1:0",
45+
]
46+
expect(BEDROCK_GLOBAL_INFERENCE_MODEL_IDS).toEqual(expectedModels)
47+
})
48+
})
49+
50+
describe("Global inference with supported models", () => {
51+
it("should apply global. prefix when global inference is enabled for Claude Sonnet 4", () => {
52+
const handler = createHandler({
53+
awsUseGlobalInference: true,
54+
apiModelId: "anthropic.claude-sonnet-4-20250514-v1:0",
55+
})
56+
57+
const model = handler.getModel()
58+
expect(model.id).toBe("global.anthropic.claude-sonnet-4-20250514-v1:0")
59+
})
60+
61+
it("should apply global. prefix when global inference is enabled for Claude Sonnet 4.5", () => {
62+
const handler = createHandler({
63+
awsUseGlobalInference: true,
64+
apiModelId: "anthropic.claude-sonnet-4-5-20250929-v1:0",
65+
})
66+
67+
const model = handler.getModel()
68+
expect(model.id).toBe("global.anthropic.claude-sonnet-4-5-20250929-v1:0")
69+
})
70+
71+
it("should apply global. prefix when global inference is enabled for Claude Opus 4", () => {
72+
const handler = createHandler({
73+
awsUseGlobalInference: true,
74+
apiModelId: "anthropic.claude-opus-4-20250514-v1:0",
75+
})
76+
77+
const model = handler.getModel()
78+
expect(model.id).toBe("global.anthropic.claude-opus-4-20250514-v1:0")
79+
})
80+
81+
it("should apply global. prefix when global inference is enabled for Claude Opus 4.1", () => {
82+
const handler = createHandler({
83+
awsUseGlobalInference: true,
84+
apiModelId: "anthropic.claude-opus-4-1-20250805-v1:0",
85+
})
86+
87+
const model = handler.getModel()
88+
expect(model.id).toBe("global.anthropic.claude-opus-4-1-20250805-v1:0")
89+
})
90+
91+
it("should apply global. prefix when global inference is enabled for Claude 3.7 Sonnet", () => {
92+
const handler = createHandler({
93+
awsUseGlobalInference: true,
94+
apiModelId: "anthropic.claude-3-7-sonnet-20250219-v1:0",
95+
})
96+
97+
const model = handler.getModel()
98+
expect(model.id).toBe("global.anthropic.claude-3-7-sonnet-20250219-v1:0")
99+
})
100+
101+
it("should apply global. prefix when global inference is enabled for Claude Haiku 4.5", () => {
102+
const handler = createHandler({
103+
awsUseGlobalInference: true,
104+
apiModelId: "anthropic.claude-haiku-4-5-20251001-v1:0",
105+
})
106+
107+
const model = handler.getModel()
108+
expect(model.id).toBe("global.anthropic.claude-haiku-4-5-20251001-v1:0")
109+
})
110+
})
111+
112+
describe("Global inference with unsupported models", () => {
113+
it("should NOT apply global. prefix for unsupported Claude 3 Sonnet model", () => {
114+
const handler = createHandler({
115+
awsUseGlobalInference: true,
116+
apiModelId: "anthropic.claude-3-sonnet-20240229-v1:0",
117+
})
118+
119+
const model = handler.getModel()
120+
expect(model.id).toBe("anthropic.claude-3-sonnet-20240229-v1:0")
121+
})
122+
123+
it("should NOT apply global. prefix for unsupported Claude 3 Haiku model", () => {
124+
const handler = createHandler({
125+
awsUseGlobalInference: true,
126+
apiModelId: "anthropic.claude-3-haiku-20240307-v1:0",
127+
})
128+
129+
const model = handler.getModel()
130+
expect(model.id).toBe("anthropic.claude-3-haiku-20240307-v1:0")
131+
})
132+
133+
it("should NOT apply global. prefix for Amazon Nova models", () => {
134+
const handler = createHandler({
135+
awsUseGlobalInference: true,
136+
apiModelId: "amazon.nova-pro-v1:0",
137+
})
138+
139+
const model = handler.getModel()
140+
expect(model.id).toBe("amazon.nova-pro-v1:0")
141+
})
142+
143+
it("should NOT apply global. prefix for Llama models", () => {
144+
const handler = createHandler({
145+
awsUseGlobalInference: true,
146+
apiModelId: "meta.llama3-1-70b-instruct-v1:0",
147+
})
148+
149+
const model = handler.getModel()
150+
expect(model.id).toBe("meta.llama3-1-70b-instruct-v1:0")
151+
})
152+
})
153+
154+
describe("Global inference priority over cross-region inference", () => {
155+
it("should prioritize global inference over cross-region inference when both are enabled", () => {
156+
const handler = createHandler({
157+
awsUseGlobalInference: true,
158+
awsUseCrossRegionInference: true,
159+
awsRegion: "us-east-1",
160+
apiModelId: "anthropic.claude-sonnet-4-20250514-v1:0",
161+
})
162+
163+
const model = handler.getModel()
164+
// Should use global. prefix, not us. prefix
165+
expect(model.id).toBe("global.anthropic.claude-sonnet-4-20250514-v1:0")
166+
})
167+
168+
it("should fall back to cross-region inference when global is disabled", () => {
169+
const handler = createHandler({
170+
awsUseGlobalInference: false,
171+
awsUseCrossRegionInference: true,
172+
awsRegion: "us-east-1",
173+
apiModelId: "anthropic.claude-sonnet-4-20250514-v1:0",
174+
})
175+
176+
const model = handler.getModel()
177+
// Should use us. prefix for cross-region inference
178+
expect(model.id).toBe("us.anthropic.claude-sonnet-4-20250514-v1:0")
179+
})
180+
181+
it("should apply no prefix when both global and cross-region are disabled", () => {
182+
const handler = createHandler({
183+
awsUseGlobalInference: false,
184+
awsUseCrossRegionInference: false,
185+
awsRegion: "us-east-1",
186+
apiModelId: "anthropic.claude-sonnet-4-20250514-v1:0",
187+
})
188+
189+
const model = handler.getModel()
190+
// Should have no prefix
191+
expect(model.id).toBe("anthropic.claude-sonnet-4-20250514-v1:0")
192+
})
193+
})
194+
195+
describe("Global inference with custom ARNs", () => {
196+
it("should parse global inference from ARN", () => {
197+
const handler = createHandler({
198+
awsCustomArn:
199+
"arn:aws:bedrock:us-east-1:123456789012:inference-profile/global.anthropic.claude-sonnet-4-20250514-v1:0",
200+
})
201+
202+
const model = handler.getModel()
203+
expect(model.id).toBe("global.anthropic.claude-sonnet-4-20250514-v1:0")
204+
})
205+
206+
it("should distinguish between global and cross-region prefixes in ARNs", () => {
207+
// Test global inference ARN
208+
const globalHandler = createHandler({
209+
awsCustomArn:
210+
"arn:aws:bedrock:us-east-1:123456789012:inference-profile/global.anthropic.claude-sonnet-4-20250514-v1:0",
211+
})
212+
const globalModel = globalHandler.getModel()
213+
expect(globalModel.id).toBe("global.anthropic.claude-sonnet-4-20250514-v1:0")
214+
215+
// Test cross-region inference ARN
216+
const crossRegionHandler = createHandler({
217+
awsCustomArn:
218+
"arn:aws:bedrock:us-east-1:123456789012:inference-profile/us.anthropic.claude-3-sonnet-20240229-v1:0",
219+
})
220+
const crossRegionModel = crossRegionHandler.getModel()
221+
expect(crossRegionModel.id).toBe("us.anthropic.claude-3-sonnet-20240229-v1:0")
222+
})
223+
})
224+
225+
describe("parseBaseModelId function", () => {
226+
it("should remove global. prefix from model IDs", () => {
227+
const handler = createHandler()
228+
const parseBaseModelId = (handler as any).parseBaseModelId.bind(handler)
229+
230+
expect(parseBaseModelId("global.anthropic.claude-sonnet-4-20250514-v1:0")).toBe(
231+
"anthropic.claude-sonnet-4-20250514-v1:0",
232+
)
233+
expect(parseBaseModelId("global.anthropic.claude-opus-4-20250514-v1:0")).toBe(
234+
"anthropic.claude-opus-4-20250514-v1:0",
235+
)
236+
})
237+
238+
it("should remove cross-region prefixes from model IDs", () => {
239+
const handler = createHandler()
240+
const parseBaseModelId = (handler as any).parseBaseModelId.bind(handler)
241+
242+
expect(parseBaseModelId("us.anthropic.claude-3-sonnet-20240229-v1:0")).toBe(
243+
"anthropic.claude-3-sonnet-20240229-v1:0",
244+
)
245+
expect(parseBaseModelId("eu.anthropic.claude-3-sonnet-20240229-v1:0")).toBe(
246+
"anthropic.claude-3-sonnet-20240229-v1:0",
247+
)
248+
expect(parseBaseModelId("apac.anthropic.claude-3-sonnet-20240229-v1:0")).toBe(
249+
"anthropic.claude-3-sonnet-20240229-v1:0",
250+
)
251+
})
252+
253+
it("should prioritize global. prefix removal over cross-region prefixes", () => {
254+
const handler = createHandler()
255+
const parseBaseModelId = (handler as any).parseBaseModelId.bind(handler)
256+
257+
// Even if there's a model ID that somehow has both (shouldn't happen in practice),
258+
// global. should be removed first
259+
expect(parseBaseModelId("global.us.some-model-id")).toBe("us.some-model-id")
260+
})
261+
262+
it("should return model ID unchanged if no prefix is present", () => {
263+
const handler = createHandler()
264+
const parseBaseModelId = (handler as any).parseBaseModelId.bind(handler)
265+
266+
expect(parseBaseModelId("anthropic.claude-3-sonnet-20240229-v1:0")).toBe(
267+
"anthropic.claude-3-sonnet-20240229-v1:0",
268+
)
269+
expect(parseBaseModelId("amazon.nova-pro-v1:0")).toBe("amazon.nova-pro-v1:0")
270+
})
271+
})
272+
})

0 commit comments

Comments
 (0)