fix(amazonq): Wrong customization being pushed to LSP on start

Problem

On start of the extension when we called getSelectedCustomization() it would
return the wrong customization. This caused the wrong customization to be
pushed to the LSP.

The reason was due to notifyNewCustomizations not being called before getSelectedCustomization()
was used.

Solution

Ensure the call that uses notifyNewCustomizations() runs before getSelectedCustomization is used.

I also refactored to compartmentalize the relevant code.

Signed-off-by: nkomonen-amazon <[email protected]>

Author: nkomonen-amazon

packages/amazonq/src/lsp/chat/activation.ts

@@ -11,7 +11,7 @@ import { registerLanguageServerEventListener, registerMessageListeners } from '.
 import { Commands, getLogger, globals, undefinedIfEmpty } from 'aws-core-vscode/shared'
 import { activate as registerLegacyChatListeners } from '../../app/chat/activation'
 import { DefaultAmazonQAppInitContext } from 'aws-core-vscode/amazonq'
-import { AuthUtil, getSelectedCustomization } from 'aws-core-vscode/codewhisperer'
+import { AuthUtil, getSelectedCustomization, notifyNewCustomizations } from 'aws-core-vscode/codewhisperer'
 import {
     DidChangeConfigurationNotification,
     updateConfigurationRequestType,
@@ -25,11 +25,8 @@ export async function activate(languageClient: LanguageClient, encryptionKey: Bu
         type: 'profile',
         profileArn: AuthUtil.instance.regionProfileManager.activeRegionProfile?.arn,
     })
-    // We need to push the cached customization on startup explicitly
-    await pushConfigUpdate(languageClient, {
-        type: 'customization',
-        customization: getSelectedCustomization(),
-    })
+
+    await initializeCustomizations()
 
     const provider = new AmazonQChatViewProvider(mynahUIPath)
 
@@ -86,10 +83,7 @@ export async function activate(languageClient: LanguageClient, encryptionKey: Bu
             await provider.refreshWebview()
         }),
         Commands.register('aws.amazonq.updateCustomizations', () => {
-            void pushConfigUpdate(languageClient, {
-                type: 'customization',
-                customization: undefinedIfEmpty(getSelectedCustomization().arn),
-            })
+            pushCustomizationToServer(languageClient)
         }),
         globals.logOutputChannel.onDidChangeLogLevel((logLevel) => {
             getLogger('amazonqLsp').info(`Local log level changed to ${logLevel}, notifying LSP`)
@@ -98,15 +92,48 @@ export async function activate(languageClient: LanguageClient, encryptionKey: Bu
             })
         })
     )
+
+    /**
+     * Initialize customizations on extension startup
+     */
+    async function initializeCustomizations() {
+        /**
+         * Even though this function is called "notify", it has a side effect that first restores the
+         * cached customization. So for {@link getSelectedCustomization()} to work as expected, we must
+         * call {@link notifyNewCustomizations} first.
+         *
+         * TODO: Separate restoring and notifying, or just rename the function to something better
+         */
+        if (AuthUtil.instance.isIdcConnection() && AuthUtil.instance.isConnected()) {
+            await notifyNewCustomizations()
+        }
+
+        /**
+         * HACK: We must explicitly push the customization here since restoring the customization from cache
+         * does not currently trigger a push to server.
+         *
+         * TODO: Always push to server whenever restoring from cache.
+         */
+        pushCustomizationToServer(languageClient)
+    }
+
+    function pushCustomizationToServer(languageClient: LanguageClient) {
+        void pushConfigUpdate(languageClient, {
+            type: 'customization',
+            customization: undefinedIfEmpty(getSelectedCustomization().arn),
+        })
+    }
 }
 
 /**
- * Push a config value to the language server, effectively updating it with the
+ * Request/Notify a config value to the language server, effectively updating it with the
  * latest configuration from the client.
  *
  * The issue is we need to push certain configs to different places, since there are
  * different handlers for specific configs. So this determines the correct place to
  * push the given config.
+ *
+ * TODO: Move this to somewhere more appropriate
  */
 async function pushConfigUpdate(client: LanguageClient, config: QConfigs) {
     switch (config.type) {

packages/core/src/codewhisperer/activation.ts

@@ -72,7 +72,7 @@ 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 { onProfileChangedListener } from './util/customizationUtil'
 import { CodeWhispererCommandBackend, CodeWhispererCommandDeclarations } from './commands/gettingStartedPageCommands'
 import { SecurityIssueHoverProvider } from './service/securityIssueHoverProvider'
 import { SecurityIssueCodeActionProvider } from './service/securityIssueCodeActionProvider'
@@ -354,9 +354,6 @@ export async function activate(context: ExtContext): Promise<void> {
         { emit: false, functionId: { name: 'activateCwCore' } }
     )
 
-    if (AuthUtil.instance.isIdcConnection() && AuthUtil.instance.isConnected()) {
-        await notifyNewCustomizations()
-    }
     if (AuthUtil.instance.isBuilderIdConnection()) {
         await CodeScansState.instance.setScansEnabled(false)
     }

packages/core/src/codewhisperer/index.ts

@@ -107,6 +107,7 @@ export {
     baseCustomization,
     onProfileChangedListener,
     CustomizationProvider,
+    notifyNewCustomizations,
 } from './util/customizationUtil'
 export { Container } from './service/serviceContainer'
 export * from './util/gitUtil'