fix(amazonq): If no matching Q developer profile in last-used global state, default to single profile if only 1 available

Author: opieter-aws

packages/amazonq/test/unit/codewhisperer/region/regionProfileManager.test.ts

@@ -9,6 +9,7 @@ import { AuthUtil, RegionProfile, RegionProfileManager, defaultServiceConfig } f
 import { globals } from 'aws-core-vscode/shared'
 import { constants } from 'aws-core-vscode/auth'
 import { createTestAuthUtil } from 'aws-core-vscode/test'
+import { randomUUID } from 'crypto'
 
 const enterpriseSsoStartUrl = 'https://enterprise.awsapps.com/start'
 const region = 'us-east-1'
@@ -158,7 +159,7 @@ describe('RegionProfileManager', async function () {
         })
     })
 
-    describe('persistence', function () {
+    describe('persistSelectedRegionProfile', function () {
         it('persistSelectedRegionProfile', async function () {
             await setupConnection('idc')
             await regionProfileManager.switchRegionProfile(profileFoo, 'user')
@@ -177,14 +178,13 @@ describe('RegionProfileManager', async function () {
 
             assert.strictEqual(state[AuthUtil.instance.profileName], profileFoo)
         })
+    })
 
-        it(`restoreRegionProfile`, async function () {
-            sinon.stub(regionProfileManager, 'listRegionProfile').resolves([profileFoo])
+    describe('restoreRegionProfile', function () {
+        beforeEach(async function () {
             await setupConnection('idc')
-            if (!AuthUtil.instance.isConnected()) {
-                fail('connection should not be undefined')
-            }
-
+        })
+        it('restores region profile if profile name matches', async function () {
             const state = {} as any
             state[AuthUtil.instance.profileName] = profileFoo
 
@@ -194,6 +194,51 @@ describe('RegionProfileManager', async function () {
 
             assert.strictEqual(regionProfileManager.activeRegionProfile, profileFoo)
         })
+
+        it('returns early when no profiles exist', async function () {
+            const state = {} as any
+            state[AuthUtil.instance.profileName] = undefined
+
+            await globals.globalState.update('aws.amazonq.regionProfiles', state)
+
+            await regionProfileManager.restoreRegionProfile()
+            assert.strictEqual(regionProfileManager.activeRegionProfile, undefined)
+        })
+
+        it('returns early when no profile name matches, and multiple profiles exist', async function () {
+            const state = {} as any
+            state[AuthUtil.instance.profileName] = undefined
+            state[randomUUID()] = profileFoo
+
+            await globals.globalState.update('aws.amazonq.regionProfiles', state)
+
+            await regionProfileManager.restoreRegionProfile()
+            assert.strictEqual(regionProfileManager.activeRegionProfile, undefined)
+        })
+
+        it('uses single profile when no profile name matches', async function () {
+            const state = {} as any
+            state[randomUUID()] = profileFoo
+
+            await globals.globalState.update('aws.amazonq.regionProfiles', state)
+
+            await regionProfileManager.restoreRegionProfile()
+
+            assert.strictEqual(regionProfileManager.activeRegionProfile, profileFoo)
+        })
+
+        it('handles cross-validation failure', async function () {
+            const state = {
+                [AuthUtil.instance.profileName]: profileFoo,
+            }
+            sinon.stub(regionProfileManager, 'loadPersistedRegionProfiles').returns(state)
+            sinon.stub(regionProfileManager, 'getProfiles').resolves([]) // No matching profile
+            const invalidateStub = sinon.stub(regionProfileManager, 'invalidateProfile')
+
+            await regionProfileManager.restoreRegionProfile()
+
+            assert.ok(invalidateStub.calledWith(profileFoo.arn))
+        })
     })
 
     describe('invalidate', function () {

packages/core/src/codewhisperer/region/regionProfileManager.ts

@@ -38,7 +38,7 @@ const endpoints = createConstantMap({
     'eu-central-1': 'https://q.eu-central-1.amazonaws.com/',
 })
 
-const getRegionProfile = () =>
+const getRegionProfiles = () =>
     globals.globalState.tryGet<{ [label: string]: RegionProfile }>('aws.amazonq.regionProfiles', Object, {})
 
 /**
@@ -86,9 +86,9 @@ export class RegionProfileManager {
     // This is a poller that handles synchornization of selected region profiles between different IDE windows.
     // It checks for changes in global state of region profile, invoking the change handler to switch profiles
     public globalStatePoller = GlobalStatePoller.create({
-        getState: getRegionProfile,
+        getState: getRegionProfiles,
         changeHandler: async () => {
-            const profile = this.loadPersistedRegionProfle()
+            const profile = this.loadPersistedRegionProfiles()
             void this._switchRegionProfile(profile[this.authProvider.profileName], 'reload')
         },
         pollIntervalInMs: 2000,
@@ -285,10 +285,23 @@ export class RegionProfileManager {
 
     // Note: should be called after [this.authProvider.isConnected()] returns non null
     async restoreRegionProfile() {
-        const previousSelected = this.loadPersistedRegionProfle()[this.authProvider.profileName] || undefined
-        if (!previousSelected) {
+        const profiles = this.loadPersistedRegionProfiles()
+        if (!profiles || Object.keys(profiles).length === 0) {
             return
         }
+
+        let previousSelected = profiles[this.authProvider.profileName]
+
+        // If no profile matches auth profileName and there are multiple profiles, return so user can select
+        if (!previousSelected && Object.keys(profiles).length > 1) {
+            return
+        }
+
+        // If no profile matches auth profileName but there's only one profile, use that one
+        if (!previousSelected && Object.keys(profiles).length === 1) {
+            previousSelected = Object.values(profiles)[0]
+        }
+
         // cross-validation
         this.getProfiles()
             .then(async (profiles) => {
@@ -319,8 +332,8 @@ export class RegionProfileManager {
         await this.switchRegionProfile(previousSelected, 'reload')
     }
 
-    private loadPersistedRegionProfle(): { [label: string]: RegionProfile } {
-        return getRegionProfile()
+    public loadPersistedRegionProfiles(): { [label: string]: RegionProfile } {
+        return getRegionProfiles()
     }
 
     async persistSelectRegionProfile() {
@@ -330,7 +343,7 @@ export class RegionProfileManager {
         }
 
         // persist connectionId to profileArn
-        const previousPersistedState = getRegionProfile()
+        const previousPersistedState = getRegionProfiles()
 
         previousPersistedState[this.authProvider.profileName] = this.activeRegionProfile
         await globals.globalState.update('aws.amazonq.regionProfiles', previousPersistedState)
@@ -379,7 +392,7 @@ export class RegionProfileManager {
                 this._activeRegionProfile = undefined
             }
 
-            const profiles = this.loadPersistedRegionProfle()
+            const profiles = this.loadPersistedRegionProfiles()
             const updatedProfiles = Object.fromEntries(
                 Object.entries(profiles).filter(([connId, profile]) => profile.arn !== arn)
             )