Merge pull request #1075 from Will-ShaoHua/feature/vector

Author: Will-ShaoHua

docs/api.md

@@ -0,0 +1,18 @@
+# AWS Toolkit API
+
+Details about any publicly accessible functionalities exposed through [extension commands](https://code.visualstudio.com/api/references/vscode-api#commands) or [exported APIs](https://code.visualstudio.com/api/references/vscode-api#extensions).
+
+## Pseudo (Internal-only) API
+
+### Commands
+
+#### `aws.codeWhisperer.connect`
+
+**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.
+
+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.

package.nls.json

@@ -216,5 +216,14 @@
     "AWS.walkthrough.gettingStarted.connect": "Connect to AWS",
     "AWS.walkthrough.gettingStarted.changeRegions": "Change AWS Regions",
     "AWS.walkthrough.gettingStarted.setupToolchain": "Configure your toolchain",
-    "AWS.command.codewhisperer.title": "CodeWhisperer Invoke Service"
+    "AWS.command.codewhisperer.title": "CodeWhisperer Invoke Service",
+    "AWS.explorerNode.selectCustomization.label": "Select a Customization",
+    "AWS.codewhisperer.customization.base.label": "CodeWhisperer foundation (Default)",
+    "AWS.codewhisperer.customization.base.description": "default",
+    "AWS.codewhisperer.customization.base.detail": "Receive suggestions from CodeWhisperer base model",
+    "AWS.codewhisperer.customization.selected": "Connected",
+    "AWS.codewhisperer.customization.quickPick.title": "Select a Customization",
+    "AWS.codewhisperer.customization.quickPick.placeholder": "You have access to the following customizations",
+    "AWS.codewhisperer.customization.notification.new_customizations.select": "Select Customization",
+    "AWS.codewhisperer.customization.notification.new_customizations.learn_more": "Learn More"
 }

src/codewhisperer/activation.ts

@@ -33,6 +33,9 @@ import {
     showIntroduction,
     reconnect,
     refreshStatusBar,
+    selectCustomizationPrompt,
+    notifyNewCustomizationsCmd,
+    connectWithCustomization,
 } from './commands/basicCommands'
 import { sleep } from '../shared/utilities/timeoutUtils'
 import { ReferenceLogViewProvider } from './service/referenceLogViewProvider'
@@ -49,6 +52,7 @@ import { AuthUtil } from './util/authUtil'
 import { ImportAdderProvider } from './service/importAdderProvider'
 import { TelemetryHelper } from './util/telemetryHelper'
 import { openUrl } from '../shared/utilities/vsCodeUtils'
+import { notifyNewCustomizations } from './util/customizationUtil'
 import { CodeWhispererCommandBackend, CodeWhispererCommandDeclarations } from './commands/gettingStartedPageCommands'
 const performance = globalThis.performance ?? require('perf_hooks').performance
 
@@ -155,6 +159,8 @@ export async function activate(context: ExtContext): Promise<void> {
         }),
         // show introduction
         showIntroduction.register(),
+        // direct CodeWhisperer connection setup with customization
+        connectWithCustomization.register(),
         // toggle code suggestions
         toggleCodeSuggestions.register(context.extensionContext.globalState),
         // enable code suggestions
