Merge remote-tracking branch 'upstream/master' into profile-cache

Author: Will-ShaoHua

packages/amazonq/.changes/next-release/Feature-2a696d00-e8c8-44a4-ab2c-2204b4d8e31d.json

@@ -65,7 +65,7 @@ describe('RegionProfileManager', function () {
             const mockClient = {
                 listAvailableProfiles: listProfilesStub,
             }
-            const createClientStub = sinon.stub(sut, '_createQClient').resolves(mockClient)
+            const createClientStub = sinon.stub(sut, 'createQClient').resolves(mockClient)
 
             const r = await sut.listRegionProfile()
 
@@ -234,65 +234,13 @@ describe('RegionProfileManager', function () {
     })
 
     describe('createQClient', function () {
-        it(`should configure the endpoint and region from a profile`, async function () {
-            await setupConnection('idc')
-
-            const iadClient = await sut.createQClient({
-                name: 'foo',
-                region: 'us-east-1',
-                arn: 'arn',
-                description: 'description',
-            })
-
-            assert.deepStrictEqual(iadClient.config.region, 'us-east-1')
-            assert.deepStrictEqual(iadClient.endpoint.href, 'https://q.us-east-1.amazonaws.com/')
-
-            const fraClient = await sut.createQClient({
-                name: 'bar',
-                region: 'eu-central-1',
-                arn: 'arn',
-                description: 'description',
-            })
-
-            assert.deepStrictEqual(fraClient.config.region, 'eu-central-1')
-            assert.deepStrictEqual(fraClient.endpoint.href, 'https://q.eu-central-1.amazonaws.com/')
-        })
-
-        it(`should throw if the region is not supported or recognizable by Q`, async function () {
-            await setupConnection('idc')
-
-            await assert.rejects(
-                async () => {
-                    await sut.createQClient({
-                        name: 'foo',
-                        region: 'ap-east-1',
-                        arn: 'arn',
-                        description: 'description',
-                    })
-                },
-                { message: /trying to initiatize Q client with unrecognizable region/ }
-            )
-
-            await assert.rejects(
-                async () => {
-                    await sut.createQClient({
-                        name: 'foo',
-                        region: 'unknown-somewhere',
-                        arn: 'arn',
-                        description: 'description',
-                    })
-                },
-                { message: /trying to initiatize Q client with unrecognizable region/ }
-            )
-        })
-
         it(`should configure the endpoint and region correspondingly`, async function () {
             await setupConnection('idc')
             await sut.switchRegionProfile(profileFoo, 'user')
             assert.deepStrictEqual(sut.activeRegionProfile, profileFoo)
             const conn = authUtil.conn as SsoConnection
 
-            const client = await sut._createQClient('eu-central-1', 'https://amazon.com/', conn)
+            const client = await sut.createQClient('eu-central-1', 'https://amazon.com/', conn)
 
             assert.deepStrictEqual(client.config.region, 'eu-central-1')
             assert.deepStrictEqual(client.endpoint.href, 'https://amazon.com/')

packages/amazonq/test/unit/codewhisperer/region/regionProfileManager.test.ts

@@ -72,7 +72,12 @@ import { AuthUtil } from './util/authUtil'
 import { ImportAdderProvider } from './service/importAdderProvider'
 import { TelemetryHelper } from './util/telemetryHelper'
 import { openUrl } from '../shared/utilities/vsCodeUtils'
-import { notifyNewCustomizations, onProfileChangedListener } from './util/customizationUtil'
+import {
+    getAvailableCustomizationsList,
+    getSelectedCustomization,
+    notifyNewCustomizations,
+    switchToBaseCustomizationAndNotify,
+} from './util/customizationUtil'
 import { CodeWhispererCommandBackend, CodeWhispererCommandDeclarations } from './commands/gettingStartedPageCommands'
 import { SecurityIssueHoverProvider } from './service/securityIssueHoverProvider'
 import { SecurityIssueCodeActionProvider } from './service/securityIssueCodeActionProvider'
@@ -338,7 +343,26 @@ export async function activate(context: ExtContext): Promise<void> {
             SecurityIssueCodeActionProvider.instance
         ),
         vscode.commands.registerCommand('aws.amazonq.openEditorAtRange', openEditorAtRange),
-        auth.regionProfileManager.onDidChangeRegionProfile(onProfileChangedListener)
+        auth.regionProfileManager.onDidChangeRegionProfile(() => {
+            // Validate user still has access to the selected customization.
+            const selectedCustomization = getSelectedCustomization()
+            // No need to validate base customization which has empty arn.
+            if (selectedCustomization.arn.length > 0) {
+                getAvailableCustomizationsList()
+                    .then(async (customizations) => {
+                        const r = customizations.find((it) => it.arn === selectedCustomization.arn)
+                        if (!r) {
+                            await switchToBaseCustomizationAndNotify()
+                        }
+                    })
+                    .catch((e) => {
+                        getLogger().error(
+                            `encounter error while validating selected customization on profile change: %s`,
+                            (e as Error).message
+                        )
+                    })
+            }
+        })
     )
 
     // run the auth startup code with context for telemetry

packages/core/src/codewhisperer/activation.ts

@@ -7,17 +7,19 @@ import { AWSError, Credentials, Service } from 'aws-sdk'
 import globals from '../../shared/extensionGlobals'
 import * as CodeWhispererClient from './codewhispererclient'
 import * as CodeWhispererUserClient from './codewhispereruserclient'
-import { SendTelemetryEventRequest } from './codewhispereruserclient'
+import { ListAvailableCustomizationsResponse, SendTelemetryEventRequest } from './codewhispereruserclient'
 import { ServiceOptions } from '../../shared/awsClientBuilder'
 import { hasVendedIamCredentials } from '../../auth/auth'
 import { CodeWhispererSettings } from '../util/codewhispererSettings'
 import { PromiseResult } from 'aws-sdk/lib/request'
 import { AuthUtil } from '../util/authUtil'
 import { isSsoConnection } from '../../auth/connection'
+import { pageableToCollection } from '../../shared/utilities/collectionUtils'
 import apiConfig = require('./service-2.json')
 import userApiConfig = require('./user-service-2.json')
 import { session } from '../util/codeWhispererSession'
 import { getLogger } from '../../shared/logger/logger'
+import { indent } from '../../shared/utilities/textUtilities'
 import { getClientId, getOptOutPreference, getOperatingSystem } from '../../shared/telemetry/util'
 import { extensionVersion, getServiceEnvVarConfig } from '../../shared/vscode/env'
 import { DevSettings } from '../../shared/settings'
@@ -217,6 +219,28 @@ export class DefaultCodeWhispererClient {
             .promise()
     }
 
+    public async listAvailableCustomizations(): Promise<ListAvailableCustomizationsResponse[]> {
+        const client = await this.createUserSdkClient()
+        const profile = AuthUtil.instance.regionProfileManager.activeRegionProfile
+        const requester = async (request: CodeWhispererUserClient.ListAvailableCustomizationsRequest) =>
+            client.listAvailableCustomizations(request).promise()
+        return pageableToCollection(requester, { profileArn: profile?.arn }, 'nextToken')
+            .promise()
+            .then((resps) => {
+                let logStr = 'amazonq: listAvailableCustomizations API request:'
+                for (const resp of resps) {
+                    const requestId = resp.$response.requestId
+                    logStr += `\n${indent('RequestID: ', 4)}${requestId},\n${indent('Customizations:', 4)}`
+                    for (const [index, c] of resp.customizations.entries()) {
+                        const entry = `${index.toString().padStart(2, '0')}: ${c.name?.trim()}`
+                        logStr += `\n${indent(entry, 8)}`
+                    }
+                }
+                getLogger().debug(logStr)
+                return resps
+            })
+    }
+
     public async sendTelemetryEvent(request: SendTelemetryEventRequest) {
         const requestWithCommonFields: SendTelemetryEventRequest = {
             ...request,

packages/core/src/codewhisperer/client/codewhisperer.ts

@@ -99,13 +99,7 @@ export * as diagnosticsProvider from './service/diagnosticsProvider'
 export * from './ui/codeWhispererNodes'
 export { SecurityScanError, SecurityScanTimedOutError } from '../codewhisperer/models/errors'
 export * as CodeWhispererConstants from '../codewhisperer/models/constants'
-export {
-    getSelectedCustomization,
-    setSelectedCustomization,
-    baseCustomization,
-    onProfileChangedListener,
-    CustomizationProvider,
-} from './util/customizationUtil'
+export { getSelectedCustomization, setSelectedCustomization, baseCustomization } from './util/customizationUtil'
 export { Container } from './service/serviceContainer'
 export * from './util/gitUtil'
 export * from './ui/prompters'

packages/core/src/codewhisperer/index.ts

@@ -48,12 +48,7 @@ const endpoints = createConstantMap({
  * 'update' -> plugin auto select the profile on users' behalf as there is only 1 profile
  * 'reload' -> on plugin restart, plugin will try to reload previous selected profile
  */
-export type ProfileSwitchIntent = 'user' | 'auth' | 'update' | 'reload' | 'customization'
-
-export type ProfileChangedEvent = {
-    profile: RegionProfile | undefined
-    intent: ProfileSwitchIntent
-}
+export type ProfileSwitchIntent = 'user' | 'auth' | 'update' | 'reload'
 
 // Only "valid" state will have non null profiles
 type CachedApiResult = {
@@ -64,7 +59,7 @@ type CachedApiResult = {
 export class RegionProfileManager {
     private static logger = getLogger()
     private _activeRegionProfile: RegionProfile | undefined
-    private _onDidChangeRegionProfile = new vscode.EventEmitter<ProfileChangedEvent>()
+    private _onDidChangeRegionProfile = new vscode.EventEmitter<RegionProfile | undefined>()
     public readonly onDidChangeRegionProfile = this._onDidChangeRegionProfile.event
 
     // Store the last API results (for UI propuse) so we don't need to call service again if doesn't require "latest" result
@@ -135,7 +130,7 @@ export class RegionProfileManager {
         await this.cacheApiResult(undefined)
 
         for (const [region, endpoint] of endpoints.entries()) {
-            const client = await this._createQClient(region, endpoint, conn as SsoConnection)
+            const client = await this.createQClient(region, endpoint, conn as SsoConnection)
             const requester = async (request: CodeWhispererUserClient.ListAvailableProfilesRequest) =>
                 client.listAvailableProfiles(request).promise()
             const request: CodeWhispererUserClient.ListAvailableProfilesRequest = {}
@@ -187,7 +182,7 @@ export class RegionProfileManager {
         const ssoConn = this.connectionProvider() as SsoConnection
 
         // only prompt to users when users switch from A profile to B profile
-        if (source !== 'customization' && this.activeRegionProfile !== undefined && regionProfile !== undefined) {
+        if (this.activeRegionProfile !== undefined && regionProfile !== undefined) {
             const response = await showConfirmationMessage({
                 prompt: localize(
                     'AWS.amazonq.profile.confirmation',
@@ -229,16 +224,13 @@ export class RegionProfileManager {
             })
         }
 
-        await this._switchRegionProfile(regionProfile, source)
+        await this._switchRegionProfile(regionProfile)
     }
 
-    private async _switchRegionProfile(regionProfile: RegionProfile | undefined, source: ProfileSwitchIntent) {
+    private async _switchRegionProfile(regionProfile: RegionProfile | undefined) {
         this._activeRegionProfile = regionProfile
 
-        this._onDidChangeRegionProfile.fire({
-            profile: regionProfile,
-            intent: source,
-        })
+        this._onDidChangeRegionProfile.fire(regionProfile)
         // dont show if it's a default (fallback)
         if (regionProfile && this.profiles.length > 1) {
             void vscode.window.showInformationMessage(`You are using the ${regionProfile.name} profile for Q.`).then()
@@ -391,21 +383,7 @@ export class RegionProfileManager {
         }
     }
 
-    // TODO: Should maintain sdk client in a better way
-    async createQClient(profile: RegionProfile): Promise<CodeWhispererUserClient> {
-        const conn = this.connectionProvider()
-        if (conn === undefined || !isSsoConnection(conn)) {
-            throw new Error('No valid SSO connection')
-        }
-        const endpoint = endpoints.get(profile.region)
-        if (!endpoint) {
-            throw new Error(`trying to initiatize Q client with unrecognizable region ${profile.region}`)
-        }
-        return this._createQClient(profile.region, endpoint, conn)
-    }
-
-    // Visible for testing only, do not use this directly, please use createQClient(profile)
-    async _createQClient(region: string, endpoint: string, conn: SsoConnection): Promise<CodeWhispererUserClient> {
+    async createQClient(region: string, endpoint: string, conn: SsoConnection): Promise<CodeWhispererUserClient> {
         const token = (await conn.getToken()).accessToken
         const serviceOption: ServiceOptions = {
             apiConfig: userApiConfig,