fix: update Q profile and customizations on language-servers crash restart

Author: floralph

packages/amazonq/src/lsp/client.ts

@@ -254,6 +254,59 @@ async function initializeAuth(client: LanguageClient): Promise<AmazonQLspAuth> {
     return auth
 }
 
+// jscpd:ignore-start
+async function initializeLanguageServerConfiguration(client: LanguageClient, context: string = 'startup') {
+    const logger = getLogger('amazonqLsp')
+
+    if (AuthUtil.instance.isConnectionValid()) {
+        logger.info(`[${context}] Initializing language server configuration`)
+// jscpd:ignore-end
+
+        try {
+            // Send profile configuration
+            logger.debug(`[${context}] Sending profile configuration to language server`)
+            await sendProfileToLsp(client)
+            logger.debug(`[${context}] Profile configuration sent successfully`)
+
+            // Send customization configuration
+            logger.debug(`[${context}] Sending customization configuration to language server`)
+            await pushConfigUpdate(client, {
+                type: 'customization',
+                customization: getSelectedCustomization(),
+            })
+            logger.debug(`[${context}] Customization configuration sent successfully`)
+
+            logger.info(`[${context}] Language server configuration completed successfully`)
+        } catch (error) {
+            logger.error(`[${context}] Failed to initialize language server configuration: ${error}`)
+            throw error
+        }
+    } else {
+        logger.warn(
+            `[${context}] Connection invalid, skipping language server configuration - this will cause authentication failures`
+        )
+        const activeConnection = AuthUtil.instance.auth.activeConnection
+        const connectionState = activeConnection
+            ? AuthUtil.instance.auth.getConnectionState(activeConnection)
+            : 'no-connection'
+        logger.warn(`[${context}] Connection state: ${connectionState}`)
+    }
+}
+
+async function sendProfileToLsp(client: LanguageClient) {
+    const logger = getLogger('amazonqLsp')
+    const profileArn = AuthUtil.instance.regionProfileManager.activeRegionProfile?.arn
+
+    logger.debug(`Sending profile to LSP: ${profileArn || 'undefined'}`)
+
+    await pushConfigUpdate(client, {
+        type: 'profile',
+        profileArn: profileArn,
+    })
+
+    logger.debug(`Profile sent to LSP successfully`)
+}
+
 async function onLanguageServerReady(
     extensionContext: vscode.ExtensionContext,
     auth: AmazonQLspAuth,
@@ -296,14 +349,7 @@ async function onLanguageServerReady(
     // We manually push the cached values the first time since event handlers, which should push, may not have been setup yet.
     // Execution order is weird and should be fixed in the flare implementation.
     // TODO: Revisit if we need this if we setup the event handlers properly
-    if (AuthUtil.instance.isConnectionValid()) {
-        await sendProfileToLsp(client)
-
-        await pushConfigUpdate(client, {
-            type: 'customization',
-            customization: getSelectedCustomization(),
-        })
-    }
+    await initializeLanguageServerConfiguration(client, 'startup')
 
     toDispose.push(
         inlineManager,
@@ -405,13 +451,6 @@ async function onLanguageServerReady(
         // Set this inside onReady so that it only triggers on subsequent language server starts (not the first)
         onServerRestartHandler(client, auth)
     )
-
-    async function sendProfileToLsp(client: LanguageClient) {
-        await pushConfigUpdate(client, {
-            type: 'profile',
-            profileArn: AuthUtil.instance.regionProfileManager.activeRegionProfile?.arn,
-        })
-    }
 }
 
 /**
@@ -431,8 +470,21 @@ function onServerRestartHandler(client: LanguageClient, auth: AmazonQLspAuth) {
         // TODO: Port this metric override to common definitions
         telemetry.languageServer_crash.emit({ id: 'AmazonQ' })
 
-        // Need to set the auth token in the again
-        await auth.refreshConnection(true)
+        const logger = getLogger('amazonqLsp')
+        logger.info('[crash-recovery] Language server crash detected, reinitializing authentication')
+
+        try {
+            // Send bearer token
+            logger.debug('[crash-recovery] Refreshing connection and sending bearer token')
+            await auth.refreshConnection(true)
+            logger.debug('[crash-recovery] Bearer token sent successfully')
+
+            // Send profile and customization configuration
+            await initializeLanguageServerConfiguration(client, 'crash-recovery')
+            logger.info('[crash-recovery] Authentication reinitialized successfully')
+        } catch (error) {
+            logger.error(`[crash-recovery] Failed to reinitialize after crash: ${error}`)
+        }
     })
 }
 

packages/amazonq/src/lsp/config.ts

@@ -3,7 +3,7 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 import * as vscode from 'vscode'
-import { DevSettings, getServiceEnvVarConfig, BaseLspInstaller } from 'aws-core-vscode/shared'
+import { DevSettings, getServiceEnvVarConfig, BaseLspInstaller, getLogger } from 'aws-core-vscode/shared'
 import { LanguageClient } from 'vscode-languageclient'
 import {
     DidChangeConfigurationNotification,
@@ -68,23 +68,31 @@ export function toAmazonQLSPLogLevel(logLevel: vscode.LogLevel): LspLogLevel {
  * push the given config.
  */
 export async function pushConfigUpdate(client: LanguageClient, config: QConfigs) {
+    const logger = getLogger('amazonqLsp')
+
     switch (config.type) {
         case 'profile':
+            logger.debug(`Pushing profile configuration: ${config.profileArn || 'undefined'}`)
             await client.sendRequest(updateConfigurationRequestType.method, {
                 section: 'aws.q',
                 settings: { profileArn: config.profileArn },
             })
+            logger.debug(`Profile configuration pushed successfully`)
             break
         case 'customization':
+            logger.debug(`Pushing customization configuration: ${config.customization || 'undefined'}`)
             client.sendNotification(DidChangeConfigurationNotification.type.method, {
                 section: 'aws.q',
                 settings: { customization: config.customization },
             })
+            logger.debug(`Customization configuration pushed successfully`)
             break
         case 'logLevel':
+            logger.debug(`Pushing log level configuration`)
             client.sendNotification(DidChangeConfigurationNotification.type.method, {
                 section: 'aws.logLevel',
             })
+            logger.debug(`Log level configuration pushed successfully`)
             break
     }
 }