feat(amazonq): Add region profile manager functionality (aws#7036)

## Problem
The [first version](aws#6958)
of migration of all auth for vscode to Flare using the identity server
did not support the recently introduced RegionProfileManager

## Solution
These code changes bring back the RegionProfileManager functionality, in
the new auth setup. Integration tests and unit tests to be fixed after
all references are updated in a follow-up PR to keep the changes
manageable. **CI is expected to fail**.

---

- 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.

Author: opieter-aws

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

@@ -8,13 +8,6 @@ import { getIcon } from '../../shared/icons'
 import { DataQuickPickItem } from '../../shared/ui/pickerPrompter'
 import { CodeWhispererConfig, RegionProfile } from '../models/model'
 import { showConfirmationMessage } from '../../shared/utilities/messages'
-import {
-    Connection,
-    isBuilderIdConnection,
-    isIdcSsoConnection,
-    isSsoConnection,
-    SsoConnection,
-} from '../../auth/connection'
 import globals from '../../shared/extensionGlobals'
 import { once } from '../../shared/utilities/functionUtils'
 import CodeWhispererUserClient from '../client/codewhispereruserclient'
@@ -28,6 +21,7 @@ import { parse } from '@aws-sdk/util-arn-parser'
 import { isAwsError, ToolkitError } from '../../shared/errors'
 import { telemetry } from '../../shared/telemetry/telemetry'
 import { localize } from '../../shared/utilities/vsCodeUtils'
+import { AuthUtil } from '../util/authUtil'
 
 // TODO: is there a better way to manage all endpoint strings in one place?
 export const defaultServiceConfig: CodeWhispererConfig = {
@@ -60,21 +54,19 @@ export class RegionProfileManager {
     private _profiles: RegionProfile[] = []
 
     get activeRegionProfile() {
-        const conn = this.connectionProvider()
-        if (isBuilderIdConnection(conn)) {
+        if (AuthUtil.instance.isBuilderIdConnection()) {
             return undefined
         }
         return this._activeRegionProfile
     }
 
     get clientConfig(): CodeWhispererConfig {
-        const conn = this.connectionProvider()
-        if (!conn) {
+        if (!AuthUtil.instance.isConnected()) {
             throw new ToolkitError('trying to get client configuration without credential')
         }
 
         // builder id should simply use default IAD
-        if (isBuilderIdConnection(conn)) {
+        if (AuthUtil.instance.isBuilderIdConnection()) {
             return defaultServiceConfig
         }
 
@@ -102,18 +94,17 @@ export class RegionProfileManager {
         return this._profiles
     }
 
-    constructor(private readonly connectionProvider: () => Connection | undefined) {}
+    constructor() {}
 
-    async listRegionProfile(): Promise<RegionProfile[]> {
+    async listRegionProfiles(): Promise<RegionProfile[]> {
         this._profiles = []
 
-        const conn = this.connectionProvider()
-        if (conn === undefined || !isSsoConnection(conn)) {
+        if (!AuthUtil.instance.isConnected() || !AuthUtil.instance.isSsoSession()) {
             return []
         }
         const availableProfiles: RegionProfile[] = []
         for (const [region, endpoint] of endpoints.entries()) {
-            const client = await this.createQClient(region, endpoint, conn as SsoConnection)
+            const client = await this.createQClient(region, endpoint)
             const requester = async (request: CodeWhispererUserClient.ListAvailableProfilesRequest) =>
                 client.listAvailableProfiles(request).promise()
             const request: CodeWhispererUserClient.ListAvailableProfilesRequest = {}
@@ -138,7 +129,7 @@ export class RegionProfileManager {
                 availableProfiles.push(...mappedPfs)
             } catch (e) {
                 const logMsg = isAwsError(e) ? `requestId=${e.requestId}; message=${e.message}` : (e as Error).message
-                RegionProfileManager.logger.error(`failed to listRegionProfile: ${logMsg}`)
+                RegionProfileManager.logger.error(`failed to listRegionProfiles: ${logMsg}`)
                 throw e
             }
 
@@ -150,18 +141,14 @@ export class RegionProfileManager {
     }
 
     async switchRegionProfile(regionProfile: RegionProfile | undefined, source: ProfileSwitchIntent) {
-        const conn = this.connectionProvider()
-        if (conn === undefined || !isIdcSsoConnection(conn)) {
+        if (!AuthUtil.instance.isConnected() || !AuthUtil.instance.isIdcConnection()) {
             return
         }
 
         if (regionProfile && this.activeRegionProfile && regionProfile.arn === this.activeRegionProfile.arn) {
             return
         }
 
-        // TODO: make it typesafe
-        const ssoConn = this.connectionProvider() as SsoConnection
-
         // only prompt to users when users switch from A profile to B profile
         if (this.activeRegionProfile !== undefined && regionProfile !== undefined) {
             const response = await showConfirmationMessage({
@@ -179,9 +166,9 @@ export class RegionProfileManager {
                 telemetry.amazonq_didSelectProfile.emit({
                     source: source,
                     amazonQProfileRegion: this.activeRegionProfile?.region ?? 'not-set',
-                    ssoRegion: ssoConn.ssoRegion,
+                    ssoRegion: AuthUtil.instance.connection?.region,
                     result: 'Cancelled',
-                    credentialStartUrl: ssoConn.startUrl,
+                    credentialStartUrl: AuthUtil.instance.connection?.startUrl,
                     profileCount: this.profiles.length,
                 })
                 return
@@ -198,9 +185,9 @@ export class RegionProfileManager {
             telemetry.amazonq_didSelectProfile.emit({
                 source: source,
                 amazonQProfileRegion: regionProfile?.region ?? 'not-set',
-                ssoRegion: ssoConn.ssoRegion,
+                ssoRegion: AuthUtil.instance.connection?.region,
                 result: 'Succeeded',
-                credentialStartUrl: ssoConn.startUrl,
+                credentialStartUrl: AuthUtil.instance.connection?.startUrl,
                 profileCount: this.profiles.length,
             })
         }
@@ -222,20 +209,19 @@ export class RegionProfileManager {
     }
 
     restoreProfileSelection = once(async () => {
-        const conn = this.connectionProvider()
-        if (conn) {
-            await this.restoreRegionProfile(conn)
+        if (AuthUtil.instance.isConnected()) {
+            await this.restoreRegionProfile()
         }
     })
 
     // Note: should be called after [AuthUtil.instance.conn] returns non null
-    async restoreRegionProfile(conn: Connection) {
-        const previousSelected = this.loadPersistedRegionProfle()[conn.id] || undefined
+    async restoreRegionProfile() {
+        const previousSelected = this.loadPersistedRegionProfle()[AuthUtil.instance.profileName] || undefined
         if (!previousSelected) {
             return
         }
         // cross-validation
-        this.listRegionProfile()
+        this.listRegionProfiles()
             .then(async (profiles) => {
                 const r = profiles.find((it) => it.arn === previousSelected.arn)
                 if (!r) {
@@ -275,10 +261,8 @@ export class RegionProfileManager {
     }
 
     async persistSelectRegionProfile() {
-        const conn = this.connectionProvider()
-
         // default has empty arn and shouldn't be persisted because it's just a fallback
-        if (!conn || this.activeRegionProfile === undefined) {
+        if (!AuthUtil.instance.isConnected() || this.activeRegionProfile === undefined) {
             return
         }
 
@@ -289,15 +273,15 @@ export class RegionProfileManager {
             {}
         )
 
-        previousPersistedState[conn.id] = this.activeRegionProfile
+        previousPersistedState[AuthUtil.instance.profileName] = 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.listRegionProfiles()
         } catch (e) {
             return [
                 {
@@ -344,8 +328,15 @@ export class RegionProfileManager {
         }
     }
 
-    async createQClient(region: string, endpoint: string, conn: SsoConnection): Promise<CodeWhispererUserClient> {
-        const token = (await conn.getToken()).accessToken
+    requireProfileSelection() {
+        if (AuthUtil.instance.isBuilderIdConnection()) {
+            return false
+        }
+        return AuthUtil.instance.isIdcConnection() && this.activeRegionProfile === undefined
+    }
+
+    async createQClient(region: string, endpoint: string): Promise<CodeWhispererUserClient> {
+        const token = await AuthUtil.instance.getToken()
         const serviceOption: ServiceOptions = {
             apiConfig: userApiConfig,
             region: region,

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

@@ -17,9 +17,10 @@ import { showAmazonQWalkthroughOnce } from '../../amazonq/onboardingPage/walkthr
 import { setContext } from '../../shared/vscode/setContext'
 import { openUrl } from '../../shared/utilities/vsCodeUtils'
 import { telemetry } from '../../shared/telemetry/telemetry'
-import { AuthStateEvent, AuthStates, LanguageClientAuth, LoginTypes, SsoLogin } from '../../auth/auth2'
+import { AuthStateEvent, LanguageClientAuth, LoginTypes, SsoLogin } from '../../auth/auth2'
 import { builderIdStartUrl } from '../../auth/sso/constants'
 import { VSCODE_EXTENSION_ID } from '../../shared/extensions'
+import { RegionProfileManager } from '../region/regionProfileManager'
 
 const localize = nls.loadMessageBundle()
 
@@ -34,6 +35,7 @@ export const amazonQScopes = [...codeWhispererChatScopes, ...scopesGumby, ...sco
  */
 export class AuthUtil {
     public readonly profileName = VSCODE_EXTENSION_ID.amazonq
+    public readonly regionProfileManager: RegionProfileManager
 
     // IAM login currently not supported
     private session: SsoLogin
@@ -53,6 +55,11 @@ export class AuthUtil {
     private constructor(private readonly lspAuth: LanguageClientAuth) {
         this.session = new SsoLogin(this.profileName, this.lspAuth)
         this.onDidChangeConnectionState((e: AuthStateEvent) => this.stateChangeHandler(e))
+
+        this.regionProfileManager = new RegionProfileManager()
+        this.regionProfileManager.onDidChangeRegionProfile(async () => {
+            await this.setVscodeContextProps()
+        })
     }
 
     isSsoSession() {
@@ -104,29 +111,32 @@ export class AuthUtil {
     }
 
     isConnected() {
-        return this.getAuthState() === AuthStates.CONNECTED
+        return this.getAuthState() === 'connected'
     }
 
     isConnectionExpired() {
-        return this.getAuthState() === AuthStates.EXPIRED
+        return this.getAuthState() === 'expired'
     }
 
     isBuilderIdConnection() {
         return this.connection?.startUrl === builderIdStartUrl
     }
 
     isIdcConnection() {
-        return this.connection?.startUrl && this.connection?.startUrl !== builderIdStartUrl
+        return Boolean(this.connection?.startUrl && this.connection?.startUrl !== builderIdStartUrl)
     }
 
     onDidChangeConnectionState(handler: (e: AuthStateEvent) => any) {
         return this.session.onDidChangeConnectionState(handler)
     }
 
     public async setVscodeContextProps(state = this.getAuthState()) {
-        await setContext('aws.codewhisperer.connected', state === AuthStates.CONNECTED)
-        await setContext('aws.amazonq.showLoginView', state !== AuthStates.CONNECTED) // Login view also handles expired state.
-        await setContext('aws.codewhisperer.connectionExpired', state === AuthStates.EXPIRED)
+        await setContext('aws.codewhisperer.connected', state === 'connected')
+        const showAmazonQLoginView =
+            !this.isConnected() || this.isConnectionExpired() || this.regionProfileManager.requireProfileSelection()
+        await setContext('aws.amazonq.showLoginView', showAmazonQLoginView)
+        await setContext('aws.amazonq.connectedSsoIdc', this.isIdcConnection())
+        await setContext('aws.codewhisperer.connectionExpired', state === 'expired')
     }
 
     private reauthenticatePromptShown: boolean = false
@@ -214,10 +224,16 @@ export class AuthUtil {
     }
 
     private async refreshState(state = this.getAuthState()) {
-        if (state === AuthStates.EXPIRED || state === AuthStates.NOT_CONNECTED) {
+        if (state === 'expired' || state === 'notConnected') {
+            if (this.isIdcConnection()) {
+                await this.regionProfileManager.invalidateProfile(this.regionProfileManager.activeRegionProfile?.arn)
+            }
             this.lspAuth.deleteBearerToken()
         }
-        if (state === AuthStates.CONNECTED) {
+        if (state === 'connected') {
+            if (this.isIdcConnection()) {
+                await this.regionProfileManager.restoreProfileSelection()
+            }
             const bearerTokenParams = (await this.session.getToken()).updateCredentialsParams
             await this.lspAuth.updateBearerToken(bearerTokenParams)
         }
@@ -229,7 +245,7 @@ export class AuthUtil {
             Commands.tryExecute('aws.amazonq.updateReferenceLog'),
         ])
 
-        if (state === AuthStates.CONNECTED && this.isIdcConnection()) {
+        if (state === 'connected' && this.isIdcConnection()) {
             void vscode.commands.executeCommand('aws.amazonq.notifyNewCustomizations')
         }
     }