fix(bedrock): address PR feedback

Author: allquixotic

packages/types/src/__tests__/bedrock.spec.ts

@@ -8,7 +8,9 @@ import {
 	BEDROCK_NATIVE_1M_CONTEXT_MODEL_IDS,
 	bedrockModels,
 	expandBedrockTargetsWith1MVariants,
+	guessBedrockModelInfoFromId,
 	hasBedrock1MContextIndicator,
+	resolveBedrockMaxOutputTokensOverride,
 	resolveBedrockModelInfo,
 	stripBedrock1MContextSuffix,
 } from "../providers/bedrock.js"
@@ -65,6 +67,11 @@ describe("Bedrock model catalog", () => {
 		expect((bedrockModels["anthropic.claude-haiku-4-5-20251001-v1:0"] as ModelInfo).promptCacheTtl).toBe("1h")
 		expect((bedrockModels["anthropic.claude-opus-4-5-20251101-v1:0"] as ModelInfo).promptCacheTtl).toBe("1h")
 	})
+
+	it("prefers specific guessed Opus patterns before generic Claude 4 patterns", () => {
+		const guessed = guessBedrockModelInfoFromId("arn:aws:bedrock:us-west-2::foundation-model/claude-4-opus-custom")
+		expect(guessed.maxTokens).toBe(4096)
+	})
 })
 
 describe("resolveBedrockModelInfo", () => {
@@ -94,6 +101,24 @@ describe("resolveBedrockModelInfo", () => {
 		})
 		expect(info.maxTokens).toBe(32_000)
 	})
