Improves model info detection for custom Bedrock ARNs

Adds heuristics to better estimate model capabilities when using unknown or custom model ARNs, including context window and max tokens. Allows user overrides for key model parameters via provider settings, improving flexibility and reliability for non-standard model integrations.

Fixes #3712

Author: adamhill

src/api/providers/bedrock.ts

@@ -29,6 +29,7 @@ import { convertToBedrockConverseMessages as sharedConverter } from "../transfor
 
 const BEDROCK_DEFAULT_TEMPERATURE = 0.3
 const BEDROCK_MAX_TOKENS = 4096
+const BEDROCK_DEFAULT_CONTEXT = 128_000 // PATCH: used for unknown custom ARNs
 
 /************************************************************************************
  *
@@ -186,6 +187,54 @@ export class AwsBedrockHandler extends BaseProvider implements SingleCompletionH
 		this.client = new BedrockRuntimeClient(clientConfig)
 	}
 
+	// PATCH: Helper to guess model info from custom modelId string if not in bedrockModels
+	private guessModelInfoFromId(modelId: string): Partial<SharedModelInfo> {
+		// Try to match Claude 3.7, 3.5, Opus, Haiku, etc.
+		const id = modelId.toLowerCase()
+
+		if (id.includes("claude-3-7")) {
+			return {
+				maxTokens: 8192,
+				contextWindow: 200_000,
+				supportsImages: true,
+				supportsPromptCache: true,
+			}
+		}
+		if (id.includes("claude-3-5")) {
+			return {
+				maxTokens: 8192,
+				contextWindow: 200_000,
+				supportsImages: true,
+				supportsPromptCache: true,
+			}
+		}
+		if (id.includes("claude-3-opus")) {
+			return {
+				maxTokens: 4096,
+				contextWindow: 200_000,
+				supportsImages: true,
+				supportsPromptCache: true,
+			}
+		}
+		if (id.includes("claude-3-haiku")) {
+			return {
+				maxTokens: 4096,
+				contextWindow: 200_000,
+				supportsImages: true,
+				supportsPromptCache: true,
+			}
+		}
+		// PATCH: Add more heuristics as needed here
+
+		// Default fallback
+		return {
+			maxTokens: BEDROCK_MAX_TOKENS,
+			contextWindow: BEDROCK_DEFAULT_CONTEXT,
+			supportsImages: false,
+			supportsPromptCache: false,
+		}
+	}
+
 	override async *createMessage(systemPrompt: string, messages: Anthropic.Messages.MessageParam[]): ApiStream {
 		let modelConfig = this.getModel()
 		// Handle cross-region inference
@@ -629,16 +678,24 @@ export class AwsBedrockHandler extends BaseProvider implements SingleCompletionH
 				info: JSON.parse(JSON.stringify(bedrockModels[bedrockDefaultPromptRouterModelId])),
 			}
 		} else {
+			// PATCH: Use heuristics for model info, then allow overrides from ProviderSettings
+			const guessed = this.guessModelInfoFromId(modelId)
 			model = {
 				id: bedrockDefaultModelId,
-				info: JSON.parse(JSON.stringify(bedrockModels[bedrockDefaultModelId])),
+				info: {
+					...JSON.parse(JSON.stringify(bedrockModels[bedrockDefaultModelId])),
+					...guessed,
+				},
 			}
 		}
 
-		// If modelMaxTokens is explicitly set in options, override the default
+		// PATCH: Always allow user to override detected/guessed maxTokens and contextWindow
 		if (this.options.modelMaxTokens && this.options.modelMaxTokens > 0) {
 			model.info.maxTokens = this.options.modelMaxTokens
 		}
+		if (this.options.awsModelContextWindow && this.options.awsModelContextWindow > 0) {
+			model.info.contextWindow = this.options.awsModelContextWindow
+		}
 
 		return model
 	}
@@ -673,8 +730,7 @@ export class AwsBedrockHandler extends BaseProvider implements SingleCompletionH
 			}
 		}
 
-		modelConfig.info.maxTokens = modelConfig.info.maxTokens || BEDROCK_MAX_TOKENS
-
+		// PATCH: Don't override maxTokens/contextWindow here; handled in getModelById (and includes user overrides)
 		return modelConfig as { id: BedrockModelId | string; info: SharedModelInfo }
 	}
 

src/exports/roo-code.d.ts

@@ -253,6 +253,7 @@ type ProviderSettings = {
 	awsProfile?: string | undefined
 	awsUseProfile?: boolean | undefined
 	awsCustomArn?: string | undefined
+	awsModelContextWindow?: number | undefined
 	vertexKeyFile?: string | undefined
 	vertexJsonCredentials?: string | undefined
 	vertexProjectId?: string | undefined
@@ -656,6 +657,7 @@ type IpcMessage =
 								awsProfile?: string | undefined
 								awsUseProfile?: boolean | undefined
 								awsCustomArn?: string | undefined
+								awsModelContextWindow?: number | undefined
 								vertexKeyFile?: string | undefined
 								vertexJsonCredentials?: string | undefined
 								vertexProjectId?: string | undefined
@@ -1131,6 +1133,7 @@ type TaskCommand =
 					awsProfile?: string | undefined
 					awsUseProfile?: boolean | undefined
 					awsCustomArn?: string | undefined
+					awsModelContextWindow?: number | undefined
 					vertexKeyFile?: string | undefined
 					vertexJsonCredentials?: string | undefined
 					vertexProjectId?: string | undefined

src/exports/types.ts

@@ -257,6 +257,7 @@ type ProviderSettings = {
 	awsProfile?: string | undefined
 	awsUseProfile?: boolean | undefined
 	awsCustomArn?: string | undefined
+	awsModelContextWindow?: number | undefined
 	vertexKeyFile?: string | undefined
 	vertexJsonCredentials?: string | undefined
 	vertexProjectId?: string | undefined
@@ -670,6 +671,7 @@ type IpcMessage =
 								awsProfile?: string | undefined
 								awsUseProfile?: boolean | undefined
 								awsCustomArn?: string | undefined
+								awsModelContextWindow?: number | undefined
 								vertexKeyFile?: string | undefined
 								vertexJsonCredentials?: string | undefined
 								vertexProjectId?: string | undefined
@@ -1147,6 +1149,7 @@ type TaskCommand =
 					awsProfile?: string | undefined
 					awsUseProfile?: boolean | undefined
 					awsCustomArn?: string | undefined
+					awsModelContextWindow?: number | undefined
 					vertexKeyFile?: string | undefined
 					vertexJsonCredentials?: string | undefined
 					vertexProjectId?: string | undefined

src/schemas/index.ts

@@ -465,6 +465,7 @@ const bedrockSchema = apiModelIdProviderModelSchema.extend({
 	awsProfile: z.string().optional(),
 	awsUseProfile: z.boolean().optional(),
 	awsCustomArn: z.string().optional(),
+	awsModelContextWindow: z.number().optional(),
 })
 
 const vertexSchema = apiModelIdProviderModelSchema.extend({
@@ -651,6 +652,7 @@ const providerSettingsRecord: ProviderSettingsRecord = {
 	awsProfile: undefined,
 	awsUseProfile: undefined,
 	awsCustomArn: undefined,
+	awsModelContextWindow: undefined,
 	// Google Vertex
 	vertexKeyFile: undefined,
 	vertexJsonCredentials: undefined,