Skip to content

Commit a4c91c6

Browse files
kcwhiteKevin Whitedaniel-lxs
authored
Adding support for custom VPC endpoints when using AWS Bedrock models (#3947)
* feat: Add custom VPC endpoint support for AWS Bedrock * fix: Fix TypeScript error in Bedrock.tsx * fix: Update VPC endpoint UI to match Cline's implementation * Fix AWS Bedrock VPC endpoint UI implementation - Changed checkbox label to 'Use custom VPC endpoint' to match Cline - Fixed conditional rendering to show text field when checkbox is checked - Ensured placeholder text appears correctly - Maintained proper styling for consistency * Fix AWS Bedrock VPC endpoint UI implementation to match Cline exactly - Added state variable to track checkbox selection - Fixed conditional rendering to show/hide text field based on checkbox state - Maintained proper styling and placeholder text * Fix AWS Bedrock VPC endpoint UI implementation with proper event handling - Fixed checkbox onChange handler to accept boolean directly instead of event object - Added unit tests to verify the behavior - Maintained proper styling and placeholder text * Update Bedrock VPC endpoint tests with proper test IDs * Improve AWS Bedrock VPC endpoint text field alignment - Removed left margin from text field to align with checkbox - Maintained proper styling and placeholder text * Preserve AWS Bedrock VPC endpoint URL when toggling checkbox - Added awsBedrockEndpointEnabled field to schema - Modified Bedrock provider to check both endpoint URL and enabled flag - Updated UI to preserve endpoint URL when checkbox is toggled - Maintained proper alignment with checkbox * Implement AWS Bedrock Custom VPC Endpoint functionality * fix: update ApiConfiguration to ProviderSettings in Bedrock tests and regenerate types * fix: update all instances of ApiConfiguration to ProviderSettings in Bedrock tests * Fixed broken unit test * Add changeset for Bedrock VPC endpoint support * informative placeholder * Bug fixes * Fixed failing tests * Add example URLs to Bedrock VPC endpoint section and update tests * Fix truncated test assertion in Bedrock.test.tsx that was breaking the UI * Refactor mock components in Bedrock.test.tsx for improved data-testid handling * feat(i18n): add VPC endpoint translations for AWS Bedrock settings * test: update Bedrock component tests for internationalized strings --------- Co-authored-by: Kevin White <[email protected]> Co-authored-by: Daniel <[email protected]> Co-authored-by: Daniel Riccio <[email protected]>
1 parent 72d06f5 commit a4c91c6

File tree

24 files changed

+729
-5
lines changed

24 files changed

+729
-5
lines changed

‎.changeset/bedrock-vpc-endpoint.md‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
---
2+
"roo-cline": patch
3+
---
4+
5+
Add AWS Bedrock VPC endpoint support, allowing users to connect to Bedrock through private VPC endpoints. This includes UI configuration options, updated types, and comprehensive test coverage.

‎packages/types/src/provider-settings.ts‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -100,6 +100,8 @@ const bedrockSchema = apiModelIdProviderModelSchema.extend({
100100
awsProfile: z.string().optional(),
101101
awsUseProfile: z.boolean().optional(),
102102
awsCustomArn: z.string().optional(),
103+
awsBedrockEndpointEnabled: z.boolean().optional(),
104+
awsBedrockEndpoint: z.string().optional(),
103105
})
104106

105107
const vertexSchema = apiModelIdProviderModelSchema.extend({
@@ -283,6 +285,8 @@ export const PROVIDER_SETTINGS_KEYS = keysOf<ProviderSettings>()([
283285
"awsProfile",
284286
"awsUseProfile",
285287
"awsCustomArn",
288+
"awsBedrockEndpointEnabled",
289+
"awsBedrockEndpoint",
286290
// Google Vertex
287291
"vertexKeyFile",
288292
"vertexJsonCredentials",

‎src/api/providers/__tests__/bedrock-vpc-endpoint.test.ts‎

Lines changed: 178 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,178 @@
1+
// Mock AWS SDK credential providers
2+
jest.mock("@aws-sdk/credential-providers", () => {
3+
const mockFromIni = jest.fn().mockReturnValue({
4+
accessKeyId: "profile-access-key",
5+
secretAccessKey: "profile-secret-key",
6+
})
7+
return { fromIni: mockFromIni }
8+
})
9+
10+
// Mock BedrockRuntimeClient and ConverseStreamCommand
11+
const mockBedrockRuntimeClient = jest.fn()
12+
const mockSend = jest.fn().mockResolvedValue({
13+
stream: [],
14+
})
15+
16+
jest.mock("@aws-sdk/client-bedrock-runtime", () => ({
17+
BedrockRuntimeClient: mockBedrockRuntimeClient.mockImplementation(() => ({
18+
send: mockSend,
19+
})),
20+
ConverseStreamCommand: jest.fn(),
21+
ConverseCommand: jest.fn(),
22+
}))
23+
24+
import { AwsBedrockHandler } from "../bedrock"
25+
26+
describe("AWS Bedrock VPC Endpoint Functionality", () => {
27+
beforeEach(() => {
28+
// Clear all mocks before each test
29+
jest.clearAllMocks()
30+
})
31+
32+
// Test Scenario 1: Input Validation Test
33+
describe("VPC Endpoint URL Validation", () => {
34+
it("should configure client with endpoint URL when both URL and enabled flag are provided", () => {
35+
// Create handler with endpoint URL and enabled flag
36+
new AwsBedrockHandler({
37+
apiModelId: "anthropic.claude-3-5-sonnet-20241022-v2:0",
38+
awsAccessKey: "test-access-key",
39+
awsSecretKey: "test-secret-key",
40+
awsRegion: "us-east-1",
41+
awsBedrockEndpoint: "https://bedrock-vpc.example.com",
42+
awsBedrockEndpointEnabled: true,
43+
})
44+
45+
// Verify the client was created with the correct endpoint
46+
expect(mockBedrockRuntimeClient).toHaveBeenCalledWith(
47+
expect.objectContaining({
48+
region: "us-east-1",
49+
endpoint: "https://bedrock-vpc.example.com",
50+
}),
51+
)
52+
})
53+
54+
it("should not configure client with endpoint URL when URL is provided but enabled flag is false", () => {
55+
// Create handler with endpoint URL but disabled flag
56+
new AwsBedrockHandler({
57+
apiModelId: "anthropic.claude-3-5-sonnet-20241022-v2:0",
58+
awsAccessKey: "test-access-key",
59+
awsSecretKey: "test-secret-key",
60+
awsRegion: "us-east-1",
61+
awsBedrockEndpoint: "https://bedrock-vpc.example.com",
62+
awsBedrockEndpointEnabled: false,
63+
})
64+
65+
// Verify the client was created without the endpoint
66+
expect(mockBedrockRuntimeClient).toHaveBeenCalledWith(
67+
expect.objectContaining({
68+
region: "us-east-1",
69+
}),
70+
)
71+
72+
// Verify the endpoint property is not present
73+
const clientConfig = mockBedrockRuntimeClient.mock.calls[0][0]
74+
expect(clientConfig).not.toHaveProperty("endpoint")
75+
})
76+
})
77+
78+
// Test Scenario 2: Edge Case Tests
79+
describe("Edge Cases", () => {
80+
it("should handle empty endpoint URL gracefully", () => {
81+
// Create handler with empty endpoint URL but enabled flag
82+
new AwsBedrockHandler({
83+
apiModelId: "anthropic.claude-3-5-sonnet-20241022-v2:0",
84+
awsAccessKey: "test-access-key",
85+
awsSecretKey: "test-secret-key",
86+
awsRegion: "us-east-1",
87+
awsBedrockEndpoint: "",
88+
awsBedrockEndpointEnabled: true,
89+
})
90+
91+
// Verify the client was created without the endpoint (since it's empty)
92+
expect(mockBedrockRuntimeClient).toHaveBeenCalledWith(
93+
expect.objectContaining({
94+
region: "us-east-1",
95+
}),
96+
)
97+
98+
// Verify the endpoint property is not present
99+
const clientConfig = mockBedrockRuntimeClient.mock.calls[0][0]
100+
expect(clientConfig).not.toHaveProperty("endpoint")
101+
})
102+
103+
it("should handle undefined endpoint URL gracefully", () => {
104+
// Create handler with undefined endpoint URL but enabled flag
105+
new AwsBedrockHandler({
106+
apiModelId: "anthropic.claude-3-5-sonnet-20241022-v2:0",
107+
awsAccessKey: "test-access-key",
108+
awsSecretKey: "test-secret-key",
109+
awsRegion: "us-east-1",
110+
awsBedrockEndpoint: undefined,
111+
awsBedrockEndpointEnabled: true,
112+
})
113+
114+
// Verify the client was created without the endpoint
115+
expect(mockBedrockRuntimeClient).toHaveBeenCalledWith(
116+
expect.objectContaining({
117+
region: "us-east-1",
118+
}),
119+
)
120+
121+
// Verify the endpoint property is not present
122+
const clientConfig = mockBedrockRuntimeClient.mock.calls[0][0]
123+
expect(clientConfig).not.toHaveProperty("endpoint")
124+
})
125+
})
126+
127+
// Test Scenario 4: Error Handling Tests
128+
describe("Error Handling", () => {
129+
it("should handle invalid endpoint URLs by passing them directly to AWS SDK", () => {
130+
// Create handler with an invalid URL format
131+
new AwsBedrockHandler({
132+
apiModelId: "anthropic.claude-3-5-sonnet-20241022-v2:0",
133+
awsAccessKey: "test-access-key",
134+
awsSecretKey: "test-secret-key",
135+
awsRegion: "us-east-1",
136+
awsBedrockEndpoint: "invalid-url-format",
137+
awsBedrockEndpointEnabled: true,
138+
})
139+
140+
// Verify the client was created with the invalid endpoint
141+
// (AWS SDK will handle the validation/errors)
142+
expect(mockBedrockRuntimeClient).toHaveBeenCalledWith(
143+
expect.objectContaining({
144+
region: "us-east-1",
145+
endpoint: "invalid-url-format",
146+
}),
147+
)
148+
})
149+
})
150+
151+
// Test Scenario 5: Persistence Tests
152+
describe("Persistence", () => {
153+
it("should maintain consistent behavior across multiple requests", async () => {
154+
// Create handler with endpoint URL and enabled flag
155+
const handler = new AwsBedrockHandler({
156+
apiModelId: "anthropic.claude-3-5-sonnet-20241022-v2:0",
157+
awsAccessKey: "test-access-key",
158+
awsSecretKey: "test-secret-key",
159+
awsRegion: "us-east-1",
160+
awsBedrockEndpoint: "https://bedrock-vpc.example.com",
161+
awsBedrockEndpointEnabled: true,
162+
})
163+
164+
// Reset mock to clear the constructor call
165+
mockBedrockRuntimeClient.mockClear()
166+
167+
// Make a request
168+
try {
169+
await handler.completePrompt("Test prompt")
170+
} catch (error) {
171+
// Ignore errors, we're just testing the client configuration
172+
}
173+
174+
// Verify the client was configured with the endpoint
175+
expect(mockSend).toHaveBeenCalled()
176+
})
177+
})
178+
})

‎src/api/providers/bedrock.ts‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -169,6 +169,9 @@ export class AwsBedrockHandler extends BaseProvider implements SingleCompletionH
169169

170170
const clientConfig: BedrockRuntimeClientConfig = {
171171
region: this.options.awsRegion,
172+
// Add the endpoint configuration when specified and enabled
173+
...(this.options.awsBedrockEndpoint &&
174+
this.options.awsBedrockEndpointEnabled && { endpoint: this.options.awsBedrockEndpoint }),
172175
}
173176

174177
if (this.options.awsUseProfile && this.options.awsProfile) {

‎src/core/config/__tests__/importExport.test.ts‎

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -231,10 +231,8 @@ describe("importExport", () => {
231231
customModesManager: mockCustomModesManager,
232232
})
233233

234-
expect(result).toEqual({
235-
success: false,
236-
error: "Expected property name or '}' in JSON at position 2",
237-
})
234+
expect(result.success).toBe(false)
235+
expect(result.error).toMatch(/^Expected property name or '}' in JSON at position 2/)
238236
expect(fs.readFile).toHaveBeenCalledWith("/mock/path/settings.json", "utf-8")
239237
expect(mockProviderSettingsManager.import).not.toHaveBeenCalled()
240238
expect(mockContextProxy.setValues).not.toHaveBeenCalled()

‎webview-ui/src/components/settings/providers/Bedrock.tsx‎

Lines changed: 32 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
import { useCallback } from "react"
1+
import { useCallback, useState, useEffect } from "react"
22
import { Checkbox } from "vscrui"
33
import { VSCodeTextField, VSCodeRadio, VSCodeRadioGroup } from "@vscode/webview-ui-toolkit/react"
44

@@ -17,6 +17,12 @@ type BedrockProps = {
1717

1818
export const Bedrock = ({ apiConfiguration, setApiConfigurationField, selectedModelInfo }: BedrockProps) => {
1919
const { t } = useAppTranslation()
20+
const [awsEndpointSelected, setAwsEndpointSelected] = useState(!!apiConfiguration?.awsBedrockEndpointEnabled)
21+
22+
// Update the endpoint enabled state when the configuration changes
23+
useEffect(() => {
24+
setAwsEndpointSelected(!!apiConfiguration?.awsBedrockEndpointEnabled)
25+
}, [apiConfiguration?.awsBedrockEndpointEnabled])
2026

2127
const handleInputChange = useCallback(
2228
<K extends keyof ProviderSettings, E>(
@@ -120,6 +126,31 @@ export const Bedrock = ({ apiConfiguration, setApiConfigurationField, selectedMo
120126
{t("settings:providers.cacheUsageNote")}
121127
</div>
122128
</div>
129+
<Checkbox
130+
checked={awsEndpointSelected}
131+
onChange={(isChecked) => {
132+
setAwsEndpointSelected(isChecked)
133+
setApiConfigurationField("awsBedrockEndpointEnabled", isChecked)
134+
}}>
135+
{t("settings:providers.awsBedrockVpc.useCustomVpcEndpoint")}
136+
</Checkbox>
137+
{awsEndpointSelected && (
138+
<>
139+
<VSCodeTextField
140+
value={apiConfiguration?.awsBedrockEndpoint || ""}
141+
style={{ width: "100%", marginTop: 3, marginBottom: 5 }}
142+
type="url"
143+
onInput={handleInputChange("awsBedrockEndpoint")}
144+
placeholder={t("settings:providers.awsBedrockVpc.vpcEndpointUrlPlaceholder")}
145+
data-testid="vpc-endpoint-input"
146+
/>
147+
<div className="text-sm text-vscode-descriptionForeground ml-6 mt-1 mb-3">
148+
{t("settings:providers.awsBedrockVpc.examples")}
149+
<div className="ml-2">• https://vpce-xxx.bedrock.region.vpce.amazonaws.com/</div>
150+
<div className="ml-2">• https://gateway.my-company.com/route/app/bedrock</div>
151+
</div>
152+
</>
153+
)}
123154
</>
124155
)
125156
}

0 commit comments

Comments
 (0)