From e58a425cdf52d30bcc992ae797dc137950c90f6b Mon Sep 17 00:00:00 2001 From: opieter-aws Date: Thu, 22 May 2025 16:14:19 -0400 Subject: [PATCH 1/3] fix(amazonq): handle existing Flare connection when migrating SSO connections --- .../unit/codewhisperer/util/authUtil.test.ts | 25 +++++ .../core/src/codewhisperer/util/authUtil.ts | 106 +++++++++++------- 2 files changed, 88 insertions(+), 43 deletions(-) diff --git a/packages/amazonq/test/unit/codewhisperer/util/authUtil.test.ts b/packages/amazonq/test/unit/codewhisperer/util/authUtil.test.ts index 97e245fecd3..d4e22e14b8d 100644 --- a/packages/amazonq/test/unit/codewhisperer/util/authUtil.test.ts +++ b/packages/amazonq/test/unit/codewhisperer/util/authUtil.test.ts @@ -225,6 +225,7 @@ describe('AuthUtil', async function () { }) describe('migrateSsoConnectionToLsp', function () { + let mockLspAuth: any let memento: any let cacheDir: string let fromRegistrationFile: string @@ -249,6 +250,9 @@ describe('AuthUtil', async function () { sinon.stub(mementoUtils, 'getEnvironmentSpecificMemento').returns(memento) sinon.stub(cache, 'getCacheDir').returns(cacheDir) + + mockLspAuth = (auth as any).lspAuth + mockLspAuth.getSsoToken.resolves(undefined) fromTokenFile = cache.getTokenCacheFile(cacheDir, 'profile1') const registrationKey = { @@ -269,6 +273,27 @@ describe('AuthUtil', async function () { sinon.restore() }) + it('skips migration if LSP token exists', async function () { + memento.get.returns({ profile1: validProfile }) + mockLspAuth.getSsoToken.resolves({ token: 'valid-token' }) + + await auth.migrateSsoConnectionToLsp('test-client') + + assert.ok(memento.update.calledWith('auth.profiles', undefined)) + assert.ok(!auth.session.updateProfile?.called) + }) + + it('proceeds with migration if LSP token check throws', async function () { + memento.get.returns({ profile1: validProfile }) + mockLspAuth.getSsoToken.rejects(new Error('Token check failed')) + const updateProfileStub = sinon.stub((auth as any).session, 'updateProfile').resolves() + + await auth.migrateSsoConnectionToLsp('test-client') + + assert.ok(updateProfileStub.calledOnce) + assert.ok(memento.update.calledWith('auth.profiles', undefined)) + }) + it('migrates valid SSO connection', async function () { memento.get.returns({ profile1: validProfile }) diff --git a/packages/core/src/codewhisperer/util/authUtil.ts b/packages/core/src/codewhisperer/util/authUtil.ts index bb1a4d11366..322d77d0a02 100644 --- a/packages/core/src/codewhisperer/util/authUtil.ts +++ b/packages/core/src/codewhisperer/util/authUtil.ts @@ -39,6 +39,7 @@ import { getEnvironmentSpecificMemento } from '../../shared/utilities/mementos' import { getCacheDir, getFlareCacheFileName, getRegistrationCacheFile, getTokenCacheFile } from '../../auth/sso/cache' import { notifySelectDeveloperProfile } from '../region/utils' import { once } from '../../shared/utilities/functionUtils' +import { CancellationTokenSource, SsoTokenSourceKind } from '@aws/language-server-runtimes/server-interface' const localize = nls.loadMessageBundle() @@ -402,58 +403,77 @@ export class AuthUtil implements IAuthProvider { if (!profiles) { return - } else { - getLogger().info(`codewhisperer: checking for old SSO connections`) - for (const [id, p] of Object.entries(profiles)) { - if (p.type === 'sso' && hasExactScopes(p.scopes ?? [], amazonQScopes)) { - toImport = p - profileId = id - if (p.metadata.connectionState === 'valid') { - break - } - } + } + + try { + // Try go get token from LSP auth. If available, skip migration and delete old auth profile + const token = await this.lspAuth.getSsoToken( + { + kind: SsoTokenSourceKind.IamIdentityCenter, + profileName: this.profileName, + }, + false, + new CancellationTokenSource().token + ) + if (token) { + getLogger().info(`codewhisperer: existing LSP auth connection found. Skipping migration`) + await memento.update(key, undefined) + return } + } catch { + getLogger().info(`codewhisperer: unable to get token from LSP auth, proceeding migration`) + } - if (toImport && profileId) { - getLogger().info(`codewhisperer: migrating SSO connection to LSP identity server...`) - - const registrationKey = { - startUrl: toImport.startUrl, - region: toImport.ssoRegion, - scopes: amazonQScopes, + getLogger().info(`codewhisperer: checking for old SSO connections`) + for (const [id, p] of Object.entries(profiles)) { + if (p.type === 'sso' && hasExactScopes(p.scopes ?? [], amazonQScopes)) { + toImport = p + profileId = id + if (p.metadata.connectionState === 'valid') { + break } + } + } - await this.session.updateProfile(registrationKey) - - const cacheDir = getCacheDir() + if (toImport && profileId) { + getLogger().info(`codewhisperer: migrating SSO connection to LSP identity server...`) - const fromRegistrationFile = getRegistrationCacheFile(cacheDir, registrationKey) - const toRegistrationFile = path.join( - cacheDir, - getFlareCacheFileName( - JSON.stringify({ - region: toImport.ssoRegion, - startUrl: toImport.startUrl, - tool: clientName, - }) - ) - ) + const registrationKey = { + startUrl: toImport.startUrl, + region: toImport.ssoRegion, + scopes: amazonQScopes, + } - const fromTokenFile = getTokenCacheFile(cacheDir, profileId) - const toTokenFile = path.join(cacheDir, getFlareCacheFileName(this.profileName)) + await this.session.updateProfile(registrationKey) - try { - await fs.rename(fromRegistrationFile, toRegistrationFile) - await fs.rename(fromTokenFile, toTokenFile) - getLogger().debug('Successfully renamed registration and token files') - } catch (err) { - getLogger().error(`Failed to rename files during migration: ${err}`) - throw err - } + const cacheDir = getCacheDir() - await memento.update(key, undefined) - getLogger().info(`codewhisperer: successfully migrated SSO connection to LSP identity server`) + const fromRegistrationFile = getRegistrationCacheFile(cacheDir, registrationKey) + const toRegistrationFile = path.join( + cacheDir, + getFlareCacheFileName( + JSON.stringify({ + region: toImport.ssoRegion, + startUrl: toImport.startUrl, + tool: clientName, + }) + ) + ) + + const fromTokenFile = getTokenCacheFile(cacheDir, profileId) + const toTokenFile = path.join(cacheDir, getFlareCacheFileName(this.profileName)) + + try { + await fs.rename(fromRegistrationFile, toRegistrationFile) + await fs.rename(fromTokenFile, toTokenFile) + getLogger().debug('Successfully renamed registration and token files') + } catch (err) { + getLogger().error(`Failed to rename files during migration: ${err}`) + throw err } + + getLogger().info(`codewhisperer: successfully migrated SSO connection to LSP identity server`) + await memento.update(key, undefined) } } } From ca51867304bdc5f3f7e0f34c1413b2f087215f7c Mon Sep 17 00:00:00 2001 From: opieter-aws Date: Thu, 22 May 2025 16:16:02 -0400 Subject: [PATCH 2/3] linftix --- .../amazonq/test/unit/codewhisperer/util/authUtil.test.ts | 4 ++-- packages/core/src/codewhisperer/util/authUtil.ts | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/packages/amazonq/test/unit/codewhisperer/util/authUtil.test.ts b/packages/amazonq/test/unit/codewhisperer/util/authUtil.test.ts index d4e22e14b8d..1795639e1e2 100644 --- a/packages/amazonq/test/unit/codewhisperer/util/authUtil.test.ts +++ b/packages/amazonq/test/unit/codewhisperer/util/authUtil.test.ts @@ -250,9 +250,9 @@ describe('AuthUtil', async function () { sinon.stub(mementoUtils, 'getEnvironmentSpecificMemento').returns(memento) sinon.stub(cache, 'getCacheDir').returns(cacheDir) - + mockLspAuth = (auth as any).lspAuth - mockLspAuth.getSsoToken.resolves(undefined) + mockLspAuth.getSsoToken.resolves(undefined) fromTokenFile = cache.getTokenCacheFile(cacheDir, 'profile1') const registrationKey = { diff --git a/packages/core/src/codewhisperer/util/authUtil.ts b/packages/core/src/codewhisperer/util/authUtil.ts index 322d77d0a02..050fd500695 100644 --- a/packages/core/src/codewhisperer/util/authUtil.ts +++ b/packages/core/src/codewhisperer/util/authUtil.ts @@ -403,8 +403,8 @@ export class AuthUtil implements IAuthProvider { if (!profiles) { return - } - + } + try { // Try go get token from LSP auth. If available, skip migration and delete old auth profile const token = await this.lspAuth.getSsoToken( From 5b22a6f3bc044af98cc39f637f12c9a5f3ad4e06 Mon Sep 17 00:00:00 2001 From: opieter-aws Date: Fri, 23 May 2025 09:12:33 -0400 Subject: [PATCH 3/3] Improve logging --- .../core/src/codewhisperer/util/authUtil.ts | 20 ++++++++++--------- packages/core/src/shared/logger/logger.ts | 1 + 2 files changed, 12 insertions(+), 9 deletions(-) diff --git a/packages/core/src/codewhisperer/util/authUtil.ts b/packages/core/src/codewhisperer/util/authUtil.ts index 050fd500695..1419eaa4772 100644 --- a/packages/core/src/codewhisperer/util/authUtil.ts +++ b/packages/core/src/codewhisperer/util/authUtil.ts @@ -65,6 +65,8 @@ export interface IAuthProvider { */ export class AuthUtil implements IAuthProvider { public readonly profileName = VSCODE_EXTENSION_ID.amazonq + protected logger = getLogger('amazonqAuth') + public readonly regionProfileManager: RegionProfileManager // IAM login currently not supported @@ -278,7 +280,7 @@ export class AuthUtil implements IAuthProvider { } private async cacheChangedHandler(event: cacheChangedEvent) { - getLogger().debug(`Auth: Cache change event received: ${event}`) + this.logger.debug(`Cache change event received: ${event}`) if (event === 'delete') { await this.logout() } else if (event === 'create') { @@ -292,7 +294,7 @@ export class AuthUtil implements IAuthProvider { await this.lspAuth.updateBearerToken(params!) return } else { - getLogger().info(`codewhisperer: connection changed to ${e.state}`) + this.logger.info(`codewhisperer: connection changed to ${e.state}`) await this.refreshState(e.state) } } @@ -416,15 +418,15 @@ export class AuthUtil implements IAuthProvider { new CancellationTokenSource().token ) if (token) { - getLogger().info(`codewhisperer: existing LSP auth connection found. Skipping migration`) + this.logger.info('existing LSP auth connection found. Skipping migration') await memento.update(key, undefined) return } } catch { - getLogger().info(`codewhisperer: unable to get token from LSP auth, proceeding migration`) + this.logger.info('unable to get token from LSP auth, proceeding migration') } - getLogger().info(`codewhisperer: checking for old SSO connections`) + this.logger.info('checking for old SSO connections') for (const [id, p] of Object.entries(profiles)) { if (p.type === 'sso' && hasExactScopes(p.scopes ?? [], amazonQScopes)) { toImport = p @@ -436,7 +438,7 @@ export class AuthUtil implements IAuthProvider { } if (toImport && profileId) { - getLogger().info(`codewhisperer: migrating SSO connection to LSP identity server...`) + this.logger.info('migrating SSO connection to LSP identity server...') const registrationKey = { startUrl: toImport.startUrl, @@ -466,13 +468,13 @@ export class AuthUtil implements IAuthProvider { try { await fs.rename(fromRegistrationFile, toRegistrationFile) await fs.rename(fromTokenFile, toTokenFile) - getLogger().debug('Successfully renamed registration and token files') + this.logger.debug('Successfully renamed registration and token files') } catch (err) { - getLogger().error(`Failed to rename files during migration: ${err}`) + this.logger.error(`Failed to rename files during migration: ${err}`) throw err } - getLogger().info(`codewhisperer: successfully migrated SSO connection to LSP identity server`) + this.logger.info('successfully migrated SSO connection to LSP identity server') await memento.update(key, undefined) } } diff --git a/packages/core/src/shared/logger/logger.ts b/packages/core/src/shared/logger/logger.ts index b398ff93162..bb94fb0dc53 100644 --- a/packages/core/src/shared/logger/logger.ts +++ b/packages/core/src/shared/logger/logger.ts @@ -21,6 +21,7 @@ export type LogTopic = | 'nextEditPrediction' | 'resourceCache' | 'telemetry' + | 'amazonqAuth' class ErrorLog { constructor(