fix(amazonq): Developer Profiles not being sent to Flare Auth as expected

packages/amazonq/src/developerProfile/util.ts

@@ -0,0 +1,49 @@
+/*!
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+import * as nls from 'vscode-nls'
+const localize = nls.loadMessageBundle()
+import vscode from 'vscode'
+import { AmazonQPromptSettings, placeholder } from 'aws-core-vscode/shared'
+import { telemetry } from 'aws-core-vscode/telemetry'
+import { selectRegionProfileCommand, toastMessage } from 'aws-core-vscode/codewhisperer'
+import { once } from 'aws-core-vscode/utils'
+
+/**
+ * Creates a toast message telling the user they need to select a Developer Profile
+ */
+export const notifySelectDeveloperProfile = once(_notifySelectDeveloperProfile)
+async function _notifySelectDeveloperProfile() {
+    const suppressId = 'amazonQSelectDeveloperProfile'
+    const settings = AmazonQPromptSettings.instance
+    const shouldShow = settings.isPromptEnabled(suppressId)
+    if (!shouldShow) {
+        return
+    }
+
+    const message = localize(
+        'aws.amazonq.profile.mustSelectMessage',
+        'You must select a Q Developer Profile for Amazon Q features to work.'
+    )
+    const selectProfile = 'Select Profile'
+    const dontShowAgain = 'Dont Show Again'
+
+    await telemetry.toolkit_showNotification.run(async () => {
+        telemetry.record({ id: 'mustSelectDeveloperProfileMessage' })
+        void vscode.window.showWarningMessage(message, selectProfile, dontShowAgain).then(async (resp) => {
+            await telemetry.toolkit_invokeAction.run(async () => {
+                if (resp === selectProfile) {
+                    // Show Profile
+                    telemetry.record({ action: 'select' })
+                    void selectRegionProfileCommand.execute(placeholder, toastMessage)
+                } else if (resp === dontShowAgain) {
+                    telemetry.record({ action: 'dontShowAgain' })
+                    await settings.disablePrompt(suppressId)
+                } else {
+                    telemetry.record({ action: 'ignore' })
+                }
+            })
+        })
+    })
+}

packages/amazonq/src/lsp/activation.ts

@@ -7,15 +7,13 @@ import vscode from 'vscode'
 import { startLanguageServer } from './client'
 import { AmazonQLspInstaller } from './lspInstaller'
 import { lspSetupStage, ToolkitError, messages } from 'aws-core-vscode/shared'
