feat(codewhisperer): set customization on connect

Problem:
Related to aws-toolkit-vscode-staging#1058, there's a need for enterprises to establish a
default customization for onboarding developers
based on their need or role.

Solution:
Provide the ability to define the customization ARN to use after connecting
to CodeWhisperer through the `aws.codeWhisperer.connect` command,
along side other customization details.

Author: Will Lovo

docs/api.md

@@ -8,9 +8,11 @@ Details about any publicly accessible functionalities exposed through [extension
 
 #### `aws.codeWhisperer.connect`
 
-**Signature**: _async (startUrl?: string, region?: string) => Promise<void>_
+**Signature**: _async (startUrl?: string, region?: string, customizationArn?: string, customizationName?: string, customizationDescription?: string) => Promise<void>_
 
-Shortcut command to directly connect to Identity Center or prompt start URL entry, as well as set a customization for CodeWhisperer requests. Customization is not yet supported.
+Shortcut command to directly connect to Identity Center or prompt start URL entry, as well as set a customization for CodeWhisperer requests.
 
-This command expects two arguments: startUrl and region (both strings).
-If these arguments are provided, they will be used. Otherwise, the commands prompts for them interactively.
+This command supports two sets of arguments:
+
+-   startUrl and region. If both arguments are provided they will be used, otherwise the command prompts for them interactively.
+-   customization{Arn, Name, Description}. If at least customizationArn is provided, the command selects this customization.

src/codewhisperer/activation.ts

@@ -33,9 +33,9 @@ import {
     showIntroduction,
     reconnect,
     refreshStatusBar,
-    selectCustomization,
+    selectCustomizationPrompt,
     notifyNewCustomizationsCmd,
-    connectIdCenter,
+    connectWithCustomization,
 } from './commands/basicCommands'
 import { sleep } from '../shared/utilities/timeoutUtils'
 import { ReferenceLogViewProvider } from './service/referenceLogViewProvider'
@@ -159,8 +159,8 @@ export async function activate(context: ExtContext): Promise<void> {
         }),
         // show introduction
         showIntroduction.register(),
-        // CodeWhisperer Identity Center connection setup
-        connectIdCenter.register(),
+        // direct CodeWhisperer connection setup with customization
+        connectWithCustomization.register(),
         // toggle code suggestions
         toggleCodeSuggestions.register(context.extensionContext.globalState),
         // enable code suggestions
@@ -184,7 +184,7 @@ export async function activate(context: ExtContext): Promise<void> {
             invokeRecommendation(vscode.window.activeTextEditor as vscode.TextEditor, client, await getConfigEntry())
         }),
         // select customization
-        selectCustomization.register(),
+        selectCustomizationPrompt.register(),
         // notify new customizations
         notifyNewCustomizationsCmd.register(),
         /**

src/codewhisperer/commands/basicCommands.ts

@@ -19,7 +19,12 @@ import { AuthUtil } from '../util/authUtil'
 import { isCloud9 } from '../../shared/extensionUtilities'
 import { InlineCompletionService } from '../service/inlineCompletionService'
 import { openUrl } from '../../shared/utilities/vsCodeUtils'
-import { notifyNewCustomizations, showCustomizationPrompt } from '../util/customizationUtil'
+import {
+    getPersistedCustomizations,
+    notifyNewCustomizations,
+    selectCustomization,
+    showCustomizationPrompt,
+} from '../util/customizationUtil'
 import { get, set } from '../util/commonUtil'
 import { CodeWhispererCommandDeclarations } from '../commands/gettingStartedPageCommands'
 import { getIcon } from '../../shared/icons'
@@ -40,7 +45,7 @@ export const toggleCodeSuggestions = Commands.declare(
         })
     }
 )
-/* 
+/*
 createGettingStartedNode(Learn) will be a childnode of CodeWhisperer
 onClick on this "Learn" Node will open the Learn CodeWhisperer Page.
 */
@@ -100,7 +105,7 @@ export const showSecurityScan = Commands.declare(
         }
 )
 
-export const selectCustomization = Commands.declare('aws.codeWhisperer.selectCustomization', () => async () => {
+export const selectCustomizationPrompt = Commands.declare('aws.codeWhisperer.selectCustomization', () => async () => {
     telemetry.ui_click.emit({ elementId: 'cw_selectCustomization_Cta' })
     showCustomizationPrompt().then()
 })
@@ -115,19 +120,37 @@ export const showSsoSignIn = Commands.declare('aws.codeWhisperer.sso', () => asy
 })
 
 // Shortcut command to directly connect to Identity Center or prompt start URL entry
-// Customization is not yet supported.
-export const connectIdCenter = Commands.declare(
+// It can optionally set a customization too.
+export const connectWithCustomization = Commands.declare(
     'aws.codeWhisperer.connect',
-    () => async (startUrl?: string, region?: string) => {
-        // This command expects two arguments: startUrl and region (both strings).
-        // If these arguments are provided, they will be used.
-        // Otherwise, the commands prompts for them interactively.
-        if (startUrl && region) {
-            await connectToEnterpriseSso(startUrl, region)
-        } else {
-            await getStartUrl()
+    () =>
+        async (
+            startUrl?: string,
+            region?: string,
+            customizationArn?: string,
+            customizationName?: string,
+            customizationDescription?: string
+        ) => {
+            // This command supports two sets of arguments:
+            //  * startUrl and region. If both arguments are provided they will be used, otherwise
+            //    the command prompts for them interactively.
+            //  * customization{Arn, Name, Description}. If at least customizationArn is provided,
+            //    the command selects this customization.
+            if (startUrl && region) {
+                await connectToEnterpriseSso(startUrl, region)
+            } else {
+                await getStartUrl()
+            }
+            if (customizationArn) {
+                const match = getPersistedCustomizations().find(c => c.arn == customizationArn)
+                const customization = {
+                    arn: customizationArn,
+                    name: customizationName ?? match?.name ?? 'unknown',
+                    description: customizationDescription ?? match?.description ?? 'unknown',
+                }
+                await selectCustomization(customization)
+            }
         }
-    }
 )
 
 export const showLearnMore = Commands.declare('aws.codeWhisperer.learnMore', () => async () => {

src/codewhisperer/explorer/codewhispererChildrenNodes.ts

@@ -14,7 +14,7 @@ import {
     showLearnMore,
     showFreeTierLimit,
     reconnect,
-    selectCustomization,
+    selectCustomizationPrompt,
 } from '../commands/basicCommands'
 import { codeScanState } from '../models/model'
 import { getNewCustomizationAvailable, getSelectedCustomization } from '../util/customizationUtil'
@@ -96,7 +96,7 @@ export const createSelectCustomizationNode = () => {
     const selectedCustomization = getSelectedCustomization()
     const newText = newCustomizationsAvailable ? 'new!      ' : ''
 
-    return selectCustomization.build().asTreeNode({
+    return selectCustomizationPrompt.build().asTreeNode({
         label: localize('AWS.explorerNode.selectCustomization.label', 'Select Customization'),
         iconPath: getIcon('vscode-extensions'),
         description: `${newText}${selectedCustomization.arn === '' ? '' : selectedCustomization.name}`,

src/codewhisperer/util/customizationUtil.ts

@@ -279,7 +279,7 @@ const createCustomizationItem = (
     } as DataQuickPickItem<string>
 }
 
-const selectCustomization = async (customization: Customization) => {
+export const selectCustomization = async (customization: Customization) => {
     // If the newly selected customization is same as the old one, do nothing
     const selectedCustomization = getSelectedCustomization()
     if (selectedCustomization.arn === customization.arn) {