Merge pull request aws#7293 from nkomonen-amazon/profiles

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

Author: opieter-aws

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/region/utils.ts

@@ -7,7 +7,6 @@ const localize = nls.loadMessageBundle()
 import { AmazonQPromptSettings } from '../../shared/settings'
 import { telemetry } from '../../shared/telemetry/telemetry'
 import vscode from 'vscode'
-import { selectRegionProfileCommand } from '../commands/basicCommands'
 import { placeholder } from '../../shared/vscode/commands2'
 import { toastMessage } from '../commands/types'
 
@@ -36,7 +35,7 @@ export async function notifySelectDeveloperProfile() {
                 if (resp === selectProfile) {
                     // Show Profile
                     telemetry.record({ action: 'select' })
-                    void selectRegionProfileCommand.execute(placeholder, toastMessage)
+                    void vscode.commands.executeCommand('aws.amazonq.selectRegionProfile', placeholder, toastMessage)
                 } else if (resp === dontShowAgain) {
                     telemetry.record({ action: 'dontShowAgain' })
                     await settings.disablePrompt(suppressId)

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,9 +93,31 @@ 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.
+     */
+    didStartSignedIn = false
+
     async restore() {
         await this.session.restore()
-        if (!this.isConnected()) {
+        this.didStartSignedIn = this.isConnected()
+
+        // HACK: We noticed that if calling `refreshState()` here when the user was already signed in, something broke.
+        //       So as a solution we only call it if they were not already signed in.
+        //
+        //       But in the case where a user was already signed in, we allow `session.restore()` to trigger `refreshState()` through
+        //       event emitters.
+        //       This is unoptimal since `refreshState()` should be able to be called multiple times and still work.
+        //
+        //       Because of this edge case, when `restore()` is called we cannot assume all Auth is setup when this function returns,
+        //       since we may still be waiting on the event emitter to trigger the expected functions.
+        //
+        //       TODO: Figure out why removing the if statement below causes things to break. Maybe we just need to
+        //             promisify the call and any subsequent callers will not make a redundant call.
+        if (!this.didStartSignedIn) {
             await this.refreshState()
         }
     }
@@ -257,20 +281,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 +311,12 @@ export class AuthUtil implements IAuthProvider {
         }
     }
 
+    private tryNotifySelectDeveloperProfile = once(async () => {
+        if (this.regionProfileManager.requireProfileSelection() && this.didStartSignedIn) {
+            await notifySelectDeveloperProfile()
+        }
+    })
+
     async getTelemetryMetadata(): Promise<TelemetryMetadata> {
         if (!this.isConnected()) {
             return {