fix(amazonq): handle existing Flare connection when migrating SSO connections (aws#7363)

opieter-aws · web-flow · commit 0ddbf4a92808 · 2025-05-23T10:03:52.000-04:00
## Problem Auth migration to LSP is not handled gracefully when a user downgrades and upgrades to auth on LSP multiple times, causing users to be logged out if they upgrade a second time ## Solution In the auth migration script, call the LSP identity server to check if a token is available. If the token is available, don't migrate the auth connection. If no token is available, migrate. Added unit tests for the case. --- - Treat all work as PUBLIC. Private `feature/x` branches will not be squash-merged at release time. - Your code changes must meet the guidelines in [CONTRIBUTING.md](https://github.com/aws/aws-toolkit-vscode/blob/master/CONTRIBUTING.md#guidelines). - License: I confirm that my contribution is made under the terms of the Apache 2.0 license.
diff --git 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
@@ -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
@@ -21,6 +21,7 @@ export type LogTopic =
     | 'nextEditPrediction'
     | 'resourceCache'
     | 'telemetry'
+    | 'amazonqAuth'
 
 class ErrorLog {
     constructor(
