fix(amazonq): validate outdated profile (#2012)

* move validate in factory

* centralize all clients

Author: evanliu048

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

@@ -42,6 +42,8 @@ import software.aws.toolkits.jetbrains.utils.isQConnected
 import software.aws.toolkits.jetbrains.utils.isQExpired
 import software.aws.toolkits.jetbrains.utils.isQWebviewsAvailable
 import software.aws.toolkits.telemetry.FeatureId
+import software.aws.toolkits.telemetry.MetricResult
+import software.aws.toolkits.telemetry.Telemetry
 import software.aws.toolkits.telemetry.UiTelemetry
 import software.aws.toolkits.telemetry.WebviewTelemetry
 import java.awt.event.ActionListener
@@ -267,6 +269,16 @@ class QWebviewBrowser(val project: Project, private val parentDisposable: Dispos
                 QRegionProfileManager.getInstance().listRegionProfiles(project).orEmpty()
             } 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)
+                }
                 emptyList()
             }
 

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

@@ -14,8 +14,6 @@ import com.intellij.ui.components.panels.Wrapper
 import com.intellij.util.ui.components.BorderLayoutPanel
 import software.aws.toolkits.core.utils.debug
 import software.aws.toolkits.core.utils.getLogger
-import software.aws.toolkits.core.utils.warn
-import software.aws.toolkits.jetbrains.core.credentials.AwsBearerTokenConnection
 import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnection
 import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManager
 import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManagerListener
@@ -28,7 +26,6 @@ import software.aws.toolkits.jetbrains.core.webview.BrowserState
 import software.aws.toolkits.jetbrains.services.amazonq.QWebviewPanel
 import software.aws.toolkits.jetbrains.services.amazonq.RefreshQChatPanelButtonPressedListener
 import software.aws.toolkits.jetbrains.services.amazonq.gettingstarted.openMeetQPage
-import software.aws.toolkits.jetbrains.services.amazonq.profile.QProfileSwitchIntent
 import software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfile
 import software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfileManager
 import software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfileSelectedListener
@@ -37,8 +34,6 @@ import software.aws.toolkits.jetbrains.utils.isQExpired
 import software.aws.toolkits.jetbrains.utils.isQWebviewsAvailable
 import software.aws.toolkits.resources.message
 import software.aws.toolkits.telemetry.FeatureId
-import software.aws.toolkits.telemetry.MetricResult
-import software.aws.toolkits.telemetry.Telemetry
 import java.awt.event.ComponentAdapter
 import java.awt.event.ComponentEvent
 
@@ -70,19 +65,7 @@ class AmazonQToolWindowFactory : ToolWindowFactory, DumbAware {
                 override fun activeConnectionChanged(newConnection: ToolkitConnection?) {
                     ToolkitConnectionManager.getInstance(project).activeConnectionForFeature(QConnection.getInstance())?.let { qConn ->
                         openMeetQPage(project)
-                        try {
-                            QRegionProfileManager.getInstance().listRegionProfiles(project)
-                        } catch (e: Exception) {
-                            LOG.warn { "Failed to call listRegionProfiles API" }
-                            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)
-                            }
-                        }
+                        QRegionProfileManager.getInstance().validateProfile(project)
                     }
                     prepareChatContent(project, qPanel)
                 }

plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqCodeTest/controller/CodeTestChatController.kt

@@ -50,13 +50,11 @@ import software.aws.toolkits.core.utils.debug
 import software.aws.toolkits.core.utils.getLogger
 import software.aws.toolkits.core.utils.info
 import software.aws.toolkits.core.utils.warn
-import software.aws.toolkits.jetbrains.core.AwsClientManager
 import software.aws.toolkits.jetbrains.core.coroutines.EDT
-import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManager
-import software.aws.toolkits.jetbrains.core.credentials.pinning.QConnection
 import software.aws.toolkits.jetbrains.core.credentials.sono.isInternalUser
 import software.aws.toolkits.jetbrains.services.amazonq.apps.AmazonQAppInitContext
 import software.aws.toolkits.jetbrains.services.amazonq.auth.AuthController
+import software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfileManager
 import software.aws.toolkits.jetbrains.services.amazonq.project.RelevantDocument
 import software.aws.toolkits.jetbrains.services.amazonqCodeTest.CodeWhispererUTGChatManager
 import software.aws.toolkits.jetbrains.services.amazonqCodeTest.ConversationState
@@ -277,9 +275,6 @@ class CodeTestChatController(
                 promptInputDisabledState = true,
             )
             // Send Request to Sync UTG API