@@ -177,6 +183,10 @@ export async function activate(context: ExtContext): Promise<void> {
         Commands.register({ id: 'aws.codeWhisperer', autoconnect: true }, async () => {
             invokeRecommendation(vscode.window.activeTextEditor as vscode.TextEditor, client, await getConfigEntry())
         }),
+        // select customization
+        selectCustomizationPrompt.register(),
+        // notify new customizations
+        notifyNewCustomizationsCmd.register(),
         /**
          * On recommendation acceptance
          */
@@ -210,6 +220,9 @@ export async function activate(context: ExtContext): Promise<void> {
     if (auth.isConnectionExpired()) {
         auth.showReauthenticatePrompt()
     }
+    if (auth.isValidEnterpriseSsoInUse()) {
+        await notifyNewCustomizations()
+    }
 
     function activateSecurityScan() {
         context.extensionContext.subscriptions.push(

src/codewhisperer/client/codewhisperer.ts

@@ -2,22 +2,22 @@
  * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
-import { AWSError, Service } from 'aws-sdk'
-import apiConfig = require('./service-2.json')
-import userApiConfig = require('./user-service-2.json')
+import { AWSError, Credentials, Service } from 'aws-sdk'
 import globals from '../../shared/extensionGlobals'
 import * as CodeWhispererClient from './codewhispererclient'
 import * as CodeWhispererUserClient from './codewhispereruserclient'
+import { ListAvailableCustomizationsResponse, SendTelemetryEventRequest } from './codewhispereruserclient'
 import * as CodeWhispererConstants from '../models/constants'
 import { ServiceOptions } from '../../shared/awsClientBuilder'
 import { isCloud9 } from '../../shared/extensionUtilities'
 import { CodeWhispererSettings } from '../util/codewhispererSettings'
 import { PromiseResult } from 'aws-sdk/lib/request'
-import { Credentials } from 'aws-sdk'
 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 { SendTelemetryEventRequest } from './codewhispereruserclient'
 import { getLogger } from '../../shared/logger'
 
 export type ProgrammingLanguage = Readonly<
@@ -192,6 +192,13 @@ export class DefaultCodeWhispererClient {
             .promise()
     }
 
+    public async listAvailableCustomizations(): Promise<ListAvailableCustomizationsResponse[]> {
+        const client = await this.createUserSdkClient()
+        const requester = async (request: CodeWhispererUserClient.ListAvailableCustomizationsRequest) =>
+            client.listAvailableCustomizations(request).promise()
+        return pageableToCollection(requester, {}, 'nextToken').promise()
+    }
+
     public async sendTelemetryEvent(request: SendTelemetryEventRequest) {
         const requestWithOptOut: SendTelemetryEventRequest = {
             ...request,

src/codewhisperer/client/user-service-2.json

@@ -75,6 +75,21 @@
                 { "shape": "AccessDeniedException" }
             ]
         },
+        "ListAvailableCustomizations": {
+            "name": "ListAvailableCustomizations",
+            "http": {
+                "method": "POST",
+                "requestUri": "/"
+            },
+            "input": { "shape": "ListAvailableCustomizationsRequest" },
+            "output": { "shape": "ListAvailableCustomizationsResponse" },
+            "errors": [
+                { "shape": "ThrottlingException" },
+                { "shape": "InternalServerException" },
+                { "shape": "ValidationException" },
+                { "shape": "AccessDeniedException" }
+            ]
+        },
         "ListCodeAnalysisFindings": {
             "name": "ListCodeAnalysisFindings",
             "http": {
@@ -164,6 +179,7 @@
             "type": "structure",
             "required": ["programmingLanguage", "acceptedCharacterCount", "totalCharacterCount", "timestamp"],
             "members": {
+                "customizationArn": { "shape": "ResourceArn" },
                 "programmingLanguage": { "shape": "ProgrammingLanguage" },
                 "acceptedCharacterCount": { "shape": "SensitiveInteger" },
                 "totalCharacterCount": { "shape": "SensitiveInteger" },
@@ -240,6 +256,34 @@
                 "kmsKeyArn": { "shape": "ResourceArn" }
             }
         },
+        "Customization": {
+            "type": "structure",
+            "required": ["arn"],
+            "members": {
+                "arn": { "shape": "ResourceArn" },
+                "name": { "shape": "CustomizationName" },
+                "description": { "shape": "Description" }
+            }
+        },
+        "CustomizationName": {
+            "type": "string",
+            "max": 100,
+            "min": 1,
+            "pattern": "[a-zA-Z][a-zA-Z0-9_-]*"
+        },
+        "Customizations": {
+            "type": "list",
+            "member": { "shape": "Customization" }
+        },
+        "Description": {
+            "type": "string",
+            "max": 256,
+            "min": 0
+        },
+        "Double": {
+            "type": "double",
+            "box": true
+        },
         "FileContext": {
             "type": "structure",
             "required": ["leftFileContent", "rightFileContent", "filename", "programmingLanguage"],
@@ -276,7 +320,8 @@
                 "maxResults": { "shape": "GenerateCompletionsRequestMaxResultsInteger" },
                 "nextToken": { "shape": "GenerateCompletionsRequestNextTokenString" },
                 "referenceTrackerConfiguration": { "shape": "ReferenceTrackerConfiguration" },
-                "supplementalContexts": { "shape": "SupplementalContextList" }
+                "supplementalContexts": { "shape": "SupplementalContextList" },
+                "customizationArn": { "shape": "ResourceArn" }
             }
         },
         "GenerateCompletionsRequestMaxResultsInteger": {
@@ -351,6 +396,27 @@
             "fault": true,
             "retryable": { "throttling": false }
         },
+        "ListAvailableCustomizationsRequest": {
+            "type": "structure",
+            "members": {
+                "maxResults": { "shape": "ListAvailableCustomizationsRequestMaxResultsInteger" },
+                "nextToken": { "shape": "Base64EncodedPaginationToken" }
+            }
+        },
+        "ListAvailableCustomizationsRequestMaxResultsInteger": {
+            "type": "integer",
+            "box": true,
+            "max": 100,
+            "min": 1
+        },
+        "ListAvailableCustomizationsResponse": {
+            "type": "structure",
+            "required": ["customizations"],
+            "members": {
+                "customizations": { "shape": "Customizations" },
+                "nextToken": { "shape": "Base64EncodedPaginationToken" }
+            }
+        },
         "ListCodeAnalysisFindingsRequest": {
             "type": "structure",
             "required": ["jobId", "codeAnalysisFindingsSchema"],
@@ -598,6 +664,7 @@
                 "requestId": { "shape": "UUID" },
                 "programmingLanguage": { "shape": "ProgrammingLanguage" },
                 "modificationPercentage": { "shape": "SensitiveDouble" },
+                "customizationArn": { "shape": "ResourceArn" },
                 "timestamp": { "shape": "Timestamp" }
             }
         },
@@ -615,6 +682,7 @@
             "members": {
                 "sessionId": { "shape": "UUID" },
                 "requestId": { "shape": "UUID" },
+                "customizationArn": { "shape": "ResourceArn" },
                 "programmingLanguage": { "shape": "ProgrammingLanguage" },
                 "completionType": { "shape": "CompletionType" },
                 "suggestionState": { "shape": "SuggestionState" },

src/codewhisperer/commands/basicCommands.ts

@@ -8,17 +8,24 @@ import { telemetry } from '../../shared/telemetry/telemetry'
 import { ExtContext } from '../../shared/extensions'
 import { Commands } from '../../shared/vscode/commands2'
 import * as CodeWhispererConstants from '../models/constants'
-import { getLogger } from '../../shared/logger'
 import { DefaultCodeWhispererClient } from '../client/codewhisperer'
 import { startSecurityScanWithProgress, confirmStopSecurityScan } from './startSecurityScan'
 import { SecurityPanelViewProvider } from '../views/securityPanelViewProvider'
 import { codeScanState } from '../models/model'
+import { connectToEnterpriseSso, getStartUrl } from '../util/getStartUrl'
 import { showConnectionPrompt } from '../util/showSsoPrompt'
 import { ReferenceLogViewProvider } from '../service/referenceLogViewProvider'
 import { AuthUtil } from '../util/authUtil'
 import { isCloud9 } from '../../shared/extensionUtilities'
 import { InlineCompletionService } from '../service/inlineCompletionService'
 import { openUrl } from '../../shared/utilities/vsCodeUtils'
+import {
+    getPersistedCustomizations,
+    notifyNewCustomizations,
+    selectCustomization,
+    showCustomizationPrompt,
+} from '../util/customizationUtil'
+import { get, set } from '../util/commonUtil'
 import { CodeWhispererCommandDeclarations } from '../commands/gettingStartedPageCommands'
 import { getIcon } from '../../shared/icons'
 import { localize } from '../../shared/utilities/vsCodeUtils'
@@ -38,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.
 */
@@ -98,28 +105,54 @@ export const showSecurityScan = Commands.declare(
         }
 )
 
+export const selectCustomizationPrompt = Commands.declare('aws.codeWhisperer.selectCustomization', () => async () => {
+    telemetry.ui_click.emit({ elementId: 'cw_selectCustomization_Cta' })
+    showCustomizationPrompt().then()
+})
+
 export const reconnect = Commands.declare('aws.codeWhisperer.reconnect', () => async () => {
     await AuthUtil.instance.reauthenticate()
 })
 
-export function get(key: string, context: vscode.Memento): any {
-    return context.get(key)
-}
-
-export async function set(key: string, value: any, context: vscode.Memento): Promise<void> {
-    await context.update(key, value).then(
-        () => {},
-        error => {
-            getLogger().verbose(`Failed to update global state: ${error}`)
-        }
-    )
-}
-
 export const showSsoSignIn = Commands.declare('aws.codeWhisperer.sso', () => async () => {
     telemetry.ui_click.emit({ elementId: 'cw_signUp_Cta' })
     await showConnectionPrompt()
 })
 
+// Shortcut command to directly connect to Identity Center or prompt start URL entry
+// It can optionally set a customization too.
+export const connectWithCustomization = Commands.declare(
+    'aws.codeWhisperer.connect',
+    () =>
+        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 () => {
     telemetry.ui_click.emit({ elementId: 'cw_learnMore_Cta' })
     openUrl(vscode.Uri.parse(CodeWhispererConstants.learnMoreUriGeneral))
@@ -149,3 +182,10 @@ export const refreshStatusBar = Commands.declare(
         }
     }
 )
+
+export const notifyNewCustomizationsCmd = Commands.declare(
+    { id: 'aws.codeWhisperer.notifyNewCustomizations', logging: false },
+    () => () => {
+        notifyNewCustomizations().then()
+    }
+)

src/codewhisperer/explorer/codewhispererChildrenNodes.ts

@@ -14,8 +14,10 @@ import {
     showLearnMore,
     showFreeTierLimit,
     reconnect,
+    selectCustomizationPrompt,
 } from '../commands/basicCommands'
 import { codeScanState } from '../models/model'
+import { getNewCustomizationAvailable, getSelectedCustomization } from '../util/customizationUtil'
 
 export const createEnableCodeSuggestionsNode = () =>
     enableCodeSuggestions.build().asTreeNode({
@@ -88,3 +90,15 @@ export const createFreeTierLimitMetNode = () => {
         description: localize('AWS.explorerNode.freeTierLimitMet.tooltip', `paused until ${nextMonth}`),
     })
 }
+
+export const createSelectCustomizationNode = () => {
+    const newCustomizationsAvailable = getNewCustomizationAvailable()
+    const selectedCustomization = getSelectedCustomization()
+    const newText = newCustomizationsAvailable ? 'new!      ' : ''
+
+    return selectCustomizationPrompt.build().asTreeNode({
+        label: localize('AWS.explorerNode.selectCustomization.label', 'Select Customization'),
+        iconPath: getIcon('vscode-extensions'),
+        description: `${newText}${selectedCustomization.arn === '' ? '' : selectedCustomization.name}`,
+    })
+}