fix(amazonq): confusing message during loading profiles

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

@@ -220,6 +220,10 @@ class QWebviewBrowser(val project: Project, private val parentDisposable: Dispos
                 )
             }
 
+            is BrowserMessage.ListProfiles -> {
+                handleListProfilesMessage()
+            }
+
             is BrowserMessage.PublishWebviewTelemetry -> {
                 publishTelemetry(message)
             }
@@ -262,60 +266,41 @@ class QWebviewBrowser(val project: Project, private val parentDisposable: Dispos
             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"
-            }
-
-            var errorMessage: String? = null
-            var profiles: List<QRegionProfile> = emptyList()
+        val stage = if (isQExpired(project)) {
+            "REAUTH"
+        } else if (isQConnected(project) && QRegionProfileManager.getInstance().isPendingProfileSelection(project)) {
+            "PROFILE_SELECT"
+        } else {
+            "START"
+        }
 
-            if (stage == "PROFILE_SELECT") {
-                try {
-                    profiles = QRegionProfileManager.getInstance().listRegionProfiles(project).orEmpty()
-                    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)
+        when (stage) {
+            "PROFILE_SELECT" -> {
+                val jsonData = """
+                    {
+                        stage: '$stage',
+                        status: 'pending'
                     }
-                }
+                """.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"}
-            }
-            """.trimIndent()
+            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()
 
-            runInEdt {
                 executeJS("window.ideClient.prepareUi($jsonData)")
             }
         }
@@ -330,6 +315,52 @@ class QWebviewBrowser(val project: Project, private val parentDisposable: Dispos
         jcefBrowser.loadHTML(getWebviewHTML(webScriptUri, query))
     }
 
+    private fun handleListProfilesMessage() {
+        ApplicationManager.getApplication().executeOnPooledThread {
+            var errorMessage = ""
+            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)
+                }
+
+                null
+            }
+
+            // auto-select the profile if users only have 1 and don't show the UI
+            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)
+                return@executeOnPooledThread
+            }
+
+            // required EDT as this entire block is executed on thread pool
+            runInEdt {
+                val jsonData = """
+                        {
+                            stage: 'PROFILE_SELECT',
+                            status: '${if (profiles != null) "succeeded" else "failed"}',
+                            profiles: ${writeValueAsString(profiles ?: "")},
+                            errorMessage: '$errorMessage'
+                        }
+                """.trimIndent()
+
+                executeJS("window.ideClient.prepareUi($jsonData)")
+            }
+        }
+    }
+
     companion object {
         private val LOG = getLogger<QWebviewBrowser>()
         private const val WEB_SCRIPT_URI = "http://webview/js/getStart.js"

plugins/core/jetbrains-community/src/software/aws/toolkits/jetbrains/core/webview/BrowserMessage.kt

@@ -27,7 +27,8 @@
     JsonSubTypes.Type(value = BrowserMessage.Reauth::class, name = "reauth"),
     JsonSubTypes.Type(value = BrowserMessage.SendUiClickTelemetry::class, name = "sendUiClickTelemetry"),
     JsonSubTypes.Type(value = BrowserMessage.SwitchProfile::class, name = "switchProfile"),
-    JsonSubTypes.Type(value = BrowserMessage.PublishWebviewTelemetry::class, name = "webviewTelemetry")
+    JsonSubTypes.Type(value = BrowserMessage.PublishWebviewTelemetry::class, name = "webviewTelemetry"),
+    JsonSubTypes.Type(value = BrowserMessage.ListProfiles::class, name = "listProfiles")
 )
 sealed interface BrowserMessage {
 
@@ -37,7 +38,7 @@
 
     data class SelectConnection(val connectionId: String) : BrowserMessage
 
     object ToggleBrowser : BrowserMessage
 
     object LoginBuilderId : BrowserMessage
 
@@ -66,6 +67,8 @@
         val arn: String,
     ) : BrowserMessage
 
+    object ListProfiles : BrowserMessage
+
     data class SendUiClickTelemetry(val signInOptionClicked: String?) : BrowserMessage
 
     data class PublishWebviewTelemetry(val event: String) : BrowserMessage

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: AuthSetupMessageFromIde | ListProfilesMessageFromIde) {
         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 as ListProfilesMessageFromIde)
+                break
+
+            default:
+                this.handleAuthSetupMessage(state as AuthSetupMessageFromIde)
+        }
+    }
+
+    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 {