Browse Source

Add support for aws credentials file or environment variables, and session token

Saoud Rizwan 1 year ago
parent
commit
0badfa2706

+ 3 - 2
src/api/bedrock.ts

@@ -13,8 +13,9 @@ export class AwsBedrockHandler implements ApiHandler {
 		this.client = new AnthropicBedrock({
 			// Authenticate by either providing the keys below or use the default AWS credential providers, such as
 			// using ~/.aws/credentials or the "AWS_SECRET_ACCESS_KEY" and "AWS_ACCESS_KEY_ID" environment variables.
-			awsAccessKey: this.options.awsAccessKey,
-			awsSecretKey: this.options.awsSecretKey,
+			...(this.options.awsAccessKey ? { awsAccessKey: this.options.awsAccessKey } : {}),
+			...(this.options.awsSecretKey ? { awsSecretKey: this.options.awsSecretKey } : {}),
+			...(this.options.awsSessionToken ? { awsSessionToken: this.options.awsSessionToken } : {}),
 
 			// awsRegion changes the aws region to which the request is made. By default, we read AWS_REGION,
 			// and if that's not present, we default to us-east-1. Note that we do not read ~/.aws/config for the region.

+ 13 - 2
src/providers/ClaudeDevProvider.ts

@@ -16,7 +16,7 @@ https://github.com/microsoft/vscode-webview-ui-toolkit-samples/blob/main/default
 https://github.com/KumarVariable/vscode-extension-sidebar-html/blob/master/src/customSidebarViewProvider.ts
 */
 
-type SecretKey = "apiKey" | "openRouterApiKey" | "awsAccessKey" | "awsSecretKey"
+type SecretKey = "apiKey" | "openRouterApiKey" | "awsAccessKey" | "awsSecretKey" | "awsSessionToken"
 type GlobalStateKey =
 	| "apiProvider"
 	| "apiModelId"
@@ -310,6 +310,7 @@ export class ClaudeDevProvider implements vscode.WebviewViewProvider {
 								openRouterApiKey,
 								awsAccessKey,
 								awsSecretKey,
+								awsSessionToken,
 								awsRegion,
 								vertexProjectId,
 								vertexRegion,
@@ -320,6 +321,7 @@ export class ClaudeDevProvider implements vscode.WebviewViewProvider {
 							await this.storeSecret("openRouterApiKey", openRouterApiKey)
 							await this.storeSecret("awsAccessKey", awsAccessKey)
 							await this.storeSecret("awsSecretKey", awsSecretKey)
+							await this.storeSecret("awsSessionToken", awsSessionToken)
 							await this.updateGlobalState("awsRegion", awsRegion)
 							await this.updateGlobalState("vertexProjectId", vertexProjectId)
 							await this.updateGlobalState("vertexRegion", vertexRegion)
@@ -609,6 +611,7 @@ export class ClaudeDevProvider implements vscode.WebviewViewProvider {
 			openRouterApiKey,
 			awsAccessKey,
 			awsSecretKey,
+			awsSessionToken,
 			awsRegion,
 			vertexProjectId,
 			vertexRegion,
@@ -623,6 +626,7 @@ export class ClaudeDevProvider implements vscode.WebviewViewProvider {
 			this.getSecret("openRouterApiKey") as Promise<string | undefined>,
 			this.getSecret("awsAccessKey") as Promise<string | undefined>,
 			this.getSecret("awsSecretKey") as Promise<string | undefined>,
+			this.getSecret("awsSessionToken") as Promise<string | undefined>,
 			this.getGlobalState("awsRegion") as Promise<string | undefined>,
 			this.getGlobalState("vertexProjectId") as Promise<string | undefined>,
 			this.getGlobalState("vertexRegion") as Promise<string | undefined>,
@@ -654,6 +658,7 @@ export class ClaudeDevProvider implements vscode.WebviewViewProvider {
 				openRouterApiKey,
 				awsAccessKey,
 				awsSecretKey,
+				awsSessionToken,
 				awsRegion,
 				vertexProjectId,
 				vertexRegion,
@@ -728,7 +733,13 @@ export class ClaudeDevProvider implements vscode.WebviewViewProvider {
 		for (const key of this.context.globalState.keys()) {
 			await this.context.globalState.update(key, undefined)
 		}
-		const secretKeys: SecretKey[] = ["apiKey", "openRouterApiKey", "awsAccessKey", "awsSecretKey"]
+		const secretKeys: SecretKey[] = [
+			"apiKey",
+			"openRouterApiKey",
+			"awsAccessKey",
+			"awsSecretKey",
+			"awsSessionToken",
+		]
 		for (const key of secretKeys) {
 			await this.storeSecret(key, undefined)
 		}

+ 1 - 0
src/shared/api.ts

@@ -6,6 +6,7 @@ export interface ApiHandlerOptions {
 	openRouterApiKey?: string
 	awsAccessKey?: string
 	awsSecretKey?: string
+	awsSessionToken?: string
 	awsRegion?: string
 	vertexProjectId?: string
 	vertexRegion?: string

+ 11 - 8
webview-ui/src/components/ApiOptions.tsx

@@ -153,6 +153,14 @@ const ApiOptions: React.FC<ApiOptionsProps> = ({ showModelOptions, apiErrorMessa
 						placeholder="Enter Secret Key...">
 						<span style={{ fontWeight: 500 }}>AWS Secret Key</span>
 					</VSCodeTextField>
+					<VSCodeTextField
+						value={apiConfiguration?.awsSessionToken || ""}
+						style={{ width: "100%" }}
+						type="password"
+						onInput={handleInputChange("awsSessionToken")}
+						placeholder="Enter Session Token...">
+						<span style={{ fontWeight: 500 }}>AWS Session Token</span>
+					</VSCodeTextField>
 					<div className="dropdown-container">
 						<label htmlFor="aws-region-dropdown">
 							<span style={{ fontWeight: 500 }}>AWS Region</span>
@@ -192,14 +200,9 @@ const ApiOptions: React.FC<ApiOptionsProps> = ({ showModelOptions, apiErrorMessa
 							marginTop: "5px",
 							color: "var(--vscode-descriptionForeground)",
 						}}>
-						These credentials are stored locally and only used to make API requests from this extension.
-						{!(apiConfiguration?.awsAccessKey && apiConfiguration?.awsSecretKey) && (
-							<VSCodeLink
-								href="https://docs.aws.amazon.com/IAM/latest/UserGuide/id_credentials_access-keys.html"
-								style={{ display: "inline" }}>
-								You can find your AWS access key and secret key here.
-							</VSCodeLink>
-						)}
+						Authenticate by either providing the keys above or use the default AWS credential providers,
+						i.e. ~/.aws/credentials or environment variables. These credentials are only used locally to
+						make API requests from this extension.
 					</p>
 				</div>
 			)}

+ 1 - 1
webview-ui/src/context/ExtensionStateContext.tsx

@@ -31,7 +31,7 @@ export const ExtensionStateContextProvider: React.FC<{ children: React.ReactNode
 			setState(message.state)
 			const config = message.state?.apiConfiguration
 			const hasKey = config
-				? [config.apiKey, config.openRouterApiKey, config.awsAccessKey, config.vertexProjectId].some(
+				? [config.apiKey, config.openRouterApiKey, config.awsRegion, config.vertexProjectId].some(
 						(key) => key !== undefined
 				  )
 				: false

+ 3 - 3
webview-ui/src/utils/validate.ts

@@ -9,8 +9,8 @@ export function validateApiConfiguration(apiConfiguration?: ApiConfiguration): s
 				}
 				break
 			case "bedrock":
-				if (!apiConfiguration.awsAccessKey || !apiConfiguration.awsSecretKey || !apiConfiguration.awsRegion) {
-					return "You must provide a valid AWS access key, secret key, and region."
+				if (!apiConfiguration.awsRegion) {
+					return "You must choose a region to use with AWS Bedrock."
 				}
 				break
 			case "openrouter":
@@ -26,4 +26,4 @@ export function validateApiConfiguration(apiConfiguration?: ApiConfiguration): s
 		}
 	}
 	return undefined
-}
+}