scriptoriumadmin/Roo-Code
1
0
Fork 0
mirror of https://github.com/RooVetGit/Roo-Code.git synced 2026-08-28 05:27:24 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Fix bedrock 1m context handling in prompt router path

Browse source
This commit is contained in:
Matt Rubens 2025-12-14 22:29:40 -05:00
parent ae65c6d730
commit 482744eb38
2 changed files with 89 additions and 0 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

80
src/api/providers/__tests__/bedrock.spec.ts
View file

@ -828,5 +828,85 @@ describe("AwsBedrockHandler", () => {
expect(model1.id).toBe(model2.id)
expect(model2.id).toBe(model3.id)
})
describe("getModelById 1M context support (prompt router path)", () => {
it("should apply 1M context window when getModelById is called with 1M context enabled", () => {
const handler = new AwsBedrockHandler({
apiModelId: "some-prompt-router",
awsAccessKey: "test",
awsSecretKey: "test",
awsRegion: "us-east-1",
awsBedrock1MContext: true,
})
// Access getModelById using type casting (it's a public method but used internally)
const modelById = handler.getModelById(BEDROCK_1M_CONTEXT_MODEL_IDS[0])
// Should have 1M context window when 1M context is enabled
expect(modelById.info.contextWindow).toBe(1_000_000)
})
it("should not apply 1M context window when getModelById is called with 1M context disabled", () => {
const handler = new AwsBedrockHandler({
apiModelId: "some-prompt-router",
awsAccessKey: "test",
awsSecretKey: "test",
awsRegion: "us-east-1",
awsBedrock1MContext: false,
})
const modelById = handler.getModelById(BEDROCK_1M_CONTEXT_MODEL_IDS[0])
// Should use default context window (200k)
expect(modelById.info.contextWindow).toBe(200_000)
})
it("should not apply 1M context window for non-Claude Sonnet 4 models via getModelById", () => {
const handler = new AwsBedrockHandler({
apiModelId: "some-prompt-router",
awsAccessKey: "test",
awsSecretKey: "test",
awsRegion: "us-east-1",
awsBedrock1MContext: true,
})
const modelById = handler.getModelById("anthropic.claude-3-5-sonnet-20241022-v2:0")
// Should use default context window for non-Sonnet 4 models
expect(modelById.info.contextWindow).toBe(200_000)
})
it("should apply 1M context window with cross-region prefix in getModelById", () => {
const handler = new AwsBedrockHandler({
apiModelId: "some-prompt-router",
awsAccessKey: "test",
awsSecretKey: "test",
awsRegion: "us-east-1",
awsBedrock1MContext: true,
})
// Test with cross-region prefixed model ID (as would come from prompt router)
const modelById = handler.getModelById(`us.${BEDROCK_1M_CONTEXT_MODEL_IDS[0]}`)
// Should still apply 1M context window (parseBaseModelId strips the prefix)
expect(modelById.info.contextWindow).toBe(1_000_000)
})
it("should allow user override of context window even with 1M context enabled in getModelById", () => {
const handler = new AwsBedrockHandler({
apiModelId: "some-prompt-router",
awsAccessKey: "test",
awsSecretKey: "test",
awsRegion: "us-east-1",
awsBedrock1MContext: true,
awsModelContextWindow: 500_000, // User override
})
const modelById = handler.getModelById(BEDROCK_1M_CONTEXT_MODEL_IDS[0])
// User override should take precedence
expect(modelById.info.contextWindow).toBe(500_000)
})
})
})
})

9
src/api/providers/bedrock.ts
View file

@ -1007,6 +1007,15 @@ export class AwsBedrockHandler extends BaseProvider implements SingleCompletionH
}
}
// Check if 1M context is enabled for Claude Sonnet 4 / 4.5
// This must be applied here as well for prompt router responses that call getModelById
if (BEDROCK_1M_CONTEXT_MODEL_IDS.includes(baseModelId as any) && this.options.awsBedrock1MContext) {
model.info = {
...model.info,
contextWindow: 1_000_000,
}
}
// Always allow user to override detected/guessed maxTokens and contextWindow
if (this.options.modelMaxTokens && this.options.modelMaxTokens > 0) {
model.info.maxTokens = this.options.modelMaxTokens