Support AWS profile to configure Bedrock Authentication

Added support for configurations under ~/.aws/credentials or ~/.aws/config.
This commit is contained in:
Lunchb0ne
2025-01-16 18:01:49 +00:00
parent 084599c9d0
commit 7a61e6ab74
6 changed files with 145 additions and 27 deletions

View File

@@ -0,0 +1,5 @@
---
"roo-cline": minor
---
Added suport for configuring Bedrock provider with AWS Profiles. Useful for users with SSO or other integrations who don't have access to long term credentials.

View File

@@ -1,7 +1,16 @@
// Mock AWS SDK credential providers
jest.mock("@aws-sdk/credential-providers", () => ({
fromIni: jest.fn().mockReturnValue({
accessKeyId: "profile-access-key",
secretAccessKey: "profile-secret-key",
}),
}))
import { AwsBedrockHandler } from "../bedrock"
import { MessageContent } from "../../../shared/api"
import { BedrockRuntimeClient } from "@aws-sdk/client-bedrock-runtime"
import { Anthropic } from "@anthropic-ai/sdk"
import { fromIni } from "@aws-sdk/credential-providers"
describe("AwsBedrockHandler", () => {
let handler: AwsBedrockHandler
@@ -30,6 +39,65 @@ describe("AwsBedrockHandler", () => {
})
expect(handlerWithoutCreds).toBeInstanceOf(AwsBedrockHandler)
})
it("should initialize with AWS profile credentials", () => {
const handlerWithProfile = new AwsBedrockHandler({
apiModelId: "anthropic.claude-3-5-sonnet-20241022-v2:0",
awsRegion: "us-east-1",
awsUseProfile: true,
awsProfile: "test-profile",
})
expect(handlerWithProfile).toBeInstanceOf(AwsBedrockHandler)
expect(handlerWithProfile["options"].awsUseProfile).toBe(true)
expect(handlerWithProfile["options"].awsProfile).toBe("test-profile")
})
it("should initialize with AWS profile enabled but no profile set", () => {
const handlerWithoutProfile = new AwsBedrockHandler({
apiModelId: "anthropic.claude-3-5-sonnet-20241022-v2:0",
awsRegion: "us-east-1",
awsUseProfile: true,
})
expect(handlerWithoutProfile).toBeInstanceOf(AwsBedrockHandler)
expect(handlerWithoutProfile["options"].awsUseProfile).toBe(true)
expect(handlerWithoutProfile["options"].awsProfile).toBeUndefined()
})
})
describe("AWS SDK client configuration", () => {
it("should configure client with profile credentials when profile mode is enabled", async () => {
// Import the fromIni function to mock it
jest.mock("@aws-sdk/credential-providers", () => ({
fromIni: jest.fn().mockReturnValue({
accessKeyId: "profile-access-key",
secretAccessKey: "profile-secret-key",
}),
}))
const handlerWithProfile = new AwsBedrockHandler({
apiModelId: "anthropic.claude-3-5-sonnet-20241022-v2:0",
awsRegion: "us-east-1",
awsUseProfile: true,
awsProfile: "test-profile",
})
// Mock a simple API call to verify credentials are used
const mockResponse = {
output: new TextEncoder().encode(JSON.stringify({ content: "test" })),
}
const mockSend = jest.fn().mockResolvedValue(mockResponse)
handlerWithProfile["client"] = {
send: mockSend,
} as unknown as BedrockRuntimeClient
await handlerWithProfile.completePrompt("test")
// Verify the client was configured with profile credentials
expect(mockSend).toHaveBeenCalled()
expect(fromIni).toHaveBeenCalledWith({
profile: "test-profile",
})
})
})
describe("createMessage", () => {

View File

@@ -4,6 +4,7 @@ import {
ConverseCommand,
BedrockRuntimeClientConfig,
} from "@aws-sdk/client-bedrock-runtime"
import { fromIni } from "@aws-sdk/credential-providers"
import { Anthropic } from "@anthropic-ai/sdk"
import { ApiHandler, SingleCompletionHandler } from "../"
import { ApiHandlerOptions, BedrockModelId, ModelInfo, bedrockDefaultModelId, bedrockModels } from "../../shared/api"
@@ -50,13 +51,17 @@ export class AwsBedrockHandler implements ApiHandler, SingleCompletionHandler {
constructor(options: ApiHandlerOptions) {
this.options = options
// Only include credentials if they actually exist
const clientConfig: BedrockRuntimeClientConfig = {
region: this.options.awsRegion || "us-east-1",
}
if (this.options.awsAccessKey && this.options.awsSecretKey) {
// Create credentials object with all properties at once
if (this.options.awsUseProfile && this.options.awsProfile) {
// Use profile-based credentials if enabled and profile is set
clientConfig.credentials = fromIni({
profile: this.options.awsProfile,
})
} else if (this.options.awsAccessKey && this.options.awsSecretKey) {
// Use direct credentials if provided
clientConfig.credentials = {
accessKeyId: this.options.awsAccessKey,
secretAccessKey: this.options.awsSecretKey,

View File

@@ -56,6 +56,8 @@ type GlobalStateKey =
| "glamaModelInfo"
| "awsRegion"
| "awsUseCrossRegionInference"
| "awsProfile"
| "awsUseProfile"
| "vertexProjectId"
| "vertexRegion"
| "lastShownAnnouncementId"
@@ -1147,6 +1149,8 @@ export class ClineProvider implements vscode.WebviewViewProvider {
awsSessionToken,
awsRegion,
awsUseCrossRegionInference,
awsProfile,
awsUseProfile,
vertexProjectId,
vertexRegion,
openAiBaseUrl,
@@ -1180,6 +1184,8 @@ export class ClineProvider implements vscode.WebviewViewProvider {
await this.storeSecret("awsSessionToken", awsSessionToken)
await this.updateGlobalState("awsRegion", awsRegion)
await this.updateGlobalState("awsUseCrossRegionInference", awsUseCrossRegionInference)
await this.updateGlobalState("awsProfile", awsProfile)
await this.updateGlobalState("awsUseProfile", awsUseProfile)
await this.updateGlobalState("vertexProjectId", vertexProjectId)
await this.updateGlobalState("vertexRegion", vertexRegion)
await this.updateGlobalState("openAiBaseUrl", openAiBaseUrl)
@@ -1795,6 +1801,8 @@ export class ClineProvider implements vscode.WebviewViewProvider {
awsSessionToken,
awsRegion,
awsUseCrossRegionInference,
awsProfile,
awsUseProfile,
vertexProjectId,
vertexRegion,
openAiBaseUrl,
@@ -1857,6 +1865,8 @@ export class ClineProvider implements vscode.WebviewViewProvider {
this.getSecret("awsSessionToken") as Promise<string | undefined>,
this.getGlobalState("awsRegion") as Promise<string | undefined>,
this.getGlobalState("awsUseCrossRegionInference") as Promise<boolean | undefined>,
this.getGlobalState("awsProfile") as Promise<string | undefined>,
this.getGlobalState("awsUseProfile") as Promise<boolean | undefined>,
this.getGlobalState("vertexProjectId") as Promise<string | undefined>,
this.getGlobalState("vertexRegion") as Promise<string | undefined>,
this.getGlobalState("openAiBaseUrl") as Promise<string | undefined>,
@@ -1936,6 +1946,8 @@ export class ClineProvider implements vscode.WebviewViewProvider {
awsSessionToken,
awsRegion,
awsUseCrossRegionInference,
awsProfile,
awsUseProfile,
vertexProjectId,
vertexRegion,
openAiBaseUrl,

View File

@@ -33,6 +33,8 @@ export interface ApiHandlerOptions {
awsUseCrossRegionInference?: boolean
awsUsePromptCache?: boolean
awspromptCacheId?: string
awsProfile?: string
awsUseProfile?: boolean
vertexProjectId?: string
vertexRegion?: string
openAiBaseUrl?: string

View File

@@ -342,30 +342,56 @@ const ApiOptions = ({ apiErrorMessage, modelIdErrorMessage }: ApiOptionsProps) =
{selectedProvider === "bedrock" && (
<div style={{ display: "flex", flexDirection: "column", gap: 5 }}>
<VSCodeTextField
value={apiConfiguration?.awsAccessKey || ""}
style={{ width: "100%" }}
type="password"
onInput={handleInputChange("awsAccessKey")}
placeholder="Enter Access Key...">
<span style={{ fontWeight: 500 }}>AWS Access Key</span>
</VSCodeTextField>
<VSCodeTextField
value={apiConfiguration?.awsSecretKey || ""}
style={{ width: "100%" }}
type="password"
onInput={handleInputChange("awsSecretKey")}
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>
<VSCodeRadioGroup
value={apiConfiguration?.awsUseProfile ? "profile" : "credentials"}
onChange={(e) => {
const value = (e.target as HTMLInputElement)?.value
const useProfile = value === "profile"
handleInputChange("awsUseProfile")({
target: { value: useProfile },
})
}}>
<VSCodeRadio value="credentials">AWS Credentials</VSCodeRadio>
<VSCodeRadio value="profile">AWS Profile</VSCodeRadio>
</VSCodeRadioGroup>
{/* AWS Profile Config Block */}
{apiConfiguration?.awsUseProfile ? (
<VSCodeTextField
value={apiConfiguration?.awsProfile || ""}
style={{ width: "100%" }}
onInput={handleInputChange("awsProfile")}
placeholder="Enter profile name">
<span style={{ fontWeight: 500 }}>AWS Profile Name</span>
</VSCodeTextField>
) : (
<>
{/* AWS Credentials Config Block */}
<VSCodeTextField
value={apiConfiguration?.awsAccessKey || ""}
style={{ width: "100%" }}
type="password"
onInput={handleInputChange("awsAccessKey")}
placeholder="Enter Access Key...">
<span style={{ fontWeight: 500 }}>AWS Access Key</span>
</VSCodeTextField>
<VSCodeTextField
value={apiConfiguration?.awsSecretKey || ""}
style={{ width: "100%" }}
type="password"
onInput={handleInputChange("awsSecretKey")}
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>