-            val connection = ToolkitConnectionManager.getInstance(project).activeConnectionForFeature(QConnection.getInstance())
-                // this should never happen because it should have been handled upstream by [AuthController]
-                ?: error("connection was found to be null")
             val contextExtractor = ActiveFileContextExtractor.create(fqnWebviewAdapter = null, project = project)
             val activeFileContext = ActiveFileContext(
                 fileContext = FileContext(
@@ -301,7 +296,7 @@ class CodeTestChatController(
                 useRelevantDocuments = false,
             )
 
-            val client = AwsClientManager.getInstance().getClient<CodeWhispererStreamingAsyncClient>(connection.getConnectionSettings())
+            val client = QRegionProfileManager.getInstance().getQClient<CodeWhispererStreamingAsyncClient>(project)
             val request = requestData.toChatRequest()
             client.generateAssistantResponse(request, responseHandler).await()
             // TODO: Need to send isCodeBlockSelected field

plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/cwc/clients/chat/v1/ChatSessionV1.kt

@@ -36,10 +36,8 @@ import software.amazon.awssdk.services.codewhispererstreaming.model.UserInputMes
 import software.amazon.awssdk.services.codewhispererstreaming.model.UserIntent
 import software.aws.toolkits.core.utils.getLogger
 import software.aws.toolkits.core.utils.info
-import software.aws.toolkits.jetbrains.core.AwsClientManager
 import software.aws.toolkits.jetbrains.core.coroutines.getCoroutineBgContext
-import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManager
-import software.aws.toolkits.jetbrains.core.credentials.pinning.QConnection
+import software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfileManager
 import software.aws.toolkits.jetbrains.services.amazonq.project.RelevantDocument
 import software.aws.toolkits.jetbrains.services.cwc.ChatConstants
 import software.aws.toolkits.jetbrains.services.cwc.clients.chat.ChatSession
@@ -170,11 +168,7 @@ class ChatSessionV1(
 
         try {
             withTimeout(ChatConstants.REQUEST_TIMEOUT_MS.toLong()) {
-                val connection = ToolkitConnectionManager.getInstance(project).activeConnectionForFeature(QConnection.getInstance())
-                    // this should never happen because it should have been handled upstream by [AuthController]
-                    ?: error("connection was found to be null")
-
-                val client = AwsClientManager.getInstance().getClient<CodeWhispererStreamingAsyncClient>(connection.getConnectionSettings())
+                val client = QRegionProfileManager.getInstance().getQClient<CodeWhispererStreamingAsyncClient>(project)
                 val request = data.toChatRequest()
                 logger.info { "Request from tab: ${data.tabId}, conversationId: $conversationId, request: $request" }
                 client.generateAssistantResponse(request, responseHandler).await()

plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/profile/QRegionProfileManager.kt

@@ -52,16 +52,14 @@ class QRegionProfileManager : PersistentStateComponent<QProfileState>, Disposabl
 
         if (profiles == null || profiles.none { it.arn == selected.arn }) {
             invalidateProfile(selected.arn)
+            switchProfile(project, null, intent = QProfileSwitchIntent.Reload)
             Telemetry.amazonq.profileState.use { span ->
                 span.source(QProfileSwitchIntent.Reload.value)
                     .amazonQProfileRegion(selected.region)
                     .ssoRegion(conn?.region)
                     .credentialStartUrl(conn?.startUrl)
                     .result(MetricResult.Failed)
             }
-            project.messageBus
-                .syncPublisher(QRegionProfileSelectedListener.TOPIC)
-                .onProfileSelected(project, null)
         }
     }
 
@@ -94,39 +92,39 @@ class QRegionProfileManager : PersistentStateComponent<QProfileState>, Disposabl
 
     fun hasValidConnectionButNoActiveProfile(project: Project): Boolean = getIdcConnectionOrNull(project) != null && activeProfile(project) == null
 
-    fun switchProfile(project: Project, newProfile: QRegionProfile, intent: QProfileSwitchIntent) {
+    fun switchProfile(project: Project, newProfile: QRegionProfile?, intent: QProfileSwitchIntent) {
         val conn = getIdcConnectionOrNull(project) ?: return
 
-        if (newProfile.arn.isEmpty()) return
-
         val oldProfile = connectionIdToActiveProfile[conn.id]
         if (oldProfile == newProfile) return
 
         connectionIdToActiveProfile[conn.id] = newProfile
         LOG.debug { "Switch from profile $oldProfile to $newProfile for project ${project.name}" }
 
-        if (intent == QProfileSwitchIntent.User || intent == QProfileSwitchIntent.Auth) {
-            notifyInfo(
-                title = message("action.q.profile.usage.text"),
-                content = message("action.q.profile.usage", newProfile.profileName),
-                project = project
-            )
-
-            Telemetry.amazonq.didSelectProfile.use { span ->
-                span.source(intent.value)
-                    .amazonQProfileRegion(newProfile.region)
-                    .profileCount(connectionIdToProfileList[conn.id])
-                    .ssoRegion(conn.region)
-                    .credentialStartUrl(conn.startUrl)
-                    .result(MetricResult.Succeeded)
-            }
-        } else {
-            Telemetry.amazonq.profileState.use { span ->
-                span.source(intent.value)
-                    .amazonQProfileRegion(newProfile.region)
-                    .ssoRegion(conn.region)
-                    .credentialStartUrl(conn.startUrl)
-                    .result(MetricResult.Succeeded)
+        if (newProfile != null) {
+            if (intent == QProfileSwitchIntent.User || intent == QProfileSwitchIntent.Auth) {
+                notifyInfo(
+                    title = message("action.q.profile.usage.text"),
+                    content = message("action.q.profile.usage", newProfile.profileName),
+                    project = project
+                )
+
+                Telemetry.amazonq.didSelectProfile.use { span ->
+                    span.source(intent.value)
+                        .amazonQProfileRegion(newProfile.region)
+                        .profileCount(connectionIdToProfileList[conn.id])
+                        .ssoRegion(conn.region)
+                        .credentialStartUrl(conn.startUrl)
+                        .result(MetricResult.Succeeded)
+                }
+            } else {
+                Telemetry.amazonq.profileState.use { span ->
+                    span.source(intent.value)
+                        .amazonQProfileRegion(newProfile.region)
+                        .ssoRegion(conn.region)
+                        .credentialStartUrl(conn.startUrl)
+                        .result(MetricResult.Succeeded)
+                }
             }
         }