feat(amazonq): support customization through language server

zixlin7 · zixlin7 · commit 785c9856cff5 · 2025-03-26T17:08:48.000-07:00
diff --git a/packages/amazonq/src/app/inline/completion.ts b/packages/amazonq/src/app/inline/completion.ts
@@ -26,9 +26,12 @@ import { SessionManager } from './sessionManager'
 import { RecommendationService } from './recommendationService'
 import {
     CodeWhispererConstants,
+    CodeWhispererSource,
+    connectCustomizationHandler,
     ReferenceHoverProvider,
     ReferenceInlineProvider,
     ReferenceLogViewProvider,
+    selectCustomizationHandler,
 } from 'aws-core-vscode/codewhisperer'
 
 export class InlineCompletionManager implements Disposable {
@@ -60,6 +63,31 @@ export class InlineCompletionManager implements Disposable {
         }
     }
 
+    public registerCustomization(client: LanguageClient) {
+        commands.registerCommand(
+            'aws.amazonq.connect',
+            async (
+                source: string,
+                startUrl?: string,
+                region?: string,
+                customizationArn?: string,
+                customizationNamePrefix?: string
+            ) => {
+                return await connectCustomizationHandler(client)(
+                    source,
+                    startUrl,
+                    region,
+                    customizationArn,
+                    customizationNamePrefix
+                )
+            }
+        )
+
+        commands.registerCommand('aws.amazonq.select', async (customization: any, source: CodeWhispererSource) => {
+            return await selectCustomizationHandler(client)(customization, source)
+        })
+    }
+
     public registerInlineCompletion() {
         const onInlineAcceptance = async (
             sessionId: string,
diff --git a/packages/amazonq/src/lsp/client.ts b/packages/amazonq/src/lsp/client.ts
@@ -14,6 +14,7 @@ import { ConnectionMetadata } from '@aws/language-server-runtimes/protocol'
 import { Settings, oidcClientName, createServerOptions, globals, Experiments, getLogger } from 'aws-core-vscode/shared'
 import { activate } from './chat/activation'
 import { AmazonQResourcePaths } from './lspInstaller'
+import { notifyNewCustomizations } from 'aws-core-vscode/codewhisperer'
 
 const localize = nls.loadMessageBundle()
 
@@ -96,9 +97,13 @@ export async function startLanguageServer(
         await auth.init()
         const inlineManager = new InlineCompletionManager(client)
         inlineManager.registerInlineCompletion()
+        inlineManager.registerCustomization(client)
         if (Experiments.instance.get('amazonqChatLSP', false)) {
             activate(client, encryptionKey, resourcePaths.mynahUI)
         }
+        if (AuthUtil.instance.isValidEnterpriseSsoInUse()) {
+            await notifyNewCustomizations(client)
+        }
 
         // Request handler for when the server wants to know about the clients auth connnection
         client.onRequest<ConnectionMetadata, Error>(notificationTypes.getConnectionMetadata.method, () => {
diff --git a/packages/amazonq/test/unit/amazonq/apps/inline/completion.test.ts b/packages/amazonq/test/unit/amazonq/apps/inline/completion.test.ts
@@ -253,7 +253,26 @@ describe('InlineCompletionManager', () => {
                 }
             })
         })
-    })
+    }),
+        describe('registerCustomization', () => {
+            beforeEach(() => {
+                manager.registerCustomization(languageClient)
+            })
+            it('should register connect customization command', async () => {
+                const connectCustomizationCall = registerCommandStub
+                    .getCalls()
+                    .find((call) => call.args[0] === 'aws.amazonq.connect')
+
+                assert(connectCustomizationCall, 'Connect customization command should be registered')
+            })
+            it('should register select customization command', async () => {
+                const connectCustomizationCall = registerCommandStub
+                    .getCalls()
+                    .find((call) => call.args[0] === 'aws.amazonq.select')
+
+                assert(connectCustomizationCall, 'Select customization command should be registered')
+            })
+        })
 
     describe('AmazonQInlineCompletionItemProvider', () => {
         describe('provideInlineCompletionItems', () => {
diff --git a/packages/core/src/codewhisperer/commands/basicCommands.ts b/packages/core/src/codewhisperer/commands/basicCommands.ts
@@ -69,6 +69,7 @@ import { UserWrittenCodeTracker } from '../tracker/userWrittenCodeTracker'
 import { parsePatch } from 'diff'
 import { createCodeIssueGroupingStrategyPrompter } from '../ui/prompters'
 import { cancel, confirm } from '../../shared/localizedText'
+import { LanguageClient } from 'vscode-languageclient'
 
 const MessageTimeOut = 5_000
 
@@ -237,15 +238,18 @@ export const showFileScan = Commands.declare(
         }
 )
 
-export const selectCustomizationPrompt = Commands.declare(
-    { id: 'aws.amazonq.selectCustomization', compositeKey: { 1: 'source' } },
-    () => async (_: VsCodeCommandArg, source: CodeWhispererSource) => {
+export const selectCustomizationHandler =
+    (client?: LanguageClient) => async (_: VsCodeCommandArg, source: CodeWhispererSource) => {
         if (isBuilderIdConnection(AuthUtil.instance.conn)) {
             throw new Error(`Select Customizations are not supported with the Amazon Builder ID connection.`)
         }
         telemetry.ui_click.emit({ elementId: 'cw_selectCustomization_Cta' })
-        void showCustomizationPrompt().then()
+        void showCustomizationPrompt(client).then()
     }
+
+export const selectCustomizationPrompt = Commands.declare(
+    { id: 'aws.amazonq.selectCustomization', compositeKey: { 1: 'source' } },
+    () => selectCustomizationHandler()
 )
 
 export const reconnect = Commands.declare(
@@ -263,6 +267,53 @@ export const showSsoSignIn = Commands.declare('aws.amazonq.sso', () => async ()
 // It can optionally set a customization too based on given values to match on
 
 // This command is only declared and registered in Amazon Q if Q exists
+
+export const connectCustomizationHandler =
+    (client?: LanguageClient) =>
+    async (
+        source: string,
+        startUrl?: string,
+        region?: string,
+        customizationArn?: string,
+        customizationNamePrefix?: string
+    ) => {
+        SsoAccessTokenProvider.authSource = source
+        if (startUrl && region) {
+            await connectToEnterpriseSso(startUrl, region)
+        } else {
+            await getStartUrl()
+        }
+
+        // No customization match information given, exit early.
+        if (!customizationArn && !customizationNamePrefix) {
+            return
+        }
+
+        let persistedCustomizations = getPersistedCustomizations()
+
+        // Check if any customizations have already been persisted.
+        // If not, call `notifyNewCustomizations` to handle it then recheck.
+        if (persistedCustomizations.length === 0) {
+            await notifyNewCustomizations(client)
+            persistedCustomizations = getPersistedCustomizations()
+        }
+
+        // If given an ARN, assume a specific customization is desired and find an entry that matches it. Ignores the prefix logic.
+        // Otherwise if only a prefix is given, find an entry that matches it.
+        // Backwards compatible with previous implementation.
+        const match = customizationArn
+            ? persistedCustomizations.find((c) => c.arn === customizationArn)
+            : persistedCustomizations.find((c) => c.name?.startsWith(customizationNamePrefix as string))
+
+        // If no match is found, nothing to do :)
+        if (!match) {
+            getLogger().error(`No customization match found: arn=${customizationArn} prefix=${customizationNamePrefix}`)
+            return
+        }
+        // Since we selected based on a match, we'll reuse the persisted values.
+        await selectCustomization(match, client)
+    }
+
 export const connectWithCustomization = Commands.declare(
     { id: 'aws.codeWhisperer.connect', compositeKey: { 0: 'source' } },
     /**
@@ -273,52 +324,7 @@ export const connectWithCustomization = Commands.declare(
      * customizationArn: select customization by ARN. If provided, `customizationNamePrefix` is ignored.
      * customizationNamePrefix: select customization by prefix, if `customizationArn` is `undefined`.
      */
-    () =>
-        async (
-            source: string,
-            startUrl?: string,
-            region?: string,
-            customizationArn?: string,
-            customizationNamePrefix?: string
-        ) => {
-            SsoAccessTokenProvider.authSource = source
-            if (startUrl && region) {
-                await connectToEnterpriseSso(startUrl, region)
-            } else {
-                await getStartUrl()
-            }
-
-            // No customization match information given, exit early.
-            if (!customizationArn && !customizationNamePrefix) {
-                return
-            }
-
-            let persistedCustomizations = getPersistedCustomizations()
-
-            // Check if any customizations have already been persisted.
-            // If not, call `notifyNewCustomizations` to handle it then recheck.
-            if (persistedCustomizations.length === 0) {
-                await notifyNewCustomizations()
-                persistedCustomizations = getPersistedCustomizations()
-            }
-
-            // If given an ARN, assume a specific customization is desired and find an entry that matches it. Ignores the prefix logic.
-            // Otherwise if only a prefix is given, find an entry that matches it.
-            // Backwards compatible with previous implementation.
-            const match = customizationArn
-                ? persistedCustomizations.find((c) => c.arn === customizationArn)
-                : persistedCustomizations.find((c) => c.name?.startsWith(customizationNamePrefix as string))
-
-            // If no match is found, nothing to do :)
-            if (!match) {
-                getLogger().error(
-                    `No customization match found: arn=${customizationArn} prefix=${customizationNamePrefix}`
-                )
-                return
-            }
-            // Since we selected based on a match, we'll reuse the persisted values.
-            await selectCustomization(match)
-        }
+    () => connectCustomizationHandler()
 )
 
 export const showLearnMore = Commands.declare(
diff --git a/packages/core/src/codewhisperer/index.ts b/packages/core/src/codewhisperer/index.ts
@@ -98,7 +98,12 @@ 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 } from './util/customizationUtil'
+export {
+    getSelectedCustomization,
+    setSelectedCustomization,
+    baseCustomization,
+    notifyNewCustomizations,
+} from './util/customizationUtil'
 export { Container } from './service/serviceContainer'
 export * from './util/gitUtil'
 export * from './ui/prompters'
diff --git a/packages/core/src/codewhisperer/ui/codeWhispererNodes.ts b/packages/core/src/codewhisperer/ui/codeWhispererNodes.ts
@@ -27,6 +27,7 @@ import { AuthUtil } from '../util/authUtil'
 import { submitFeedback } from '../../feedback/vue/submitFeedback'
 import { focusAmazonQPanel } from '../../codewhispererChat/commands/registerCommands'
 import { isWeb } from '../../shared/extensionGlobals'
+import { Experiments } from '../../shared/settings'
 
 export function createAutoSuggestions(running: boolean): DataQuickPickItem<'autoSuggestions'> {
     const labelResume = localize('AWS.codewhisperer.resumeCodeWhispererNode.label', 'Resume Auto-Suggestions')
@@ -129,11 +130,14 @@ export function createSelectCustomization(): DataQuickPickItem<'selectCustomizat
     const description =
         newCustomizationsAmount > 0 ? `${newCustomizationsAmount} new available` : `Using ${selectedCustomization.name}`
 
+    const selectCustomizationCommand = Experiments.instance.get('amazonqLSP', true)
+        ? () => vscode.commands.executeCommand('aws.amazonq.select')
+        : () => selectCustomizationPrompt.execute(placeholder, cwQuickPickSource)
     return {
         data: 'selectCustomization',
         label: codicon`${icon} ${label}`,
         description: description,
-        onClick: () => selectCustomizationPrompt.execute(placeholder, cwQuickPickSource),
+        onClick: selectCustomizationCommand,
     } as DataQuickPickItem<'selectCustomization'>
 }
 
diff --git a/packages/core/src/codewhisperer/util/customizationUtil.ts b/packages/core/src/codewhisperer/util/customizationUtil.ts