-import { AuthUtil } from 'aws-core-vscode/codewhisperer'
 
 export async function activate(ctx: vscode.ExtensionContext) {
     try {
         await lspSetupStage('all', async () => {
             const installResult = await new AmazonQLspInstaller().resolve()
             return await lspSetupStage('launch', () => startLanguageServer(ctx, installResult.resourcePaths))
         })
-        await AuthUtil.instance.restore()
     } catch (err) {
         const e = err as ToolkitError
         void messages.showViewLogsMessage(`Failed to launch Amazon Q language server: ${e.message}`)

packages/amazonq/src/lsp/client.ts

@@ -175,8 +175,23 @@ export async function startLanguageServer(
     toDispose.push(disposable)
 
     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()
+
+    await postStartLanguageServer(client, resourcePaths, toDispose)
 
+    return client
+}
+
+async function postStartLanguageServer(
+    client: LanguageClient,
+    resourcePaths: AmazonQResourcePaths,
+    toDispose: vscode.Disposable[]
+) {
     // Request handler for when the server wants to know about the clients auth connnection. Must be registered before the initial auth init call
     client.onRequest<ConnectionMetadata, Error>(auth2.notificationTypes.getConnectionMetadata.method, () => {
         return {
@@ -347,8 +362,6 @@ export async function startLanguageServer(
         // Set this inside onReady so that it only triggers on subsequent language server starts (not the first)
         onServerRestartHandler(client)
     )
-
-    return client
 }
 
 /**

packages/core/src/auth/auth2.ts

@@ -325,10 +325,14 @@ export class SsoLogin implements BaseLogin {
     }
 
     private updateConnectionState(state: AuthState) {
-        if (this.connectionState !== state) {
-            this.eventEmitter.fire({ id: this.profileName, state })
+        const oldState = this.connectionState
+        const newState = state
+
+        this.connectionState = newState
+
+        if (oldState !== newState) {
+            this.eventEmitter.fire({ id: this.profileName, state: this.connectionState })
         }
-        this.connectionState = state
     }
 
     private ssoTokenChangedHandler(params: SsoTokenChangedParams) {

packages/core/src/codewhisperer/activation.ts

@@ -91,7 +91,6 @@ import { setContext } from '../shared/vscode/setContext'
 import { syncSecurityIssueWebview } from './views/securityIssue/securityIssueWebview'
 import { detectCommentAboveLine } from '../shared/utilities/commentUtils'
 import { activateEditTracking } from './nextEditPrediction/activation'
-import { notifySelectDeveloperProfile } from './region/utils'
 
 let localize: nls.LocalizeFunc
 
@@ -351,10 +350,6 @@ export async function activate(context: ExtContext): Promise<void> {
                     await AuthUtil.instance.notifySessionConfiguration()
                 }
             }
-
-            if (AuthUtil.instance.regionProfileManager.requireProfileSelection()) {
-                await notifySelectDeveloperProfile()
-            }
         },
         { emit: false, functionId: { name: 'activateCwCore' } }
     )

packages/core/src/codewhisperer/util/authUtil.ts

@@ -30,6 +30,8 @@ import { builderIdStartUrl, internalStartUrl } from '../../auth/sso/constants'
 import { VSCODE_EXTENSION_ID } from '../../shared/extensions'
 import { RegionProfileManager } from '../region/regionProfileManager'
 import { AuthFormId } from '../../login/webview/vue/types'
+import { notifySelectDeveloperProfile } from '../region/utils'
+import { once } from '../../shared/utilities/functionUtils'
 
 const localize = nls.loadMessageBundle()
 
@@ -91,8 +93,20 @@ export class AuthUtil implements IAuthProvider {
         return this.session.loginType === LoginTypes.SSO
     }
 
+    /**
+     * HACK: Ideally we'd put {@link notifySelectDeveloperProfile} in to {@link restore}.
+     *       But because {@link refreshState} is only called if !isConnected, we cannot do it since
+     *       {@link notifySelectDeveloperProfile} needs {@link refreshState} to run so it can set
+     *       the Bearer Token in the LSP first.
+     */
+    startedConnected = false
+
     async restore() {
         await this.session.restore()
+        this.startedConnected = this.isConnected()
+
+        // HACK: The LSP should allow us to call refreshState multiple times, but for some reason it breaks.
+        // isConnected somehow allows it to work. Assumption is that the LSP does not handle redundant calls nicely.
         if (!this.isConnected()) {
             await this.refreshState()
         }
@@ -257,20 +271,24 @@ export class AuthUtil implements IAuthProvider {
 
     private async refreshState(state = this.getAuthState()) {
         if (state === 'expired' || state === 'notConnected') {
+            this.lspAuth.deleteBearerToken()
             if (this.isIdcConnection()) {
                 await this.regionProfileManager.invalidateProfile(this.regionProfileManager.activeRegionProfile?.arn)
                 await this.regionProfileManager.clearCache()
             }
-            this.lspAuth.deleteBearerToken()
         }
         if (state === 'connected') {
+            const bearerTokenParams = (await this.session.getToken()).updateCredentialsParams
+            await this.lspAuth.updateBearerToken(bearerTokenParams)
+
             if (this.isIdcConnection()) {
                 await this.regionProfileManager.restoreProfileSelection()
             }
-            const bearerTokenParams = (await this.session.getToken()).updateCredentialsParams
-            await this.lspAuth.updateBearerToken(bearerTokenParams)
         }
 
+        // regardless of state, send message at startup if user needs to select a Developer Profile
+        void this.tryNotifySelectDeveloperProfile()
+
         vsCodeState.isFreeTierLimitReached = false
         await this.setVscodeContextProps(state)
         await Promise.all([
@@ -283,6 +301,12 @@ export class AuthUtil implements IAuthProvider {
         }
     }
 
+    private tryNotifySelectDeveloperProfile = once(async () => {
+        if (this.regionProfileManager.requireProfileSelection() && this.startedConnected) {
+            await notifySelectDeveloperProfile()
+        }
+    })
+
     async getTelemetryMetadata(): Promise<TelemetryMetadata> {
         if (!this.isConnected()) {
             return {