codewhisperer: Allow endpoint/region overrides for dev mode (#4324)

jpinkney-aws · web-flow · commit d933144cc6e0 · 2024-01-31T08:09:31.000-05:00
diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md
@@ -342,6 +342,17 @@ Example:
 }
 ```
 
+Overrides specifically for CodeWhisperer/Amazon Q can be set using the `aws.dev.codewhispererService` setting. This is a JSON object consisting of keys/values required to override API calls to CodeWhisperer/Amazon Q: `region` and `endpoint`. If this setting is present, then all keys need to be explicitly provided.
+
+Example:
+
+```json
+"aws.dev.codewhispererService": {
+    "region": "us-west-2",
+    "endpoint": "https://codewhisperer-gamma.example.com"
+}
+```
+
 ### Environment variables
 
 Environment variables can be used to modify the behaviour of VSCode. The following are environment variables that can be used to configure the extension:
diff --git a/src/amazonqFeatureDev/client/featureDev.ts b/src/amazonqFeatureDev/client/featureDev.ts
@@ -15,18 +15,19 @@ import * as FeatureDevProxyClient from './featuredevproxyclient'
 import apiConfig = require('./codewhispererruntime-2022-11-11.json')
 import { featureName } from '../constants'
 import { ApiError, ContentLengthError, UnknownApiError } from '../errors'
-import { endpoint, region } from '../../codewhisperer/models/constants'
 import { isAwsError, isCodeWhispererStreamingServiceException } from '../../shared/errors'
+import { getCodewhispererConfig } from '../../codewhisperer/client/codewhisperer'
 
 // Create a client for featureDev proxy client based off of aws sdk v2
 export async function createFeatureDevProxyClient(): Promise<FeatureDevProxyClient> {
     const bearerToken = await AuthUtil.instance.getBearerToken()
+    const cwsprConfig = getCodewhispererConfig()
     return (await globals.sdkClientBuilder.createAwsService(
         Service,
         {
             apiConfig: apiConfig,
-            region,
-            endpoint,
+            region: cwsprConfig.region,
+            endpoint: cwsprConfig.endpoint,
             token: new Token({ token: bearerToken }),
             // SETTING TO 0 FOR BETA. RE-ENABLE FOR RE-INVENT
             maxRetries: 0,
@@ -42,9 +43,10 @@ export async function createFeatureDevProxyClient(): Promise<FeatureDevProxyClie
 // Create a client for featureDev streaming based off of aws sdk v3
 async function createFeatureDevStreamingClient(): Promise<CodeWhispererStreaming> {
     const bearerToken = await AuthUtil.instance.getBearerToken()
+    const cwsprConfig = getCodewhispererConfig()
     const streamingClient = new CodeWhispererStreaming({
-        region,
-        endpoint,
+        region: cwsprConfig.region,
+        endpoint: cwsprConfig.endpoint,
         token: { token: bearerToken },
         // SETTING max attempts to 0 FOR BETA. RE-ENABLE FOR RE-INVENT
         // Implement exponential back off starting with a base of 500ms (500 + attempt^10)
diff --git a/src/codewhisperer/client/codewhisperer.ts b/src/codewhisperer/client/codewhisperer.ts
@@ -13,7 +13,6 @@ import {
     ListFeatureEvaluationsResponse,
     SendTelemetryEventRequest,
 } from './codewhispereruserclient'
-import * as CodeWhispererConstants from '../models/constants'
 import { ServiceOptions } from '../../shared/awsClientBuilder'
 import { hasVendedIamCredentials } from '../../auth/auth'
 import { CodeWhispererSettings } from '../util/codewhispererSettings'
@@ -31,6 +30,21 @@ import { getOptOutPreference } from '../util/commonUtil'
 import * as os from 'os'
 import { getClientId } from '../../shared/telemetry/util'
 import { extensionVersion } from '../../shared/vscode/env'
+import { DevSettings } from '../../shared/settings'
+
+export interface CodeWhispererConfig {
+    readonly region: string
+    readonly endpoint: string
+}
+
+export const defaultServiceConfig: CodeWhispererConfig = {
+    region: 'us-east-1',
+    endpoint: 'https://codewhisperer.us-east-1.amazonaws.com/',
+}
+
+export function getCodewhispererConfig(): CodeWhispererConfig {
+    return DevSettings.instance.getServiceConfig('codewhispererService', defaultServiceConfig)
+}
 
 export type ProgrammingLanguage = Readonly<
     CodeWhispererClient.ProgrammingLanguage | CodeWhispererUserClient.ProgrammingLanguage
@@ -80,14 +94,14 @@ export type Imports = CodeWhispererUserClient.Imports
 export class DefaultCodeWhispererClient {
     private async createSdkClient(): Promise<CodeWhispererClient> {
         const isOptedOut = CodeWhispererSettings.instance.isOptoutEnabled()
-
+        const cwsprConfig = getCodewhispererConfig()
         return (await globals.sdkClientBuilder.createAwsService(
             Service,
             {
                 apiConfig: apiConfig,
-                region: CodeWhispererConstants.region,
+                region: cwsprConfig.region,
                 credentials: await AuthUtil.instance.getCredentials(),
-                endpoint: CodeWhispererConstants.endpoint,
+                endpoint: cwsprConfig.endpoint,
                 onRequestSetup: [
                     req => {
                         if (req.operation === 'listRecommendations') {
@@ -123,12 +137,13 @@ export class DefaultCodeWhispererClient {
         session.setFetchCredentialStart()
         const bearerToken = await AuthUtil.instance.getBearerToken()
         session.setSdkApiCallStart()
+        const cwsprConfig = getCodewhispererConfig()
         return (await globals.sdkClientBuilder.createAwsService(
             Service,
             {
                 apiConfig: userApiConfig,
-                region: CodeWhispererConstants.region,
-                endpoint: CodeWhispererConstants.endpoint,
+                region: cwsprConfig.region,
+                endpoint: cwsprConfig.endpoint,
                 credentials: new Credentials({ accessKeyId: 'xxx', secretAccessKey: 'xxx' }),
                 onRequestSetup: [
                     req => {
diff --git a/src/codewhisperer/models/constants.ts b/src/codewhisperer/models/constants.ts
@@ -3,13 +3,6 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-/**
- * SDK Client
- */
-export const endpoint = 'https://codewhisperer.us-east-1.amazonaws.com/'
-
-export const region = 'us-east-1'
-
 /**
  * Automated and manual trigger
  */
diff --git a/src/shared/clients/codecatalystClient.ts b/src/shared/clients/codecatalystClient.ts
@@ -47,7 +47,7 @@ export const defaultServiceConfig: CodeCatalystConfig = {
 }
 
 export function getCodeCatalystConfig(): CodeCatalystConfig {
-    return DevSettings.instance.getCodeCatalystConfig(defaultServiceConfig)
+    return DevSettings.instance.getServiceConfig('codecatalystService', defaultServiceConfig)
 }
 
 export interface DevEnvironment extends CodeCatalyst.DevEnvironmentSummary {
diff --git a/src/shared/settings.ts b/src/shared/settings.ts
@@ -6,6 +6,7 @@
 import * as vscode from 'vscode'
 import * as packageJson from '../../package.json'
 import * as codecatalyst from './clients/codecatalystClient'
+import * as codewhisperer from '../codewhisperer/client/codewhisperer'
 import { getLogger } from './logger'
 import { cast, FromDescriptor, Record, TypeConstructor, TypeDescriptor } from './utilities/typeConstructors'
 import { assertHasProps, ClassToInterfaceType, keys } from './utilities/tsUtils'
@@ -662,11 +663,18 @@ const devSettings = {
     renderDebugDetails: Boolean,
     endpoints: Record(String, String),
     codecatalystService: Record(String, String),
+    codewhispererService: Record(String, String),
     ssoCacheDirectory: String,
 }
 type ResolvedDevSettings = FromDescriptor<typeof devSettings>
 type AwsDevSetting = keyof ResolvedDevSettings
 
+type ServiceClients = keyof ServiceTypeMap
+interface ServiceTypeMap {
+    codecatalystService: codecatalyst.CodeCatalystConfig
+    codewhispererService: codewhisperer.CodeWhispererConfig
+}
+
 /**
  * Developer settings are intended to be used by developers of this codebase for anything
  * that may be potentially useful during active development and/or testing. Examples include
@@ -721,14 +729,14 @@ export class DevSettings extends Settings.define('aws.dev', devSettings) {
         return Object.keys(this.activeSettings).length > 0
     }
 
-    public getCodeCatalystConfig(
-        defaultConfig: codecatalyst.CodeCatalystConfig
-    ): Readonly<codecatalyst.CodeCatalystConfig> {
-        const devSetting = 'codecatalystService'
-        const devConfig = this.get(devSetting, {})
+    public getServiceConfig<T extends ServiceClients>(
+        devSetting: T,
+        defaultConfig: ServiceTypeMap[T]
+    ): Readonly<ServiceTypeMap[T]> {
+        const devConfig = this.get<ServiceClients>(devSetting, {})
 
         if (Object.keys(devConfig).length === 0) {
-            this.logCodeCatalystConfigOnce('default')
+            this.logConfigOnce(devSetting, 'default')
             return defaultConfig
         }
 
@@ -739,8 +747,8 @@ export class DevSettings extends Settings.define('aws.dev', devSettings) {
             throw ToolkitError.chain(err, `Dev setting '${devSetting}' has missing or invalid properties.`)
         }
 
-        this.logCodeCatalystConfigOnce(JSON.stringify(devConfig, undefined, 4))
-        return devConfig as unknown as codecatalyst.CodeCatalystConfig
+        this.logConfigOnce(devSetting, JSON.stringify(devConfig, undefined, 4))
+        return devConfig as unknown as ServiceTypeMap[T]
     }
 
     public override get<K extends AwsDevSetting>(key: K, defaultValue: ResolvedDevSettings[K]) {
@@ -770,8 +778,8 @@ export class DevSettings extends Settings.define('aws.dev', devSettings) {
         }
     }
 
-    private logCodeCatalystConfigOnce = onceChanged(val => {
-        getLogger().info(`using CodeCatalyst service configuration: ${val}`)
+    private logConfigOnce = onceChanged((serviceName, val) => {
+        getLogger().info(`using ${serviceName} service configuration: ${val}`)
     })
 
     static #instance: DevSettings
diff --git a/src/test/shared/settingsConfiguration.test.ts b/src/test/shared/settingsConfiguration.test.ts
@@ -313,41 +313,75 @@ describe('DevSetting', function () {
         })
     })
 
-    describe('getCodeCatalystConfig()', function () {
-        const devSettingName = 'aws.dev.codecatalystService'
-        const defaultConfig = {
-            region: 'default',
-            endpoint: 'default',
-            hostname: 'default',
-            gitHostname: 'default',
-        }
-
-        it('throws an error for incomplete dev configuration', async function () {
-            const testSetting = {
-                // missing region
-                endpoint: 'test_endpoint',
-                hostname: 'test_hostname',
-                gitHostname: 'test_githostname',
+    describe('getServiceConfig()', function () {
+        describe('get CodeCatalyst config', function () {
+            const devSettingName = 'aws.dev.codecatalystService'
+            const defaultConfig = {
+                region: 'default',
+                endpoint: 'default',
+                hostname: 'default',
+                gitHostname: 'default',
             }
 
-            await settings.update(devSettingName, testSetting)
-            assert.throws(() => sut.getCodeCatalystConfig(defaultConfig), ToolkitError)
+            it('throws an error for incomplete dev configuration', async function () {
+                const testSetting = {
+                    // missing region
+                    endpoint: 'test_endpoint',
+                    hostname: 'test_hostname',
+                    gitHostname: 'test_githostname',
+                }
+
+                await settings.update(devSettingName, testSetting)
+                assert.throws(() => sut.getServiceConfig('codecatalystService', defaultConfig), ToolkitError)
+            })
+
+            it('returns dev settings configuration when provided', async function () {
+                const testSetting = {
+                    region: 'test_region',
+                    endpoint: 'test_endpoint',
+                    hostname: 'test_hostname',
+                    gitHostname: 'test_githostname',
+                }
+
+                await settings.update(devSettingName, testSetting)
+                assert.deepStrictEqual(sut.getServiceConfig('codecatalystService', defaultConfig), testSetting)
+            })
+
+            it('returns default configuration when dev settings are not provided', function () {
+                assert.deepStrictEqual(sut.getServiceConfig('codecatalystService', defaultConfig), defaultConfig)
+            })
         })
 
-        it('returns dev settings configuration when provided', async function () {
-            const testSetting = {
-                region: 'test_region',
-                endpoint: 'test_endpoint',
-                hostname: 'test_hostname',
-                gitHostname: 'test_githostname',
+        describe('get Codewhisperer config', function () {
+            const devSettingName = 'aws.dev.codewhispererService'
+            const defaultConfig = {
+                region: 'default',
+                endpoint: 'default',
             }
 
-            await settings.update(devSettingName, testSetting)
-            assert.deepStrictEqual(sut.getCodeCatalystConfig(defaultConfig), testSetting)
-        })
+            it('throws an error for incomplete dev configuration', async function () {
+                const testSetting = {
+                    // missing region
+                    endpoint: 'test_endpoint',
+                }
+
+                await settings.update(devSettingName, testSetting)
+                assert.throws(() => sut.getServiceConfig('codewhispererService', defaultConfig), ToolkitError)
+            })
 
-        it('returns default configuration when dev settings are not provided', function () {
-            assert.deepStrictEqual(sut.getCodeCatalystConfig(defaultConfig), defaultConfig)
+            it('returns dev settings configuration when provided', async function () {
+                const testSetting = {
+                    region: 'test_region',
+                    endpoint: 'test_endpoint',
+                }
+
+                await settings.update(devSettingName, testSetting)
+                assert.deepStrictEqual(sut.getServiceConfig('codewhispererService', defaultConfig), testSetting)
+            })
+
+            it('returns default configuration when dev settings are not provided', function () {
+                assert.deepStrictEqual(sut.getServiceConfig('codewhispererService', defaultConfig), defaultConfig)
+            })
         })
     })
 })

Original file line number	Diff line number	Diff line change
`@@ -47,7 +47,7 @@ export const defaultServiceConfig: CodeCatalystConfig = {`
`47`	`47`	`}`
`48`	`48`
`49`	`49`	`export function getCodeCatalystConfig(): CodeCatalystConfig {`
`50`		`- return DevSettings.instance.getCodeCatalystConfig(defaultServiceConfig)`
	`50`	`+ return DevSettings.instance.getServiceConfig('codecatalystService', defaultServiceConfig)`
`51`	`51`	`}`
`52`	`52`
`53`	`53`	`export interface DevEnvironment extends CodeCatalyst.DevEnvironmentSummary {`