diff --git a/packages/amazonq/test/unit/codewhisperer/util/authUtil.test.ts b/packages/amazonq/test/unit/codewhisperer/util/authUtil.test.ts index 97e245fecd3..1795639e1e2 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 @@ -250,6 +251,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 = { startUrl: validProfile.startUrl, @@ -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..1419eaa4772 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() @@ -64,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 @@ -277,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') { @@ -291,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) } } @@ -402,58 +405,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 - } - } - } + } - if (toImport && profileId) { - getLogger().info(`codewhisperer: migrating SSO connection to LSP identity server...`) + 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) { + this.logger.info('existing LSP auth connection found. Skipping migration') + await memento.update(key, undefined) + return + } + } catch { + this.logger.info('unable to get token from LSP auth, proceeding migration') + } - const registrationKey = { - startUrl: toImport.startUrl, - region: toImport.ssoRegion, - scopes: amazonQScopes, + 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 + profileId = id + if (p.metadata.connectionState === 'valid') { + break } + } + } - await this.session.updateProfile(registrationKey) + if (toImport && profileId) { + this.logger.info('migrating SSO connection to LSP identity server...') - const cacheDir = getCacheDir() + const registrationKey = { + startUrl: toImport.startUrl, + region: toImport.ssoRegion, + scopes: amazonQScopes, + } - const fromRegistrationFile = getRegistrationCacheFile(cacheDir, registrationKey) - const toRegistrationFile = path.join( - cacheDir, - getFlareCacheFileName( - JSON.stringify({ - region: toImport.ssoRegion, - startUrl: toImport.startUrl, - tool: clientName, - }) - ) - ) + await this.session.updateProfile(registrationKey) - const fromTokenFile = getTokenCacheFile(cacheDir, profileId) - const toTokenFile = path.join(cacheDir, getFlareCacheFileName(this.profileName)) + const cacheDir = getCacheDir() - 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 - } - - 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) + this.logger.debug('Successfully renamed registration and token files') + } catch (err) { + this.logger.error(`Failed to rename files during migration: ${err}`) + throw err } + + 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(