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

Author: opieter-aws

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 })
 

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)
         }
     }
 }