fix(amazonq): Profile needing to be selected on server restart

packages/amazonq/.changes/next-release/Bug Fix-2471c584-3904-4d90-bb7c-61efff219e43.json

@@ -0,0 +1,4 @@
+{
+	"type": "Bug Fix",
+	"description": "Fix Error: 'Amazon Q service is not signed in'"
+}

packages/amazonq/.changes/next-release/Bug Fix-91d391d4-3777-4053-9e71-15b36dfa1f67.json

@@ -0,0 +1,4 @@
+{
+	"type": "Bug Fix",
+	"description": "Fix Error: 'Amazon Q Profile is not selected for IDC connection type'"
+}

packages/amazonq/src/lsp/auth.ts

@@ -18,6 +18,7 @@ import { AuthUtil } from 'aws-core-vscode/codewhisperer'
 import { Writable } from 'stream'
 import { onceChanged } from 'aws-core-vscode/utils'
 import { getLogger, oneMinute } from 'aws-core-vscode/shared'
+import { isSsoConnection } from 'aws-core-vscode/auth'
 
 export const encryptionKey = crypto.randomBytes(32)
 
@@ -76,8 +77,8 @@ export class AmazonQLspAuth {
      * @param force bypass memoization, and forcefully update the bearer token
      */
     async refreshConnection(force: boolean = false) {
-        const activeConnection = this.authUtil.auth.activeConnection
-        if (activeConnection?.state === 'valid' && activeConnection?.type === 'sso') {
+        const activeConnection = this.authUtil.conn
+        if (this.authUtil.isConnectionValid() && isSsoConnection(activeConnection)) {
             // send the token to the language server
             const token = await this.authUtil.getBearerToken()
             await (force ? this._updateBearerToken(token) : this.updateBearerToken(token))
@@ -118,7 +119,7 @@ export class AmazonQLspAuth {
             data: jwt,
             metadata: {
                 sso: {
-                    startUrl: AuthUtil.instance.auth.startUrl,
+                    startUrl: AuthUtil.instance.startUrl,
                 },
             },
             encrypted: true,

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

@@ -12,25 +12,11 @@ import { Commands, getLogger, globals, undefinedIfEmpty } from 'aws-core-vscode/
 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 { pushConfigUpdate } from '../config'
 
 export async function activate(languageClient: LanguageClient, encryptionKey: Buffer, mynahUIPath: string) {
     const disposables = globals.context.subscriptions
 
-    // Make sure we've sent an auth profile to the language server before even initializing the UI
-    await pushConfigUpdate(languageClient, {
-        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(),
-    })
-
     const provider = new AmazonQChatViewProvider(mynahUIPath)
 
     disposables.push(
@@ -79,10 +65,6 @@ 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', () => {
@@ -99,45 +81,3 @@ 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
-    }
-}
-type ProfileConfig = {
-    type: 'profile'
-    profileArn: string | undefined
-}
-type CustomizationConfig = {
-    type: 'customization'
-    customization: string | undefined
-}
-type LogLevelConfig = {
-    type: 'logLevel'
-}
-type QConfigs = ProfileConfig | CustomizationConfig | LogLevelConfig

packages/amazonq/src/lsp/client.ts

@@ -16,7 +16,6 @@ import {
     GetConfigurationFromServerParams,
     RenameFilesParams,
     ResponseMessage,
-    updateConfigurationRequestType,
     WorkspaceFolder,
 } from '@aws/language-server-runtimes/protocol'
 import { AuthUtil, CodeWhispererSettings, getSelectedCustomization } from 'aws-core-vscode/codewhisperer'
@@ -38,7 +37,7 @@ import {
 import { processUtils } from 'aws-core-vscode/shared'
 import { activate } from './chat/activation'
 import { AmazonQResourcePaths } from './lspInstaller'
-import { ConfigSection, isValidConfigSection, toAmazonQLSPLogLevel } from './config'
+import { ConfigSection, isValidConfigSection, pushConfigUpdate, toAmazonQLSPLogLevel } from './config'
 import { telemetry } from 'aws-core-vscode/telemetry'
 
 const localize = nls.loadMessageBundle()
@@ -160,120 +159,133 @@ export async function startLanguageServer(
 
     const disposable = client.start()
     toDispose.push(disposable)
+    await client.onReady()
 
-    const auth = new AmazonQLspAuth(client)
+    const auth = await initializeAuth(client)
 
-    return client.onReady().then(async () => {
-        await auth.refreshConnection()
+    await onLanguageServerReady(auth, client, resourcePaths, toDispose)
 
-        if (Experiments.instance.get('amazonqLSPInline', false)) {
-            const inlineManager = new InlineCompletionManager(client)
-            inlineManager.registerInlineCompletion()
-            toDispose.push(
-                inlineManager,
-                Commands.register({ id: 'aws.amazonq.invokeInlineCompletion', autoconnect: true }, async () => {
-                    await vscode.commands.executeCommand('editor.action.inlineSuggest.trigger')
-                }),
-                vscode.workspace.onDidCloseTextDocument(async () => {
-                    await vscode.commands.executeCommand('aws.amazonq.rejectCodeSuggestion')
-                })
-            )
-        }
+    return client
+}
 
-        if (Experiments.instance.get('amazonqChatLSP', true)) {
-            await activate(client, encryptionKey, resourcePaths.ui)
-        }
+async function initializeAuth(client: LanguageClient): Promise<AmazonQLspAuth> {
+    const auth = new AmazonQLspAuth(client)
+    await auth.refreshConnection(true)
+    return auth
+}
 
-        const refreshInterval = auth.startTokenRefreshInterval(10 * oneSecond)
+async function onLanguageServerReady(
+    auth: AmazonQLspAuth,
+    client: LanguageClient,
+    resourcePaths: AmazonQResourcePaths,
+    toDispose: vscode.Disposable[]
+) {
+    if (Experiments.instance.get('amazonqLSPInline', false)) {
+        const inlineManager = new InlineCompletionManager(client)
+        inlineManager.registerInlineCompletion()
+        toDispose.push(
+            inlineManager,
+            Commands.register({ id: 'aws.amazonq.invokeInlineCompletion', autoconnect: true }, async () => {
+                await vscode.commands.executeCommand('editor.action.inlineSuggest.trigger')
+            }),
+            vscode.workspace.onDidCloseTextDocument(async () => {
+                await vscode.commands.executeCommand('aws.amazonq.rejectCodeSuggestion')
+            })
+        )
+    }
 
-        const sendProfileToLsp = async () => {
-            try {
-                const result = await client.sendRequest(updateConfigurationRequestType.method, {
-                    section: 'aws.q',
-                    settings: {
-                        profileArn: AuthUtil.instance.regionProfileManager.activeRegionProfile?.arn,
-                    },
-                })
-                client.info(
-                    `Client: Updated Amazon Q Profile ${AuthUtil.instance.regionProfileManager.activeRegionProfile?.arn} to Amazon Q LSP`,
-                    result
-                )
-            } catch (err) {
-                client.error('Error when setting Q Developer Profile to Amazon Q LSP', err)
-            }
-        }
+    if (Experiments.instance.get('amazonqChatLSP', true)) {
+        await activate(client, encryptionKey, resourcePaths.ui)
+    }
 
-        // send profile to lsp once.
-        void sendProfileToLsp()
+    const refreshInterval = auth.startTokenRefreshInterval(10 * oneSecond)
 
-        toDispose.push(
-            AuthUtil.instance.auth.onDidChangeActiveConnection(async () => {
-                await auth.refreshConnection()
-            }),
-            AuthUtil.instance.auth.onDidDeleteConnection(async () => {
-                client.sendNotification(notificationTypes.deleteBearerToken.method)
-            }),
-            AuthUtil.instance.regionProfileManager.onDidChangeRegionProfile(sendProfileToLsp),
-            vscode.commands.registerCommand('aws.amazonq.getWorkspaceId', async () => {
-                const requestType = new RequestType<GetConfigurationFromServerParams, ResponseMessage, Error>(
-                    'aws/getConfigurationFromServer'
-                )
-                const workspaceIdResp = await client.sendRequest(requestType.method, {
-                    section: 'aws.q.workspaceContext',
-                })
-                return workspaceIdResp
-            }),
-            vscode.workspace.onDidCreateFiles((e) => {
-                client.sendNotification('workspace/didCreateFiles', {
-                    files: e.files.map((it) => {
-                        return { uri: it.fsPath }
-                    }),
-                } as CreateFilesParams)
-            }),
-            vscode.workspace.onDidDeleteFiles((e) => {
-                client.sendNotification('workspace/didDeleteFiles', {
-                    files: e.files.map((it) => {
-                        return { uri: it.fsPath }
+    // 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(),
+        })
+    }
+
+    toDispose.push(
+        AuthUtil.instance.auth.onDidChangeActiveConnection(async () => {
+            await auth.refreshConnection()
+        }),
+        AuthUtil.instance.auth.onDidDeleteConnection(async () => {
+            client.sendNotification(notificationTypes.deleteBearerToken.method)
+        }),
+        AuthUtil.instance.regionProfileManager.onDidChangeRegionProfile(() => sendProfileToLsp(client)),
+        vscode.commands.registerCommand('aws.amazonq.getWorkspaceId', async () => {
+            const requestType = new RequestType<GetConfigurationFromServerParams, ResponseMessage, Error>(
+                'aws/getConfigurationFromServer'
+            )
+            const workspaceIdResp = await client.sendRequest(requestType.method, {
+                section: 'aws.q.workspaceContext',
+            })
+            return workspaceIdResp
+        }),
+        vscode.workspace.onDidCreateFiles((e) => {
+            client.sendNotification('workspace/didCreateFiles', {
+                files: e.files.map((it) => {
+                    return { uri: it.fsPath }
+                }),
+            } as CreateFilesParams)
+        }),
+        vscode.workspace.onDidDeleteFiles((e) => {
+            client.sendNotification('workspace/didDeleteFiles', {
+                files: e.files.map((it) => {
+                    return { uri: it.fsPath }
+                }),
+            } as DeleteFilesParams)
+        }),
+        vscode.workspace.onDidRenameFiles((e) => {
+            client.sendNotification('workspace/didRenameFiles', {
+                files: e.files.map((it) => {
+                    return { oldUri: it.oldUri.fsPath, newUri: it.newUri.fsPath }
+                }),
+            } as RenameFilesParams)
+        }),
+        vscode.workspace.onDidSaveTextDocument((e) => {
+            client.sendNotification('workspace/didSaveTextDocument', {
+                textDocument: {
+                    uri: e.uri.fsPath,
+                },
+            } as DidSaveTextDocumentParams)
+        }),
+        vscode.workspace.onDidChangeWorkspaceFolders((e) => {
+            client.sendNotification('workspace/didChangeWorkspaceFolder', {
+                event: {
+                    added: e.added.map((it) => {
+                        return {
+                            name: it.name,
+                            uri: it.uri.fsPath,
+                        } as WorkspaceFolder
                     }),
-                } as DeleteFilesParams)
-            }),
-            vscode.workspace.onDidRenameFiles((e) => {
-                client.sendNotification('workspace/didRenameFiles', {
-                    files: e.files.map((it) => {
-                        return { oldUri: it.oldUri.fsPath, newUri: it.newUri.fsPath }
+                    removed: e.removed.map((it) => {
+                        return {
+                            name: it.name,
+                            uri: it.uri.fsPath,
+                        } as WorkspaceFolder
                     }),
-                } as RenameFilesParams)
-            }),
-            vscode.workspace.onDidSaveTextDocument((e) => {
-                client.sendNotification('workspace/didSaveTextDocument', {
-                    textDocument: {
-                        uri: e.uri.fsPath,
-                    },
-                } as DidSaveTextDocumentParams)
-            }),
-            vscode.workspace.onDidChangeWorkspaceFolders((e) => {
-                client.sendNotification('workspace/didChangeWorkspaceFolder', {
-                    event: {
-                        added: e.added.map((it) => {
-                            return {
-                                name: it.name,
-                                uri: it.uri.fsPath,
-                            } as WorkspaceFolder
-                        }),
-                        removed: e.removed.map((it) => {
-                            return {
-                                name: it.name,
-                                uri: it.uri.fsPath,
-                            } as WorkspaceFolder
-                        }),
-                    },
-                } as DidChangeWorkspaceFoldersParams)
-            }),
-            { dispose: () => clearInterval(refreshInterval) },
-            // Set this inside onReady so that it only triggers on subsequent language server starts (not the first)
-            onServerRestartHandler(client, auth)
-        )
-    })
+                },
+            } as DidChangeWorkspaceFoldersParams)
+        }),
+        { dispose: () => clearInterval(refreshInterval) },
+        // 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,
+        })
+    }
 }
 
 /**