Add support for bedrock api keys (#6132)

Co-authored-by: Roo Code <[email protected]>

Author: jr

packages/types/src/global-settings.ts

@@ -164,6 +164,7 @@ export const SECRET_STATE_KEYS = [
 	"glamaApiKey",
 	"openRouterApiKey",
 	"awsAccessKey",
+	"awsApiKey",
 	"awsSecretKey",
 	"awsSessionToken",
 	"openAiApiKey",

packages/types/src/provider-settings.ts

@@ -114,6 +114,8 @@ const bedrockSchema = apiModelIdProviderModelSchema.extend({
 	awsUsePromptCache: z.boolean().optional(),
 	awsProfile: z.string().optional(),
 	awsUseProfile: z.boolean().optional(),
+	awsApiKey: z.string().optional(),
+	awsUseApiKey: z.boolean().optional(),
 	awsCustomArn: z.string().optional(),
 	awsModelContextWindow: z.number().optional(),
 	awsBedrockEndpointEnabled: z.boolean().optional(),

packages/types/src/providers/bedrock.ts

@@ -360,7 +360,7 @@ export const BEDROCK_MAX_TOKENS = 4096
 
 export const BEDROCK_DEFAULT_CONTEXT = 128_000
 
-// AWS Bedrock Inference Profile mapping based on official documentation
+// Amazon Bedrock Inference Profile mapping based on official documentation
 // https://docs.aws.amazon.com/bedrock/latest/userguide/inference-profiles-support.html
 // This mapping is pre-ordered by pattern length (descending) to ensure more specific patterns match first
 export const AWS_INFERENCE_PROFILE_MAPPING: Array<[string, string]> = [
@@ -378,7 +378,7 @@ export const AWS_INFERENCE_PROFILE_MAPPING: Array<[string, string]> = [
 	["sa-", "sa."],
 ]
 
-// AWS Bedrock supported regions for the regions dropdown
+// Amazon Bedrock supported regions for the regions dropdown
 // Based on official AWS documentation
 export const BEDROCK_REGIONS = [
 	{ value: "us-east-1", label: "us-east-1" },

pnpm-lock.yaml

@@ -16,7 +16,7 @@ vitest.mock("@aws-sdk/client-bedrock-runtime", () => {
 	}
 })
 
-describe("AWS Bedrock Inference Profiles", () => {
+describe("Amazon Bedrock Inference Profiles", () => {
 	// Helper function to create a handler with specific options
 	const createHandler = (options: Partial<ApiHandlerOptions> = {}) => {
 		const defaultOptions: ApiHandlerOptions = {

src/api/providers/__tests__/bedrock-inference-profiles.spec.ts

@@ -278,5 +278,51 @@ describe("AwsBedrockHandler - Extended Thinking", () => {
 			expect(reasoningChunks[0].text).toBe("Let me think...")
 			expect(reasoningChunks[1].text).toBe(" about this problem.")
 		})
+
+		it("should support API key authentication", async () => {
+			handler = new AwsBedrockHandler({
+				apiProvider: "bedrock",
+				apiModelId: "anthropic.claude-3-5-sonnet-20241022-v2:0",
+				awsRegion: "us-east-1",
+				awsUseApiKey: true,
+				awsApiKey: "test-api-key-token",
+			})
+
+			mockSend.mockResolvedValue({
+				stream: (async function* () {
+					yield { messageStart: { role: "assistant" } }
+					yield {
+						contentBlockStart: {
+							start: { text: "Hello from API key auth" },
+							contentBlockIndex: 0,
+						},
+					}
+					yield { metadata: { usage: { inputTokens: 100, outputTokens: 50 } } }
+				})(),
+			})
+
+			const messages = [{ role: "user" as const, content: "Test message" }]
+			const stream = handler.createMessage("System prompt", messages)
+
+			const chunks = []
+			for await (const chunk of stream) {
+				chunks.push(chunk)
+			}
+
+			// Verify the client was created with API key token
+			expect(BedrockRuntimeClient).toHaveBeenCalledWith(
+				expect.objectContaining({
+					region: "us-east-1",
+					token: { token: "test-api-key-token" },
+					authSchemePreference: ["httpBearerAuth"],
+				}),
+			)
+
+			// Verify the stream worked correctly
+			expect(mockSend).toHaveBeenCalledTimes(1)
+			const textChunks = chunks.filter((c) => c.type === "text")
+			expect(textChunks).toHaveLength(1)
+			expect(textChunks[0].text).toBe("Hello from API key auth")
+		})
 	})
 })

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

@@ -29,7 +29,7 @@ import { BedrockRuntimeClient } from "@aws-sdk/client-bedrock-runtime"
 // Get access to the mocked functions
 const mockBedrockRuntimeClient = vi.mocked(BedrockRuntimeClient)
 
-describe("AWS Bedrock VPC Endpoint Functionality", () => {
+describe("Amazon Bedrock VPC Endpoint Functionality", () => {
 	beforeEach(() => {
 		// Clear all mocks before each test
 		vi.clearAllMocks()

src/api/providers/__tests__/bedrock-vpc-endpoint.spec.ts

@@ -222,7 +222,11 @@ export class AwsBedrockHandler extends BaseProvider implements SingleCompletionH
 				this.options.awsBedrockEndpointEnabled && { endpoint: this.options.awsBedrockEndpoint }),
 		}
 
-		if (this.options.awsUseProfile && this.options.awsProfile) {
+		if (this.options.awsUseApiKey && this.options.awsApiKey) {
+			// Use API key/token-based authentication if enabled and API key is set
+			clientConfig.token = { token: this.options.awsApiKey }
+			clientConfig.authSchemePreference = ["httpBearerAuth"] // Otherwise there's no end of credential problems.
+		} else if (this.options.awsUseProfile && this.options.awsProfile) {
 			// Use profile-based credentials if enabled and profile is set
 			clientConfig.credentials = fromIni({
 				profile: this.options.awsProfile,
@@ -1078,7 +1082,7 @@ Please verify:
 				"throttl",
 				"rate",
 				"limit",
-				"bedrock is unable to process your request", // AWS Bedrock specific throttling message
+				"bedrock is unable to process your request", // Amazon Bedrock specific throttling message
 				"please wait",
 				"quota exceeded",
 				"service unavailable",
@@ -1124,7 +1128,7 @@ Suggestions:
 Please try:
 1. Contact AWS support to request a quota increase
 2. Reduce request frequency temporarily
-3. Check your AWS Bedrock quotas in the AWS console
+3. Check your Amazon Bedrock quotas in the AWS console
 4. Consider using a different model or region with available capacity
 
 `,
@@ -1139,15 +1143,15 @@ Please try:
 
 Please try:
 1. Wait a few minutes and retry
-2. Check the model status in AWS Bedrock console
+2. Check the model status in Amazon Bedrock console
 3. Verify the model is properly provisioned
 
 `,
 			logLevel: "error",
 		},
 		INTERNAL_SERVER_ERROR: {
 			patterns: ["internal server error", "internal error", "server error", "service error"],
-			messageTemplate: `AWS Bedrock internal server error. This is a temporary service issue.
+			messageTemplate: `Amazon Bedrock internal server error. This is a temporary service issue.
 
 Please try:
 1. Retry the request after a brief delay
@@ -1184,7 +1188,7 @@ Please try:
 			],
 			messageTemplate: `Parameter validation error: {errorMessage}
 
-This error indicates that the request parameters don't match AWS Bedrock's expected format.
+This error indicates that the request parameters don't match Amazon Bedrock's expected format.
 
 Common causes:
 1. Extended thinking parameter format is incorrect
@@ -1193,7 +1197,7 @@ Common causes:
 
 Please check:
 - Model supports the requested features (extended thinking, etc.)
-- Parameter format matches AWS Bedrock specification
+- Parameter format matches Amazon Bedrock specification
 - Model ID is correct for the requested features`,
 			logLevel: "error",
 		},
@@ -1218,7 +1222,7 @@ Please check:
 			return "THROTTLING"
 		}
 
-		// Check for AWS Bedrock specific throttling exception names
+		// Check for Amazon Bedrock specific throttling exception names
 		if ((error as any).name === "ThrottlingException" || (error as any).__type === "ThrottlingException") {
 			return "THROTTLING"
 		}

src/api/providers/bedrock.ts

@@ -413,8 +413,8 @@
 		"@anthropic-ai/bedrock-sdk": "^0.10.2",
 		"@anthropic-ai/sdk": "^0.37.0",
 		"@anthropic-ai/vertex-sdk": "^0.7.0",
-		"@aws-sdk/client-bedrock-runtime": "^3.779.0",
-		"@aws-sdk/credential-providers": "^3.806.0",
+		"@aws-sdk/client-bedrock-runtime": "^3.848.0",
+		"@aws-sdk/credential-providers": "^3.848.0",
 		"@google/genai": "^1.0.0",
 		"@lmstudio/sdk": "^1.1.1",
 		"@mistralai/mistralai": "^1.3.6",

src/package.json

@@ -1,6 +1,6 @@
 import { useCallback, useState, useEffect } from "react"
 import { Checkbox } from "vscrui"
-import { VSCodeTextField, VSCodeRadio, VSCodeRadioGroup } from "@vscode/webview-ui-toolkit/react"
+import { VSCodeTextField } from "@vscode/webview-ui-toolkit/react"
 
 import { type ProviderSettings, type ModelInfo, BEDROCK_REGIONS } from "@roo-code/types"
 
@@ -37,19 +37,51 @@ export const Bedrock = ({ apiConfiguration, setApiConfigurationField, selectedMo
 
 	return (
 		<>
-			<VSCodeRadioGroup
-				value={apiConfiguration?.awsUseProfile ? "profile" : "credentials"}
-				onChange={handleInputChange(
-					"awsUseProfile",
-					(e) => (e.target as HTMLInputElement).value === "profile",
-				)}>
-				<VSCodeRadio value="credentials">{t("settings:providers.awsCredentials")}</VSCodeRadio>
-				<VSCodeRadio value="profile">{t("settings:providers.awsProfile")}</VSCodeRadio>
-			</VSCodeRadioGroup>
+			<div>
+				<label className="block font-medium mb-1">Authentication Method</label>
+				<Select
+					value={
+						apiConfiguration?.awsUseApiKey
+							? "apikey"
+							: apiConfiguration?.awsUseProfile
+								? "profile"
+								: "credentials"
+					}
+					onValueChange={(value) => {
+						if (value === "apikey") {
+							setApiConfigurationField("awsUseApiKey", true)
+							setApiConfigurationField("awsUseProfile", false)
+						} else if (value === "profile") {
+							setApiConfigurationField("awsUseApiKey", false)
+							setApiConfigurationField("awsUseProfile", true)
+						} else {
+							setApiConfigurationField("awsUseApiKey", false)
+							setApiConfigurationField("awsUseProfile", false)
+						}
+					}}>
+					<SelectTrigger className="w-full">
+						<SelectValue placeholder={t("settings:common.select")} />
+					</SelectTrigger>
+					<SelectContent>
+						<SelectItem value="credentials">{t("settings:providers.awsCredentials")}</SelectItem>
+						<SelectItem value="profile">{t("settings:providers.awsProfile")}</SelectItem>
+						<SelectItem value="apikey">{t("settings:providers.awsApiKey")}</SelectItem>
+					</SelectContent>
+				</Select>
+			</div>
 			<div className="text-sm text-vscode-descriptionForeground -mt-3">
 				{t("settings:providers.apiKeyStorageNotice")}
 			</div>
-			{apiConfiguration?.awsUseProfile ? (
+			{apiConfiguration?.awsUseApiKey ? (
+				<VSCodeTextField
+					value={apiConfiguration?.awsApiKey || ""}
+					type="password"
+					onInput={handleInputChange("awsApiKey")}
+					placeholder={t("settings:placeholders.apiKey")}
+					className="w-full">
+					<label className="block font-medium mb-1">{t("settings:providers.awsApiKey")}</label>
+				</VSCodeTextField>
+			) : apiConfiguration?.awsUseProfile ? (
 				<VSCodeTextField
 					value={apiConfiguration?.awsProfile || ""}
 					onInput={handleInputChange("awsProfile")}