From 482744eb3833ca1de9aeb34e2a3094182e9ef8ca Mon Sep 17 00:00:00 2001
From: Matt Rubens <mrubens@users.noreply.github.com>
Date: Sun, 14 Dec 2025 22:29:40 -0500
Subject: [PATCH] Fix bedrock 1m context handling in prompt router path

---
 src/api/providers/__tests__/bedrock.spec.ts | 80 +++++++++++++++++++++
 src/api/providers/bedrock.ts                |  9 +++
 2 files changed, 89 insertions(+)

diff --git a/src/api/providers/__tests__/bedrock.spec.ts b/src/api/providers/__tests__/bedrock.spec.ts
index e9a0282dce..8a17f893e1 100644
--- a/src/api/providers/__tests__/bedrock.spec.ts
+++ b/src/api/providers/__tests__/bedrock.spec.ts
@@ -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)
+			})
+		})
 	})
 })
diff --git a/src/api/providers/bedrock.ts b/src/api/providers/bedrock.ts
index 4a4adfc0f4..11ef743d15 100644
--- a/src/api/providers/bedrock.ts
+++ b/src/api/providers/bedrock.ts
@@ -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