Skip to content

Commit e411922

Browse files
authored
fix(amazonq): Handle developer profile migration more gracefully (#7362)
## Problem The cached/persisted Q Developer Profile selection is stored as a map of `{connectionID: RegionProfile}`, where `connectionID` is a `randomUUID`. When migrating to Flare auth, we move away from the concept of a connectionID, and we do not have access to the latest ID of a user. The result is that we cannot restore the user's last selected region profile, and always need users who update versions to make a profile selection. ## Solution To handle this more gracefully, we will: * Use regionProfile if matching auth profile name (existing logic) * If no match, check if there is only a single RegionProfile stored in lastUsed. If so, use that one * If no match, and multiple RegionProfiles are stored in lastUsed cache, make user select Unit tests added --- - 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.
1 parent 5ff4c32 commit e411922

File tree

2 files changed

+74
-16
lines changed

2 files changed

+74
-16
lines changed

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

Lines changed: 52 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@ import { AuthUtil, RegionProfile, RegionProfileManager, defaultServiceConfig } f
99
import { globals } from 'aws-core-vscode/shared'
1010
import { constants } from 'aws-core-vscode/auth'
1111
import { createTestAuthUtil } from 'aws-core-vscode/test'
12+
import { randomUUID } from 'crypto'
1213

1314
const enterpriseSsoStartUrl = 'https://enterprise.awsapps.com/start'
1415
const region = 'us-east-1'
@@ -158,7 +159,7 @@ describe('RegionProfileManager', async function () {
158159
})
159160
})
160161

