profile loading ux

plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/QLoginWebview.kt

@@ -183,16 +183,16 @@
                     ToolkitConnectionManager.getInstance(project)
                         .activeConnectionForFeature(QConnection.getInstance()) as? AwsBearerTokenConnection
                     )?.let { connection ->
-                    runInEdt {
-                        SsoLogoutAction(connection).actionPerformed(
-                            AnActionEvent.createFromDataContext(
-                                "qBrowser",
-                                null,
-                                DataContext.EMPTY_CONTEXT
+                        runInEdt {
+                            SsoLogoutAction(connection).actionPerformed(
+                                AnActionEvent.createFromDataContext(
+                                    "qBrowser",
+                                    null,
+                                    DataContext.EMPTY_CONTEXT
+                                )
                             )
-                        )
+                        }
                     }
-                }
             }
 
             is BrowserMessage.Reauth -> {
@@ -262,60 +262,83 @@
             writeValueAsString(it)
         }
 
-        // TODO: pass "REAUTH" if connection expires
-        // Perform the potentially blocking AWS call outside the EDT to fetch available region profiles.
-        ApplicationManager.getApplication().executeOnPooledThread {
-            val stage = if (isQExpired(project)) {
-                "REAUTH"
-            } else if (isQConnected(project) && QRegionProfileManager.getInstance().isPendingProfileSelection(project)) {
-                "PROFILE_SELECT"
-            } else {
-                "START"
-            }
+        val stage = if (isQExpired(project)) {
+            "REAUTH"
+        } else if (isQConnected(project) && QRegionProfileManager.getInstance().isPendingProfileSelection(project)) {
+            "PROFILE_SELECT"
+        } else {
+            "START"
+        }
+
+        when (stage) {
+            "PROFILE_SELECT" -> {
+                ApplicationManager.getApplication().executeOnPooledThread {
+                    var errorMessage: String = ""
+                    val profiles = try {
+                        QRegionProfileManager.getInstance().listRegionProfiles(project)
+                    } catch (e: Exception) {
+                        e.message?.let {
+                            errorMessage = it
+                        }
+                        LOG.warn { "Failed to call listRegionProfiles API: $errorMessage" }
+                        val qConn = ToolkitConnectionManager.getInstance(project).activeConnectionForFeature(QConnection.getInstance())
+                        Telemetry.amazonq.didSelectProfile.use { span ->
+                            span.source(QProfileSwitchIntent.Auth.value)
+                                .amazonQProfileRegion(QRegionProfileManager.getInstance().activeProfile(project)?.region ?: "not-set")
+                                .ssoRegion((qConn as? AwsBearerTokenConnection)?.region)
+                                .credentialStartUrl((qConn as? AwsBearerTokenConnection)?.startUrl)
+                                .result(MetricResult.Failed)
+                                .reason(e.message)
+                        }
 
-            var errorMessage: String? = null
-            var profiles: List<QRegionProfile> = emptyList()
+                        null
+                    }
 
-            if (stage == "PROFILE_SELECT") {
-                try {
-                    profiles = QRegionProfileManager.getInstance().listRegionProfiles(project).orEmpty()
-                    if (profiles.size == 1) {
+                    if (profiles?.size == 1) {
                         LOG.debug { "User only have access to 1 Q profile, auto-selecting profile ${profiles.first().profileName} for ${project.name}" }
                         QRegionProfileManager.getInstance().switchProfile(project, profiles.first(), QProfileSwitchIntent.Update)
                     }
-                } catch (e: Exception) {
-                    errorMessage = e.message
-                    LOG.warn { "Failed to call listRegionProfiles API" }
-                    val qConn = ToolkitConnectionManager.getInstance(project).activeConnectionForFeature(QConnection.getInstance())
-                    Telemetry.amazonq.didSelectProfile.use { span ->
-                        span.source(QProfileSwitchIntent.Auth.value)
-                            .amazonQProfileRegion(QRegionProfileManager.getInstance().activeProfile(project)?.region ?: "not-set")
-                            .ssoRegion((qConn as? AwsBearerTokenConnection)?.region)
-                            .credentialStartUrl((qConn as? AwsBearerTokenConnection)?.startUrl)
-                            .result(MetricResult.Failed)
-                            .reason(e.message)
+
+                    // required EDT as this entire block is executed on thread pool
+                    runInEdt {
+                        val jsonData = """
+                        {
+                            stage: '$stage',
+                            status: '${if (profiles != null) "succeeded" else "failed"}',
+                            profiles: ${writeValueAsString(profiles ?: "")},
+                            errorMessage: '$errorMessage'
+                        }
+                        """.trimIndent()
+
+                        executeJS("window.ideClient.prepareUi($jsonData)")
                     }
                 }
-            }
 
-            val jsonData = """
-            {
-                stage: '$stage',
-                regions: $regions,
-                idcInfo: {
-                    profileName: '${lastLoginIdcInfo.profileName}',
-                    startUrl: '${lastLoginIdcInfo.startUrl}',
-                    region: '${lastLoginIdcInfo.region}'
-                },
-                cancellable: ${state.browserCancellable},
-                feature: '${state.feature}',
-                existConnections: ${writeValueAsString(selectionSettings.values.map { it.currentSelection }.toList())},
-                profiles: ${writeValueAsString(profiles)},
-                errorMessage: ${errorMessage?.let { "\"$it\"" } ?: "null"}
+                val jsonData = """
+                    {
+                        stage: '$stage',
+                        status: 'pending'
+                    }
+                """.trimIndent()
+                executeJS("window.ideClient.prepareUi($jsonData)")
             }
-            """.trimIndent()
 
-            runInEdt {
+            else -> {
+                val jsonData = """
+                    {
+                        stage: '$stage',
+                        regions: $regions,
+                        idcInfo: {
+                            profileName: '${lastLoginIdcInfo.profileName}',
+                            startUrl: '${lastLoginIdcInfo.startUrl}',
+                            region: '${lastLoginIdcInfo.region}'
+                        },
+                        cancellable: ${state.browserCancellable},
+                        feature: '${state.feature}',
+                        existConnections: ${writeValueAsString(selectionSettings.values.map { it.currentSelection }.toList())},
+                    }
+                """.trimIndent()
+
                 executeJS("window.ideClient.prepareUi($jsonData)")
             }
         }

plugins/core/webview/src/ideClient.ts

@@ -2,28 +2,62 @@
 // SPDX-License-Identifier: Apache-2.0
 
 import {Store} from "vuex";
-import {IdcInfo, Region, Stage, State, BrowserSetupData, AwsBearerTokenConnection, Profile} from "./model";
+import {
+    IdcInfo,
+    State,
+    AuthSetupMessageFromIde,
+    ListProfileResult,
+    ListProfileSuccessResult,
+    ListProfileFailureResult, ListProfilePendingResult, ListProfilesMessageFromIde
+} from "./model";
 import {WebviewTelemetry} from './webviewTelemetry'
 
 export class IdeClient {
     constructor(private readonly store: Store<State>) {}
 
     // TODO: design and improve the API here
 
-    prepareUi(state: BrowserSetupData) {
+    prepareUi(state: any) {
         WebviewTelemetry.instance.reset()
         console.log('browser is preparing UI with state ', state)
-        this.store.commit('setStage', state.stage)
         // hack as window.onerror don't have access to vuex store
         void ((window as any).uiState = state.stage)
         WebviewTelemetry.instance.willShowPage(state.stage)
-        this.store.commit('setSsoRegions', state.regions)
-        this.updateLastLoginIdcInfo(state.idcInfo)
-        this.store.commit("setCancellable", state.cancellable)
-        this.store.commit("setFeature", state.feature)
-        this.store.commit('setProfiles', state.profiles);
-        this.store.commit("setErrorMessage", state.errorMessage)
-        const existConnections = state.existConnections.map(it => {
+
+        this.store.commit('setStage', state.stage)
+
+        switch (state.stage) {
+            case "PROFILE_SELECT":
+                this.handleProfileSelectMessage(state)
+                break
+
+            default:
+                this.handleAuthSetupMessage(state)
+        }
+    }
+
+    private handleProfileSelectMessage(msg: ListProfilesMessageFromIde) {
+        let result: ListProfileResult | undefined
+        switch (msg.status) {
+            case 'succeeded':
+                result = new ListProfileSuccessResult(msg.profiles)
+                break
+            case 'failed':
+                result = new ListProfileFailureResult(msg.errorMessage)
+                break
+            case 'pending':
+                result = new ListProfilePendingResult()
+                break
+        }
+        this.store.commit('setProfilesResult', result)
+    }
+
+    private handleAuthSetupMessage(msg: AuthSetupMessageFromIde) {
+        this.store.commit('setSsoRegions', msg.regions)
+        this.updateLastLoginIdcInfo(msg.idcInfo)
+        this.store.commit("setCancellable", msg.cancellable)
+        this.store.commit("setFeature", msg.feature)
+        const existConnections = msg.existConnections.map(it => {
             return {
                 sessionName: it.sessionName,
                 startUrl: it.startUrl,
@@ -37,13 +71,6 @@ export class IdeClient {
         this.updateAuthorization(undefined)
     }
 
-    handleProfiles(profilesData: { profiles: Profile[] }) {
-        this.store.commit('setStage', 'PROFILE_SELECT')
-        console.debug("received profile data")
-        const availableProfiles: Profile[] = profilesData.profiles;
-        this.store.commit('setProfiles', availableProfiles);
-    }
-
     updateAuthorization(code: string | undefined) {
         this.store.commit('setAuthorizationCode', code)
         // TODO: mutage stage to AUTHing here probably makes life easier

plugins/core/webview/src/model.ts

@@ -1,17 +1,23 @@
 // Copyright 2024 Amazon.com, Inc. or its affiliates. All Rights Reserved.
 // SPDX-License-Identifier: Apache-2.0
 
-export type BrowserSetupData = {
+export type AuthSetupMessageFromIde = {
     stage: Stage,
     regions: Region[],
     idcInfo: IdcInfo,
     cancellable: boolean,
     feature: string,
     existConnections: AwsBearerTokenConnection[],
+}
+
+export type ListProfilesMessageFromIde = {
+    stage: Stage,
+    status: 'succeeded' | 'failed' | 'pending',
     profiles: Profile[],
     errorMessage: string
 }
 
+
 // plugin interface [AwsBearerTokenConnection]
 export interface AwsBearerTokenConnection {
     sessionName: string,
@@ -20,6 +26,7 @@ export interface AwsBearerTokenConnection {
     scopes: string[],
     id: string
 }
+
 export const SONO_URL = "https://view.awsapps.com/start"
 
 export type Stage =
@@ -54,9 +61,27 @@ export interface State {
     feature: Feature,
     cancellable: boolean,
     existingConnections: AwsBearerTokenConnection[],
-    profiles: Profile[],
-    selectedProfile: Profile | undefined,
-    errorMessage: string | undefined
+    listProfilesResult: ListProfileResult | undefined,
+    selectedProfile: Profile | undefined
+}
+
+export interface ListProfileResult {
+    status: 'succeeded' | 'failed' | 'pending'
+}
+
+export class ListProfileSuccessResult implements ListProfileResult {
+    status: 'succeeded' = 'succeeded'
+    constructor(readonly profiles: Profile[]) {}
+}
+
+export class ListProfileFailureResult implements ListProfileResult {
+    status: 'failed' = 'failed'
+    constructor(readonly errorMessage: string) {}
+}
+
+export class ListProfilePendingResult implements ListProfileResult {
+    status: 'pending' = 'pending'
+    constructor() {}
 }
 
 export enum LoginIdentifier {