fix(amazonq): fix enterprise users not able to sign in correctly if they have 2+ vscode instances open

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

@@ -29,6 +29,7 @@ import { isAwsError, ToolkitError } from '../../shared/errors'
 import { telemetry } from '../../shared/telemetry/telemetry'
 import { localize } from '../../shared/utilities/vsCodeUtils'
 import { Commands } from '../../shared/vscode/commands2'
+import { CachedResource, GlobalStateSchema } from '../../shared/utilities/resourceCache'
 
 // TODO: is there a better way to manage all endpoint strings in one place?
 export const defaultServiceConfig: CodeWhispererConfig = {
@@ -50,6 +51,10 @@ const endpoints = createConstantMap({
  */
 export type ProfileSwitchIntent = 'user' | 'auth' | 'update' | 'reload'
 
+interface RegionProfileGlobalState extends GlobalStateSchema<RegionProfile> {
+    previousSelection: { [label: string]: RegionProfile }
+}
+
 export class RegionProfileManager {
     private static logger = getLogger()
     private _activeRegionProfile: RegionProfile | undefined
@@ -59,6 +64,22 @@ export class RegionProfileManager {
     // Store the last API results (for UI propuse) so we don't need to call service again if doesn't require "latest" result
     private _profiles: RegionProfile[] = []
 
+    private readonly cache = new (class extends CachedResource<RegionProfile[]> {
+        constructor(private readonly profileProvider: () => Promise<RegionProfile[]>) {
+            super('aws.amazonq.regionProfiles', 60000, {
+                resource: {
+                    locked: false,
+                    timestamp: 0,
+                    result: undefined,
+                },
+            })
+        }
+
+        override resourceProvider(): Promise<RegionProfile[]> {
+            return this.profileProvider()
+        }
+    })(this.listRegionProfile.bind(this))
+
     get activeRegionProfile() {
         const conn = this.connectionProvider()
         if (isBuilderIdConnection(conn)) {
@@ -104,6 +125,14 @@ export class RegionProfileManager {
 
     constructor(private readonly connectionProvider: () => Connection | undefined) {}
 
+    async resetCache() {
+        await this.cache.releaseLock()
+    }
+
+    async getProfiles(): Promise<RegionProfile[]> {
+        return this.cache.getResource()
+    }
+
     async listRegionProfile(): Promise<RegionProfile[]> {
         this._profiles = []
 
@@ -238,7 +267,7 @@ export class RegionProfileManager {
             return
         }
         // cross-validation
-        this.listRegionProfile()
+        this.getProfiles()
             .then(async (profiles) => {
                 const r = profiles.find((it) => it.arn === previousSelected.arn)
                 if (!r) {
@@ -268,13 +297,20 @@ export class RegionProfileManager {
     }
 
     private loadPersistedRegionProfle(): { [label: string]: RegionProfile } {
-        const previousPersistedState = globals.globalState.tryGet<{ [label: string]: RegionProfile }>(
+        const previousPersistedState = globals.globalState.tryGet<RegionProfileGlobalState>(
             'aws.amazonq.regionProfiles',
             Object,
-            {}
+            {
+                previousSelection: {},
+                resource: {
+                    locked: false,
+                    result: undefined,
+                    timestamp: 0,
+                },
+            }
         )
 
-        return previousPersistedState
+        return previousPersistedState.previousSelection
     }
 
     async persistSelectRegionProfile() {
@@ -286,21 +322,28 @@ export class RegionProfileManager {
         }
 
         // persist connectionId to profileArn
-        const previousPersistedState = globals.globalState.tryGet<{ [label: string]: RegionProfile }>(
+        const previousPersistedState = globals.globalState.tryGet<RegionProfileGlobalState>(
             'aws.amazonq.regionProfiles',
             Object,
-            {}
+            {
+                previousSelection: {},
+                resource: {
+                    locked: false,
+                    result: undefined,
+                    timestamp: 0,
+                },
+            }
         )
 
-        previousPersistedState[conn.id] = this.activeRegionProfile
+        previousPersistedState.previousSelection[conn.id] = this.activeRegionProfile
         await globals.globalState.update('aws.amazonq.regionProfiles', previousPersistedState)
     }
 
     async generateQuickPickItem(): Promise<DataQuickPickItem<string>[]> {
         const selected = this.activeRegionProfile
         let profiles: RegionProfile[] = []
         try {
-            profiles = await this.listRegionProfile()
+            profiles = await this.getProfiles()
         } catch (e) {
             return [
                 {
@@ -343,7 +386,23 @@ export class RegionProfileManager {
             const updatedProfiles = Object.fromEntries(
                 Object.entries(profiles).filter(([connId, profile]) => profile.arn !== arn)
             )
-            await globals.globalState.update('aws.amazonq.regionProfiles', updatedProfiles)
+            const previousPersistedState = globals.globalState.tryGet<RegionProfileGlobalState>(
+                'aws.amazonq.regionProfiles',
+                Object,
+                {
+                    previousSelection: {},
+                    resource: {
+                        locked: false,
+                        result: undefined,
+                        timestamp: 0,
+                    },
+                }
+            )
+            const toUpdate: RegionProfileGlobalState = {
+                previousSelection: updatedProfiles,
+                resource: previousPersistedState.resource,
+            }
+            await globals.globalState.update('aws.amazonq.regionProfiles', toUpdate)
         }
     }
 

packages/core/src/codewhisperer/util/authUtil.ts

@@ -142,6 +142,7 @@ export class AuthUtil {
 
             if (!this.isConnected()) {
                 await this.regionProfileManager.invalidateProfile(this.regionProfileManager.activeRegionProfile?.arn)
+                await this.regionProfileManager.resetCache()
             }
         })
 

packages/core/src/login/webview/vue/amazonq/backend_amazonq.ts

@@ -223,7 +223,7 @@ export class AmazonQLoginWebview extends CommonAuthWebview {
      */
     override async listRegionProfiles(): Promise<RegionProfile[] | string> {
         try {
-            return await AuthUtil.instance.regionProfileManager.listRegionProfile()
+            return await AuthUtil.instance.regionProfileManager.getProfiles()
         } catch (e) {
             const conn = AuthUtil.instance.conn as SsoConnection | undefined
             telemetry.amazonq_didSelectProfile.emit({

packages/core/src/shared/globalState.ts

@@ -49,6 +49,7 @@ export type globalKey =
     | 'aws.toolkit.lsp.manifest'
     | 'aws.amazonq.customization.overrideV2'
     | 'aws.amazonq.regionProfiles'
+    | 'aws.amazonq.regionProfiles.cache'
     // Deprecated/legacy names. New keys should start with "aws.".
     | '#sessionCreationDates' // Legacy name from `ssoAccessTokenProvider.ts`.
     | 'CODECATALYST_RECONNECT'

packages/core/src/shared/utilities/resourceCache.ts

@@ -0,0 +1,127 @@
+/*!
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import globals from '../extensionGlobals'
+import { globalKey } from '../globalState'
+import { getLogger } from '../logger/logger'
+import { waitUntil } from '../utilities/timeoutUtils'
+
+interface Resource<V> {
+    result: V | undefined
+    locked: boolean
+    timestamp: number
+}
+
+// GlobalStates schema, which is used for vscode global states deserialization
+// [globals.globalState.tryGet<T>]
+export interface GlobalStateSchema<V> {
+    resource: Resource<V>
+}
+
+const logger = getLogger()
+
+function now() {
+    return globals.clock.Date.now()
+}
+
+export abstract class CachedResource<V> {
+    constructor(
+        private readonly key: globalKey,
+        private readonly expirationInMilli: number,
+        private readonly defaultValue: GlobalStateSchema<V>
+    ) {}
+
+    abstract resourceProvider(): Promise<V>
+
+    async getResource(): Promise<V> {
+        const cachedValue = await this.readResourceAndLock()
+        const resource = cachedValue?.resource
+
+        // if cache is still fresh, return
+        if (cachedValue && resource && resource.result) {
+            if (now() - resource.timestamp < this.expirationInMilli) {
+                logger.info(`cache hit`)
+                // release the lock
+                await this.updateCache(cachedValue, {
+                    ...resource,
+                    locked: false,
+                })
+                return resource.result
+            } else {
+                logger.info(`cache hit but cached value is stale, invoking service API to pull the latest response`)
+            }
+        }
+
+        // catch and error case?
+        logger.info(`cache miss, invoking service API to pull the latest response`)
+        const latest = await this.resourceProvider()
+
+        // update resource cache and release the lock
+        const r: Resource<V> = {
+            locked: false,
+            timestamp: now(),
+            result: latest,
+        }
+        await this.updateCache(cachedValue, r)
+        return latest
+    }
+
+    async readResourceAndLock(): Promise<GlobalStateSchema<V> | undefined> {
+        const _acquireLock = async () => {
+            const cachedValue = this.readCacheOrDefault()
+
+            if (!cachedValue.resource.locked) {
+                await this.updateCache(cachedValue, {
+                    ...cachedValue.resource,
+                    locked: true,
+                })
+                return cachedValue
+            }
+
+            return undefined
+        }
+
+        const lock = await waitUntil(
+            async () => {
+                const lock = await _acquireLock()
+                logger.info(`try obtain resource cache read lock %s`, lock)
+                if (lock) {
+                    return lock
+                }
+            },
+            { timeout: 15000, interval: 1500, truthy: true } // TODO: pass via ctor
+        )
+
+        return lock
+    }
+
+    async releaseLock() {
+        await globals.globalState.update(this.key, {
+            ...this.readCacheOrDefault(),
+            locked: false,
+        })
+    }
+
+    private async updateCache(cache: GlobalStateSchema<any> | undefined, resource: Resource<any>) {
+        await globals.globalState.update(this.key, {
+            ...(cache ? cache : this.readCacheOrDefault()),
+            resource: resource,
+        })
+    }
+
+    private readCacheOrDefault(): GlobalStateSchema<V> {
+        const cachedValue = globals.globalState.tryGet<GlobalStateSchema<V>>(this.key, Object, {
+            ...this.defaultValue,
+            resource: {
+                ...this.defaultValue.resource,
+                locked: false,
+                result: undefined,
+                timestamp: 0,
+            },
+        })
+
+        return cachedValue
+    }
+}