161-
describe('persistence', function () {
162+
describe('persistSelectedRegionProfile', function () {
162163
it('persistSelectedRegionProfile', async function () {
163164
await setupConnection('idc')
164165
await regionProfileManager.switchRegionProfile(profileFoo, 'user')
@@ -177,14 +178,13 @@ describe('RegionProfileManager', async function () {
177178

178179
assert.strictEqual(state[AuthUtil.instance.profileName], profileFoo)
179180
})
181+
})
180182

181-
it(`restoreRegionProfile`, async function () {
182-
sinon.stub(regionProfileManager, 'listRegionProfile').resolves([profileFoo])
183+
describe('restoreRegionProfile', function () {
184+
beforeEach(async function () {
183185
await setupConnection('idc')
184-
if (!AuthUtil.instance.isConnected()) {
185-
fail('connection should not be undefined')
186-
}
187-
186+
})
187+
it('restores region profile if profile name matches', async function () {
188188
const state = {} as any
189189
state[AuthUtil.instance.profileName] = profileFoo
190190

@@ -194,6 +194,51 @@ describe('RegionProfileManager', async function () {
194194

195195
assert.strictEqual(regionProfileManager.activeRegionProfile, profileFoo)
196196
})
197+
198+
it('returns early when no profiles exist', async function () {
199+
const state = {} as any
200+
state[AuthUtil.instance.profileName] = undefined
201+
202+
await globals.globalState.update('aws.amazonq.regionProfiles', state)
203+
204+
await regionProfileManager.restoreRegionProfile()
205+
assert.strictEqual(regionProfileManager.activeRegionProfile, undefined)
206+
})
207+
208+
it('returns early when no profile name matches, and multiple profiles exist', async function () {
209+
const state = {} as any
210+
state[AuthUtil.instance.profileName] = undefined
211+
state[randomUUID()] = profileFoo
212+
213+
await globals.globalState.update('aws.amazonq.regionProfiles', state)
214+
215+
await regionProfileManager.restoreRegionProfile()
216+
assert.strictEqual(regionProfileManager.activeRegionProfile, undefined)
217+
})
218+
219+
it('uses single profile when no profile name matches', async function () {
220+
const state = {} as any
221+
state[randomUUID()] = profileFoo
222+
223+
await globals.globalState.update('aws.amazonq.regionProfiles', state)
224+
225+
await regionProfileManager.restoreRegionProfile()
226+
227+
assert.strictEqual(regionProfileManager.activeRegionProfile, profileFoo)
228+
})
229+
230+
it('handles cross-validation failure', async function () {
231+
const state = {
232+
[AuthUtil.instance.profileName]: profileFoo,
233+
}
234+
sinon.stub(regionProfileManager, 'loadPersistedRegionProfiles').returns(state)
235+
sinon.stub(regionProfileManager, 'getProfiles').resolves([]) // No matching profile
236+
const invalidateStub = sinon.stub(regionProfileManager, 'invalidateProfile')
237+
238+
await regionProfileManager.restoreRegionProfile()
239+
240+
assert.ok(invalidateStub.calledWith(profileFoo.arn))
241+
})
197242
})
198243

199244
describe('invalidate', function () {

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

Lines changed: 22 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -38,7 +38,7 @@ const endpoints = createConstantMap({
3838
'eu-central-1': 'https://q.eu-central-1.amazonaws.com/',
3939
})
4040

41-
const getRegionProfile = () =>
41+
const getRegionProfiles = () =>
4242
globals.globalState.tryGet<{ [label: string]: RegionProfile }>('aws.amazonq.regionProfiles', Object, {})
4343

4444
/**
@@ -86,9 +86,9 @@ export class RegionProfileManager {
8686
// This is a poller that handles synchornization of selected region profiles between different IDE windows.
8787
// It checks for changes in global state of region profile, invoking the change handler to switch profiles
8888
public globalStatePoller = GlobalStatePoller.create({
89-
getState: getRegionProfile,
89+
getState: getRegionProfiles,
9090
changeHandler: async () => {
91-
const profile = this.loadPersistedRegionProfle()
91+
const profile = this.loadPersistedRegionProfiles()
9292
void this._switchRegionProfile(profile[this.authProvider.profileName], 'reload')
9393
},
9494
pollIntervalInMs: 2000,
@@ -285,10 +285,23 @@ export class RegionProfileManager {
285285

286286
// Note: should be called after [this.authProvider.isConnected()] returns non null
287287
async restoreRegionProfile() {
288-
const previousSelected = this.loadPersistedRegionProfle()[this.authProvider.profileName] || undefined
289-
if (!previousSelected) {
288+
const profiles = this.loadPersistedRegionProfiles()
289+
if (!profiles || Object.keys(profiles).length === 0) {
290290
return
291291
}
292+
293+
let previousSelected = profiles[this.authProvider.profileName]
294+
295+
// If no profile matches auth profileName and there are multiple profiles, return so user can select
296+
if (!previousSelected && Object.keys(profiles).length > 1) {
297+
return
298+
}
299+
300+
// If no profile matches auth profileName but there's only one profile, use that one
301+
if (!previousSelected && Object.keys(profiles).length === 1) {
302+
previousSelected = Object.values(profiles)[0]
303+
}
304+
292305
// cross-validation
293306
this.getProfiles()
294307
.then(async (profiles) => {
@@ -319,8 +332,8 @@ export class RegionProfileManager {
319332
await this.switchRegionProfile(previousSelected, 'reload')
320333
}
321334

322-
private loadPersistedRegionProfle(): { [label: string]: RegionProfile } {
323-
return getRegionProfile()
335+
public loadPersistedRegionProfiles(): { [label: string]: RegionProfile } {
336+
return getRegionProfiles()
324337
}
325338

326339
async persistSelectRegionProfile() {
@@ -330,7 +343,7 @@ export class RegionProfileManager {
330343
}
331344

332345
// persist connectionId to profileArn
333-
const previousPersistedState = getRegionProfile()
346+
const previousPersistedState = getRegionProfiles()
334347

335348
previousPersistedState[this.authProvider.profileName] = this.activeRegionProfile
336349
await globals.globalState.update('aws.amazonq.regionProfiles', previousPersistedState)
@@ -379,7 +392,7 @@ export class RegionProfileManager {
379392
this._activeRegionProfile = undefined
380393
}
381394

382-
const profiles = this.loadPersistedRegionProfle()
395+
const profiles = this.loadPersistedRegionProfiles()
383396
const updatedProfiles = Object.fromEntries(
384397
Object.entries(profiles).filter(([connId, profile]) => profile.arn !== arn)
385398
)

0 commit comments

Comments
 (0)