From 945c05f073f88c3d22030da7617857a68fee23ba Mon Sep 17 00:00:00 2001 From: nkomonen-amazon Date: Fri, 9 May 2025 16:49:25 -0400 Subject: [PATCH 1/8] set auth state before firing event This may not actually do anything, but we risked something reading the stale state after emitting the event because we set it after the fire(). Signed-off-by: nkomonen-amazon --- packages/core/src/auth/auth2.ts | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/packages/core/src/auth/auth2.ts b/packages/core/src/auth/auth2.ts index 39a2dd6c8f6..75bc7523f0c 100644 --- a/packages/core/src/auth/auth2.ts +++ b/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) { From 520acb05250a5f2ea021882a8bf9307271da8a1f Mon Sep 17 00:00:00 2001 From: nkomonen-amazon Date: Mon, 12 May 2025 13:01:10 -0400 Subject: [PATCH 2/8] refactor: move Q LSP startup code in to separate function Signed-off-by: nkomonen-amazon --- packages/amazonq/src/lsp/client.ts | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/packages/amazonq/src/lsp/client.ts b/packages/amazonq/src/lsp/client.ts index 7da0c5d825a..341ab75bc63 100644 --- a/packages/amazonq/src/lsp/client.ts +++ b/packages/amazonq/src/lsp/client.ts @@ -177,6 +177,16 @@ export async function startLanguageServer( await client.onReady() AuthUtil.create(new auth2.LanguageClientAuth(client, clientId, encryptionKey)) + 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(auth2.notificationTypes.getConnectionMetadata.method, () => { return { @@ -347,8 +357,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 } /** From 33a5736f51a5a4d69fc41bbb7e05bb1b3c0feb70 Mon Sep 17 00:00:00 2001 From: nkomonen-amazon Date: Mon, 12 May 2025 13:14:49 -0400 Subject: [PATCH 3/8] refactor: move auth restore to right after creating the instance Signed-off-by: nkomonen-amazon --- packages/amazonq/src/lsp/activation.ts | 2 -- packages/amazonq/src/lsp/client.ts | 5 +++++ 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/packages/amazonq/src/lsp/activation.ts b/packages/amazonq/src/lsp/activation.ts index fe1e2f15ddc..aebb4a60479 100644 --- a/packages/amazonq/src/lsp/activation.ts +++ b/packages/amazonq/src/lsp/activation.ts @@ -7,7 +7,6 @@ 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 { @@ -15,7 +14,6 @@ export async function activate(ctx: vscode.ExtensionContext) { 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}`) diff --git a/packages/amazonq/src/lsp/client.ts b/packages/amazonq/src/lsp/client.ts index 341ab75bc63..6a1d2d26f70 100644 --- a/packages/amazonq/src/lsp/client.ts +++ b/packages/amazonq/src/lsp/client.ts @@ -175,7 +175,12 @@ 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) From 6958c25c942696a8c3d390aea9f29028508a6db7 Mon Sep 17 00:00:00 2001 From: nkomonen-amazon Date: Mon, 12 May 2025 19:08:15 -0400 Subject: [PATCH 4/8] fix: update Auth LSP before anything else The Auth LSP complains if we set the developer profile before setting the bearer token, because it assumes we should not be setting a profile if a token does not exist. So to fix this, we first ensure the server is updated with the token. Signed-off-by: nkomonen-amazon --- packages/core/src/codewhisperer/util/authUtil.ts | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/packages/core/src/codewhisperer/util/authUtil.ts b/packages/core/src/codewhisperer/util/authUtil.ts index 362b1ae7157..b178ab8ca14 100644 --- a/packages/core/src/codewhisperer/util/authUtil.ts +++ b/packages/core/src/codewhisperer/util/authUtil.ts @@ -257,18 +257,20 @@ 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) } vsCodeState.isFreeTierLimitReached = false From 779501f0d51bb9ce2059b94dbaaea37e8b28032e Mon Sep 17 00:00:00 2001 From: nkomonen-amazon Date: Mon, 12 May 2025 19:19:38 -0400 Subject: [PATCH 5/8] fix: Show "Must select Developer Profile" message Problem: THe show developer profile message was appearing at the wrong time. It was due to the order that the code executed. This was because when we called to check if the profiles were restored from cache, the logic that restored them had not actually run yet. Solution: Move the logic that decides to show the Needs Profile message, to after the code that restores the Profile. It is also only called once since we only want to show the message on startup. Signed-off-by: nkomonen-amazon --- packages/amazonq/src/developerProfile/util.ts | 49 +++++++++++++++++++ packages/core/src/codewhisperer/activation.ts | 5 -- .../core/src/codewhisperer/util/authUtil.ts | 24 ++++++++- 3 files changed, 72 insertions(+), 6 deletions(-) create mode 100644 packages/amazonq/src/developerProfile/util.ts diff --git a/packages/amazonq/src/developerProfile/util.ts b/packages/amazonq/src/developerProfile/util.ts new file mode 100644 index 00000000000..0091f4fb766 --- /dev/null +++ b/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' }) + } + }) + }) + }) +} diff --git a/packages/core/src/codewhisperer/activation.ts b/packages/core/src/codewhisperer/activation.ts index e4aaaeda934..d929758da6d 100644 --- a/packages/core/src/codewhisperer/activation.ts +++ b/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 { await AuthUtil.instance.notifySessionConfiguration() } } - - if (AuthUtil.instance.regionProfileManager.requireProfileSelection()) { - await notifySelectDeveloperProfile() - } }, { emit: false, functionId: { name: 'activateCwCore' } } ) diff --git a/packages/core/src/codewhisperer/util/authUtil.ts b/packages/core/src/codewhisperer/util/authUtil.ts index b178ab8ca14..cbca5997ed4 100644 --- a/packages/core/src/codewhisperer/util/authUtil.ts +++ b/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() } @@ -258,7 +272,6 @@ 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() @@ -273,6 +286,9 @@ export class AuthUtil implements IAuthProvider { } } + // 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([ @@ -285,6 +301,12 @@ export class AuthUtil implements IAuthProvider { } } + private tryNotifySelectDeveloperProfile = once(async () => { + if (this.regionProfileManager.requireProfileSelection() && this.startedConnected) { + await notifySelectDeveloperProfile() + } + }) + async getTelemetryMetadata(): Promise { if (!this.isConnected()) { return { From e2c41970436eabae9beb562a92bf618e60d05eaa Mon Sep 17 00:00:00 2001 From: nkomonen-amazon Date: Tue, 13 May 2025 10:57:59 -0400 Subject: [PATCH 6/8] fix comments - Fix docstring explaining refreshState fragility - Fix variable name Signed-off-by: nkomonen-amazon --- .../core/src/codewhisperer/util/authUtil.ts | 24 +++++++++++++------ 1 file changed, 17 insertions(+), 7 deletions(-) diff --git a/packages/core/src/codewhisperer/util/authUtil.ts b/packages/core/src/codewhisperer/util/authUtil.ts index cbca5997ed4..84fe432131a 100644 --- a/packages/core/src/codewhisperer/util/authUtil.ts +++ b/packages/core/src/codewhisperer/util/authUtil.ts @@ -99,15 +99,25 @@ export class AuthUtil implements IAuthProvider { * {@link notifySelectDeveloperProfile} needs {@link refreshState} to run so it can set * the Bearer Token in the LSP first. */ - startedConnected = false + didStartSignedIn = 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()) { + 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() } } @@ -302,7 +312,7 @@ export class AuthUtil implements IAuthProvider { } private tryNotifySelectDeveloperProfile = once(async () => { - if (this.regionProfileManager.requireProfileSelection() && this.startedConnected) { + if (this.regionProfileManager.requireProfileSelection() && this.didStartSignedIn) { await notifySelectDeveloperProfile() } }) From 5965e1d8dd9762af5cf747cce72f5e5cfe65bd0c Mon Sep 17 00:00:00 2001 From: nkomonen-amazon Date: Tue, 13 May 2025 17:03:04 -0400 Subject: [PATCH 7/8] remove stale code Signed-off-by: nkomonen-amazon --- packages/amazonq/src/developerProfile/util.ts | 49 ------------------- 1 file changed, 49 deletions(-) delete mode 100644 packages/amazonq/src/developerProfile/util.ts diff --git a/packages/amazonq/src/developerProfile/util.ts b/packages/amazonq/src/developerProfile/util.ts deleted file mode 100644 index 0091f4fb766..00000000000 --- a/packages/amazonq/src/developerProfile/util.ts +++ /dev/null @@ -1,49 +0,0 @@ -/*! - * 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' }) - } - }) - }) - }) -} From d46ad37e976eb7c12c2a0a0c8172f0f77f6b3fc2 Mon Sep 17 00:00:00 2001 From: nkomonen-amazon Date: Tue, 13 May 2025 17:36:26 -0400 Subject: [PATCH 8/8] fix circular dependency Signed-off-by: nkomonen-amazon --- packages/core/src/codewhisperer/region/utils.ts | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/packages/core/src/codewhisperer/region/utils.ts b/packages/core/src/codewhisperer/region/utils.ts index dd988f74a30..fb768e3b710 100644 --- a/packages/core/src/codewhisperer/region/utils.ts +++ b/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)