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 7da0c5d825a..6a1d2d26f70 100644 --- a/packages/amazonq/src/lsp/client.ts +++ b/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(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 } /** 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) { 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/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) diff --git a/packages/core/src/codewhisperer/util/authUtil.ts b/packages/core/src/codewhisperer/util/authUtil.ts index 362b1ae7157..84fe432131a 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,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 { if (!this.isConnected()) { return {