+
+	it("applies max-output overrides only to the target that was probed", () => {
+		expect(
+			resolveBedrockMaxOutputTokensOverride({
+				currentTargetId: "global.anthropic.claude-opus-4-7",
+				overrideTargetId: "global.anthropic.claude-opus-4-7",
+				maxOutputTokensOverride: 128_000,
+			}),
+		).toBe(128_000)
+
+		expect(
+			resolveBedrockMaxOutputTokensOverride({
+				currentTargetId: "anthropic.claude-haiku-4-5-20251001-v1:0",
+				overrideTargetId: "global.anthropic.claude-opus-4-7",
+				maxOutputTokensOverride: 128_000,
+			}),
+		).toBeUndefined()
+	})
 })
 
 describe("expandBedrockTargetsWith1MVariants", () => {

packages/types/src/provider-settings.ts

@@ -236,8 +236,10 @@ const bedrockSchema = apiModelIdProviderModelSchema.extend({
 		.optional(),
 	awsModelContextWindow: z.number().optional(),
 	// Empirically detected (or manually entered) per-config cap on the model's max output tokens.
-	// Takes precedence over the static `bedrockModels.<id>.maxTokens` table when present.
+	// Takes precedence over the static `bedrockModels.<id>.maxTokens` table when present
+	// and scoped to the target id recorded in `awsModelMaxOutputTokensTargetId`.
 	awsModelMaxOutputTokens: z.number().optional(),
+	awsModelMaxOutputTokensTargetId: z.string().optional(),
 	awsBedrockEndpointEnabled: z.boolean().optional(),
 	awsBedrockEndpoint: z.string().optional(),
 	awsBedrock1MContext: z.boolean().optional(), // Enable 'context-1m-2025-08-07' beta for 1M context window.

packages/types/src/providers/bedrock.ts

@@ -899,38 +899,38 @@ export const shouldUseBedrock1MContext = ({
 
 export const guessBedrockModelInfoFromId = (modelId: string): Partial<ModelInfo> => {
 	const modelConfigMap: Record<string, Partial<ModelInfo>> = {
-		"claude-4": {
-			maxTokens: 8192,
+		"claude-4-opus": {
+			maxTokens: 4096,
 			contextWindow: 200_000,
 			supportsImages: true,
 			supportsPromptCache: true,
 		},
-		"claude-3-7": {
-			maxTokens: 8192,
+		"claude-3-opus": {
+			maxTokens: 4096,
 			contextWindow: 200_000,
 			supportsImages: true,
 			supportsPromptCache: true,
 		},
-		"claude-3-5": {
-			maxTokens: 8192,
+		"claude-3-haiku": {
+			maxTokens: 4096,
 			contextWindow: 200_000,
 			supportsImages: true,
 			supportsPromptCache: true,
 		},
-		"claude-4-opus": {
-			maxTokens: 4096,
+		"claude-4": {
+			maxTokens: 8192,
 			contextWindow: 200_000,
 			supportsImages: true,
 			supportsPromptCache: true,
 		},
-		"claude-3-opus": {
-			maxTokens: 4096,
+		"claude-3-7": {
+			maxTokens: 8192,
 			contextWindow: 200_000,
 			supportsImages: true,
 			supportsPromptCache: true,
 		},
-		"claude-3-haiku": {
-			maxTokens: 4096,
+		"claude-3-5": {
+			maxTokens: 8192,
 			contextWindow: 200_000,
 			supportsImages: true,
 			supportsPromptCache: true,
@@ -952,6 +952,24 @@ export const guessBedrockModelInfoFromId = (modelId: string): Partial<ModelInfo>
 	}
 }
 
+export const resolveBedrockMaxOutputTokensOverride = ({
+	currentTargetId,
+	overrideTargetId,
+	maxOutputTokensOverride,
+}: {
+	currentTargetId?: string
+	overrideTargetId?: string
+	maxOutputTokensOverride?: number
+}): number | undefined => {
+	if (!currentTargetId || !overrideTargetId || !maxOutputTokensOverride || maxOutputTokensOverride <= 0) {
+		return undefined
+	}
+
+	return stripBedrock1MContextSuffix(currentTargetId) === stripBedrock1MContextSuffix(overrideTargetId)
+		? maxOutputTokensOverride
+		: undefined
+}
+
 export const resolveBedrockModelInfo = ({
 	baseModelId,
 	targetId,

src/api/providers/__tests__/bedrock-discovery.spec.ts

@@ -0,0 +1,61 @@
+import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"
+
+const mocks = vi.hoisted(() => ({
+	send: vi.fn(),
+	sendOptions: [] as Array<{ abortSignal?: AbortSignal } | undefined>,
+}))
+
+vi.mock("@aws-sdk/client-bedrock", () => {
+	class BedrockClient {
+		send(command: unknown, options?: { abortSignal?: AbortSignal }) {
+			mocks.sendOptions.push(options)
+			return mocks.send(command, options)
+		}
+	}
+
+	class ListFoundationModelsCommand {
+		constructor(public readonly input: unknown) {}
+	}
+
+	class ListInferenceProfilesCommand {
+		constructor(public readonly input: unknown) {}
+	}
+
+	return {
+		BedrockClient,
+		ListFoundationModelsCommand,
+		ListInferenceProfilesCommand,
+	}
+})
+
+import { BEDROCK_DISCOVERY_TIMEOUT_MS, discoverBedrockTargets } from "../bedrock-discovery"
+
+describe("discoverBedrockTargets", () => {
+	beforeEach(() => {
+		mocks.send.mockReset()
+		mocks.sendOptions.length = 0
+	})
+
+	afterEach(() => {
+		vi.useRealTimers()
+	})
+
+	it("times out and aborts slow AWS discovery calls", async () => {
+		vi.useFakeTimers()
+		mocks.send.mockImplementation(() => new Promise(() => {}))
+
+		const promise = expect(
+			discoverBedrockTargets({
+				awsRegion: "us-west-2",
+				awsAccessKey: "AKIA",
+				awsSecretKey: "secret",
+			}),
+		).rejects.toThrow(/Bedrock discovery timed out/)
+
+		await vi.advanceTimersByTimeAsync(BEDROCK_DISCOVERY_TIMEOUT_MS)
+
+		await promise
+		expect(mocks.sendOptions).toHaveLength(2)
+		expect(mocks.sendOptions.every((options) => options?.abortSignal?.aborted)).toBe(true)
+	})
+})

src/api/providers/bedrock-discovery.ts

@@ -21,6 +21,8 @@ import {
 
 import { Package } from "../../shared/package"
 
+export const BEDROCK_DISCOVERY_TIMEOUT_MS = 15_000
+
 // The two AWS SDK packages we use here (`@aws-sdk/client-bedrock` for control-plane discovery
 // and `@aws-sdk/client-bedrock-runtime` for Converse probes) each ship their own discriminated
 // `Config` interface where `token`/`credentials` are tagged with package-private branded types.
@@ -152,7 +154,31 @@ const buildInferenceProfileTarget = (summary: InferenceProfileSummary): BedrockD
 	}
 }
 
-const listInferenceProfiles = async (client: BedrockClient) => {
+const withTimeout = async <T>(
+	label: string,
+	timeoutMs: number,
+	operation: (abortSignal: AbortSignal) => Promise<T>,
+): Promise<T> => {
+	const abortController = new AbortController()
+	let timeoutId: ReturnType<typeof setTimeout> | undefined
+
+	const timeoutPromise = new Promise<never>((_, reject) => {
+		timeoutId = setTimeout(() => {
+			abortController.abort()
+			reject(new Error(`${label} timed out after ${timeoutMs}ms`))
+		}, timeoutMs)
+	})
+
+	try {
+		return await Promise.race([operation(abortController.signal), timeoutPromise])
+	} finally {
+		if (timeoutId) {
+			clearTimeout(timeoutId)
+		}
+	}
+}
+
+const listInferenceProfiles = async (client: BedrockClient, abortSignal?: AbortSignal) => {
 	const results: InferenceProfileSummary[] = []
 	let nextToken: string | undefined
 
@@ -162,6 +188,7 @@ const listInferenceProfiles = async (client: BedrockClient) => {
 				nextToken,
 				maxResults: 100,
 			}),
+			{ abortSignal },
 		)
 
 		results.push(...(response.inferenceProfileSummaries ?? []))
@@ -178,10 +205,15 @@ export const discoverBedrockTargets = async (options: ProviderSettings): Promise
 
 	const client = new BedrockClient(toBedrockClientConfig(options))
 
-	const [foundationModelsResponse, inferenceProfiles] = await Promise.all([
-		client.send(new ListFoundationModelsCommand({})),
-		listInferenceProfiles(client),
-	])
+	const [foundationModelsResponse, inferenceProfiles] = await withTimeout(
+		"Bedrock discovery",
+		BEDROCK_DISCOVERY_TIMEOUT_MS,
+		(abortSignal) =>
+			Promise.all([
+				client.send(new ListFoundationModelsCommand({}), { abortSignal }),
+				listInferenceProfiles(client, abortSignal),
+			]),
+	)
 
 	const targets = [
 		...(foundationModelsResponse.modelSummaries ?? [])

src/api/providers/bedrock.ts

@@ -34,6 +34,8 @@ import {
 	ApiProviderError,
 	inferBedrockInvokeTargetKind,
 	parseBedrockBaseModelId,
+	resolveBedrockInvokeTargetId,
+	resolveBedrockMaxOutputTokensOverride,
 	resolveBedrockModelInfo,
 	shouldUseBedrock1MContext,
 	stripBedrock1MContextSuffix,
@@ -379,7 +381,9 @@ export class AwsBedrockHandler extends BaseProvider implements SingleCompletionH
 		// parseBaseModelId strips cross-region inference prefixes (e.g. `us.`, `eu.`) and the
 		// synthetic `:1m` dropdown suffix.
 		const baseModelId = this.parseBaseModelId(modelConfig.id)
-		const requiresAdaptiveThinking = BEDROCK_ADAPTIVE_THINKING_MODEL_IDS.includes(baseModelId as any)
+		const requiresAdaptiveThinking = BEDROCK_ADAPTIVE_THINKING_MODEL_IDS.includes(
+			baseModelId as (typeof BEDROCK_ADAPTIVE_THINKING_MODEL_IDS)[number],
+		)
 
 		// Determine if thinking should be enabled
 		// metadata?.thinking?.enabled: Explicitly enabled through API metadata (direct request)
@@ -443,7 +447,7 @@ export class AwsBedrockHandler extends BaseProvider implements SingleCompletionH
 		const configuredTargetForIndicator =
 			this.options.awsBedrockInvokeTarget || this.options.awsCustomArn || modelConfig.id
 		const is1MContextEnabled =
-			BEDROCK_1M_CONTEXT_MODEL_IDS.includes(baseModelId as any) &&
+			BEDROCK_1M_CONTEXT_MODEL_IDS.includes(baseModelId as (typeof BEDROCK_1M_CONTEXT_MODEL_IDS)[number]) &&
 			shouldUseBedrock1MContext({
 				targetId: configuredTargetForIndicator,
 				baseModelId,
@@ -456,11 +460,14 @@ export class AwsBedrockHandler extends BaseProvider implements SingleCompletionH
 		// the fine-grained-tool-streaming beta, so we must omit anthropic_beta entirely
 		// for those models. Older Claudes silently accept (and effectively ignore) them,
 		// so we keep the current behavior for them.
-		const skipAnthropicBetaFlags = BEDROCK_NATIVE_1M_CONTEXT_MODEL_IDS.includes(baseModelId as any)
+		const skipAnthropicBetaFlags = BEDROCK_NATIVE_1M_CONTEXT_MODEL_IDS.includes(
+			baseModelId as (typeof BEDROCK_NATIVE_1M_CONTEXT_MODEL_IDS)[number],
+		)
 
 		// Determine if service tier should be applied (checked later when building payload)
 		const useServiceTier =
-			this.options.awsBedrockServiceTier && BEDROCK_SERVICE_TIER_MODEL_IDS.includes(baseModelId as any)
+			this.options.awsBedrockServiceTier &&
+			BEDROCK_SERVICE_TIER_MODEL_IDS.includes(baseModelId as (typeof BEDROCK_SERVICE_TIER_MODEL_IDS)[number])
 		if (useServiceTier) {
 			logger.info("Service tier specified for Bedrock request", {
 				ctx: "bedrock",
@@ -1133,6 +1140,14 @@ export class AwsBedrockHandler extends BaseProvider implements SingleCompletionH
 		return parseBedrockBaseModelId(modelId)
 	}
 
+	private getScopedMaxOutputTokensOverride(): number | undefined {
+		return resolveBedrockMaxOutputTokensOverride({
+			currentTargetId: resolveBedrockInvokeTargetId(this.options),
+			overrideTargetId: this.options.awsModelMaxOutputTokensTargetId,
+			maxOutputTokensOverride: this.options.awsModelMaxOutputTokens,
+		})
+	}
+
 	//Prompt Router responses come back in a different sequence and the model used is in the response and must be fetched by name
 	getModelById(modelId: string, modelType?: string): { id: BedrockModelId | string; info: ModelInfo } {
 		let model
@@ -1143,8 +1158,8 @@ export class AwsBedrockHandler extends BaseProvider implements SingleCompletionH
 			modelMaxTokens: this.options.modelMaxTokens,
 			contextWindowOverride: this.options.awsModelContextWindow,
 			// Empirically detected per-config max output tokens (from the "Detect" probe in the
-			// settings UI) widens the static cap so downstream request builders pick it up too.
-			maxOutputTokensOverride: this.options.awsModelMaxOutputTokens,
+			// settings UI) widens the static cap only for the exact AWS target that was probed.
+			maxOutputTokensOverride: this.getScopedMaxOutputTokensOverride(),
 		})
 
 		if (resolved.baseModelId in bedrockModels) {
@@ -1227,7 +1242,9 @@ export class AwsBedrockHandler extends BaseProvider implements SingleCompletionH
 				const baseIdForGlobal = this.parseBaseModelId(modelConfig.id)
 				if (
 					this.options.awsUseGlobalInference &&
-					BEDROCK_GLOBAL_INFERENCE_MODEL_IDS.includes(baseIdForGlobal as any)
+					BEDROCK_GLOBAL_INFERENCE_MODEL_IDS.includes(
+						baseIdForGlobal as (typeof BEDROCK_GLOBAL_INFERENCE_MODEL_IDS)[number],
+					)
 				) {
 					modelConfig.id = `global.${baseIdForGlobal}`
 				}
@@ -1244,7 +1261,10 @@ export class AwsBedrockHandler extends BaseProvider implements SingleCompletionH
 		// Check if 1M context is enabled for supported Claude 4 models
 		// Use parseBaseModelId to handle cross-region inference prefixes
 		const baseModelId = this.parseBaseModelId(modelConfig.id)
-		if (BEDROCK_1M_CONTEXT_MODEL_IDS.includes(baseModelId as any) && this.options.awsBedrock1MContext) {
+		if (
+			BEDROCK_1M_CONTEXT_MODEL_IDS.includes(baseModelId as (typeof BEDROCK_1M_CONTEXT_MODEL_IDS)[number]) &&
+			this.options.awsBedrock1MContext
+		) {
 			// Update context window and pricing to 1M tier when 1M context beta is enabled
 			const tier = modelConfig.info.tiers?.[0]
 			modelConfig.info = {
@@ -1268,7 +1288,12 @@ export class AwsBedrockHandler extends BaseProvider implements SingleCompletionH
 
 		// Apply service tier pricing if specified and model supports it
 		const baseModelIdForTier = this.parseBaseModelId(modelConfig.id)
-		if (this.options.awsBedrockServiceTier && BEDROCK_SERVICE_TIER_MODEL_IDS.includes(baseModelIdForTier as any)) {
+		if (
+			this.options.awsBedrockServiceTier &&
+			BEDROCK_SERVICE_TIER_MODEL_IDS.includes(
+				baseModelIdForTier as (typeof BEDROCK_SERVICE_TIER_MODEL_IDS)[number],
+			)
+		) {
 			const pricingMultiplier = BEDROCK_SERVICE_TIER_PRICING[this.options.awsBedrockServiceTier]
 			if (pricingMultiplier && pricingMultiplier !== 1.0) {
 				// Apply pricing multiplier to all price fields