Merge pull request aws#7302 from nkomonen-amazon/customizationFix

opieter-aws · web-flow · commit f849b3feb533 · 2025-05-14T08:55:03.000-04:00
fix(amazonq): Customization/Profiles not being sent to LSP on startup
diff --git a/packages/amazonq/src/lsp/chat/activation.ts b/packages/amazonq/src/lsp/chat/activation.ts
@@ -11,11 +11,8 @@ 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 {
-    DidChangeConfigurationNotification,
-    updateConfigurationRequestType,
-} from '@aws/language-server-runtimes/protocol'
+import { AuthUtil, getSelectedCustomization, notifyNewCustomizations } from 'aws-core-vscode/codewhisperer'
+import { pushConfigUpdate } from '../config'
 
 export async function activate(languageClient: LanguageClient, encryptionKey: Buffer, mynahUIPath: string) {
     const disposables = globals.context.subscriptions
@@ -25,11 +22,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)
 
@@ -79,17 +73,10 @@ export async function activate(languageClient: LanguageClient, encryptionKey: Bu
 
     disposables.push(
         AuthUtil.instance.regionProfileManager.onDidChangeRegionProfile(async () => {
-            void pushConfigUpdate(languageClient, {
-                type: 'profile',
-                profileArn: AuthUtil.instance.regionProfileManager.activeRegionProfile?.arn,
-            })
             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,46 +85,35 @@ export async function activate(languageClient: LanguageClient, encryptionKey: Bu
             })
         })
     )
-}
 
-/**
- * Push 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.
- */
-async function pushConfigUpdate(client: LanguageClient, config: QConfigs) {
-    switch (config.type) {
-        case 'profile':
-            await client.sendRequest(updateConfigurationRequestType.method, {
-                section: 'aws.q',
-                settings: { profileArn: config.profileArn },
-            })
-            break
-        case 'customization':
-            client.sendNotification(DidChangeConfigurationNotification.type.method, {
-                section: 'aws.q',
-                settings: { customization: config.customization },
-            })
-            break
-        case 'logLevel':
-            client.sendNotification(DidChangeConfigurationNotification.type.method, {
-                section: 'aws.logLevel',
-            })
-            break
+    /**
+     * 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),
+        })
     }
 }
-type ProfileConfig = {
-    type: 'profile'
-    profileArn: string | undefined
-}
-type CustomizationConfig = {
-    type: 'customization'
-    customization: string | undefined
-}
-type LogLevelConfig = {
-    type: 'logLevel'
-}
-type QConfigs = ProfileConfig | CustomizationConfig | LogLevelConfig
diff --git a/packages/amazonq/src/lsp/client.ts b/packages/amazonq/src/lsp/client.ts
@@ -54,7 +54,7 @@ import { processUtils } from 'aws-core-vscode/shared'
 import { activate as activateChat } from './chat/activation'
 import { AmazonQResourcePaths } from './lspInstaller'
 import { auth2 } from 'aws-core-vscode/auth'
-import { ConfigSection, isValidConfigSection, toAmazonQLSPLogLevel } from './config'
+import { ConfigSection, isValidConfigSection, pushConfigUpdate, toAmazonQLSPLogLevel } from './config'
 import { telemetry } from 'aws-core-vscode/telemetry'
 
 const localize = nls.loadMessageBundle()
@@ -179,15 +179,31 @@ export async function startLanguageServer(
 
     await client.onReady()
 
-    // IMPORTANT: This sets up Auth and must be called before anything attempts to use it
-    AuthUtil.create(new auth2.LanguageClientAuth(client, clientId, encryptionKey))
-    // Ideally this would be part of AuthUtil.create() as it restores the existing Auth connection, but we have some
-    // future work which will need to delay this call
-    await AuthUtil.instance.restore()
+    /**
+     * We use the Flare Auth language server, and our Auth client depends on it.
+     * Because of this we initialize our Auth client **immediately** after the language server is ready.
+     * Doing this removes the chance of something else attempting to use the Auth client before it is ready.
+     */
+    await initializeAuth(client)
 
     await postStartLanguageServer(client, resourcePaths, toDispose)
 
     return client
+
+    async function initializeAuth(client: LanguageClient) {
+        AuthUtil.create(new auth2.LanguageClientAuth(client, clientId, encryptionKey))
+
+        /** All must be setup before {@link AuthUtil.restore} otherwise they may not trigger when expected */
+        AuthUtil.instance.regionProfileManager.onDidChangeRegionProfile(async () => {
+            void pushConfigUpdate(client, {
+                type: 'profile',
+                profileArn: AuthUtil.instance.regionProfileManager.activeRegionProfile?.arn,
+            })
+        })
+
+        // Try and restore a cached connection if exists
+        await AuthUtil.instance.restore()
+    }
 }
 
 async function postStartLanguageServer(
diff --git a/packages/amazonq/src/lsp/config.ts b/packages/amazonq/src/lsp/config.ts
@@ -5,6 +5,11 @@
 import * as vscode from 'vscode'
 import { DevSettings, getServiceEnvVarConfig } from 'aws-core-vscode/shared'
 import { LspConfig } from 'aws-core-vscode/amazonq'
+import { LanguageClient } from 'vscode-languageclient'
+import {
+    DidChangeConfigurationNotification,
+    updateConfigurationRequestType,
+} from '@aws/language-server-runtimes/protocol'
 
 export interface ExtendedAmazonQLSPConfig extends LspConfig {
     ui?: string
@@ -54,3 +59,45 @@ export function getAmazonQLspConfig(): ExtendedAmazonQLSPConfig {
 export function toAmazonQLSPLogLevel(logLevel: vscode.LogLevel): LspLogLevel {
     return lspLogLevelMapping.get(logLevel) ?? 'info'
 }
+
+/**
+ * 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.
+ */
+export async function pushConfigUpdate(client: LanguageClient, config: QConfigs) {
+    switch (config.type) {
+        case 'profile':
+            await client.sendRequest(updateConfigurationRequestType.method, {
+                section: 'aws.q',
+                settings: { profileArn: config.profileArn },
+            })
+            break
+        case 'customization':
+            client.sendNotification(DidChangeConfigurationNotification.type.method, {
+                section: 'aws.q',
+                settings: { customization: config.customization },
+            })
+            break
+        case 'logLevel':
+            client.sendNotification(DidChangeConfigurationNotification.type.method, {
+                section: 'aws.logLevel',
+            })
+            break
+    }
+}
+type ProfileConfig = {
+    type: 'profile'
+    profileArn: string | undefined
+}
+type CustomizationConfig = {
+    type: 'customization'
+    customization: string | undefined
+}
+type LogLevelConfig = {
+    type: 'logLevel'
+}
+type QConfigs = ProfileConfig | CustomizationConfig | LogLevelConfig
diff --git a/packages/core/src/codewhisperer/activation.ts b/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)
     }
diff --git a/packages/core/src/codewhisperer/index.ts b/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'
