handle followupAuth

Author: samgst-amazon

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

@@ -18,7 +18,6 @@ import kotlinx.coroutines.flow.distinctUntilChanged
 import kotlinx.coroutines.flow.launchIn
 import kotlinx.coroutines.flow.merge
 import kotlinx.coroutines.flow.onEach
-import kotlinx.coroutines.future.await
 import kotlinx.coroutines.launch
 import org.cef.browser.CefBrowser
 import org.eclipse.lsp4j.Position
@@ -33,6 +32,8 @@ import software.aws.toolkits.jetbrains.services.amazonq.lsp.encryption.JwtEncryp
 import software.aws.toolkits.jetbrains.services.amazonq.lsp.flareChat.AwsServerCapabilitiesProvider
 import software.aws.toolkits.jetbrains.services.amazonq.lsp.flareChat.ChatCommunicationManager
 import software.aws.toolkits.jetbrains.services.amazonq.lsp.flareChat.getTextDocumentIdentifier
+import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.AUTH_FOLLOW_UP_CLICKED
+import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.AuthFollowUpClickNotification
 import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.ButtonClickNotification
 import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.ButtonClickParams
 import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.ButtonClickResult
@@ -384,6 +385,13 @@ class BrowserConnector(
                     }
                 }
             }
+            AUTH_FOLLOW_UP_CLICKED -> {
+                val message = serializer.deserializeChatMessages<AuthFollowUpClickNotification>(node)
+                chatCommunicationManager.handleAuthFollowUpClicked(
+                    project,
+                    message.params
+                )
+            }
             CHAT_COPY_CODE_TO_CLIPBOARD -> {
                 handleChatNotification<CopyCodeToClipboardNotification, CopyCodeToClipboardParams>(node) { server, params ->
                     server.copyCodeToClipboard(params)

plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/flareChat/ChatCommunicationManager.kt

@@ -8,8 +8,18 @@ import com.intellij.openapi.components.Service
 import com.intellij.openapi.components.service
 import com.intellij.openapi.project.Project
 import org.eclipse.lsp4j.ProgressParams
+import software.aws.toolkits.core.utils.getLogger
+import software.aws.toolkits.core.utils.warn
+import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManager
+import software.aws.toolkits.jetbrains.core.credentials.logoutFromSsoConnection
+import software.aws.toolkits.jetbrains.core.credentials.pinning.QConnection
+import software.aws.toolkits.jetbrains.core.credentials.reauthConnectionIfNeeded
+import software.aws.toolkits.jetbrains.core.gettingstarted.editor.BearerTokenFeatureSet
+import software.aws.toolkits.jetbrains.core.gettingstarted.editor.checkBearerConnectionValidity
 import software.aws.toolkits.jetbrains.services.amazonq.lsp.AmazonQLspService
 import software.aws.toolkits.jetbrains.services.amazonq.lsp.flareChat.ProgressNotificationUtils.getObject
+import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.AuthFollowUpClickedParams
+import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.AuthFollowUpType
 import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.GetSerializedChatResult
 import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.OpenTabResult
 import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.SEND_CHAT_COMMAND_PROMPT
@@ -58,9 +68,44 @@ class ChatCommunicationManager {
         }
     }
 
+    fun handleAuthFollowUpClicked(project: Project, params: AuthFollowUpClickedParams) {
+        val incomingType = params.authFollowUpType
+        val connectionManager = ToolkitConnectionManager.getInstance(project)
+        try {
+            when (incomingType) {
+                AuthFollowUpType.RE_AUTH.value, AuthFollowUpType.MISSING_SCOPES.value -> {
+                    connectionManager.activeConnectionForFeature(QConnection.getInstance())?.let {
+                        reauthConnectionIfNeeded(project, it, isReAuth = true)
+                    }
+                    return
+                }
+                AuthFollowUpType.FULL_AUTH.value, AuthFollowUpType.USE_SUPPORTED_AUTH.value -> {
+                    // Logout by deleting token credentials
+                    val validConnection = checkBearerConnectionValidity(project, BearerTokenFeatureSet.Q)
+                    val connection = validConnection.activeConnectionBearer
+                    if (connection != null) {
+                        logoutFromSsoConnection(project, connection)
+                    } else {
+                        LOG.warn { "No valid connection found for logout" }
+                    }
+                    return
+                }
+                else -> {
+                    LOG.warn { "Unknown auth follow up type: $incomingType" }
+                    throw IllegalStateException("Error occurred while attempting to handle auth follow up: Unknown AuthFollowUpType $incomingType")
+                }
+            }
+        } catch (ex: Exception) {
+            LOG.warn(ex) { "Failed to handle authentication when auth follow up clicked" }
+            throw ex
+        }
+    }
+
     companion object {
         fun getInstance(project: Project) = project.service<ChatCommunicationManager>()
 
+        val LOG = getLogger<ChatCommunicationManager>()
+
         val pendingSerializedChatRequests = ConcurrentHashMap<String, CompletableFuture<GetSerializedChatResult>>()
         fun completeSerializedChatResponse(requestId: String, content: String) {
             pendingSerializedChatRequests.remove(requestId)?.complete(GetSerializedChatResult((content)))

plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/model/aws/chat/AuthFollowUpClicked.kt

@@ -0,0 +1,31 @@
+// Copyright 2025 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+package software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat
+
+data class AuthFollowUpClickNotification(
+    override val command: String,
+    override val params: AuthFollowUpClickedParams,
+) : ChatNotification<AuthFollowUpClickedParams>
+
+data class AuthFollowUpClickedParams(
+    val tabId: String,
+    val messageId: String,
+    val authFollowUpType: String,
+) {
+    companion object {
+        fun create(tabId: String, messageId: String, authType: AuthFollowUpType): AuthFollowUpClickedParams =
+            AuthFollowUpClickedParams(tabId, messageId, authType.value)
+    }
+}
+
+enum class AuthFollowUpType(val value: String) {
+    FULL_AUTH("full-auth"),
+    RE_AUTH("re-auth"),
+    MISSING_SCOPES("missing_scopes"),
+    USE_SUPPORTED_AUTH("use-supported-auth"),
+    ;
+
+    override fun toString(): String =
+        name.lowercase()
+}

plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/model/aws/chat/FlareChatCommands.kt

@@ -13,6 +13,7 @@ const val PROMPT_INPUT_OPTIONS_CHANGE = "aws/chat/promptInputOptionChange"
 const val CHAT_PROMPT_OPTION_ACKNOWLEDGED = "chatPromptOptionAcknowledged"
 const val CHAT_FEEDBACK = "aws/chat/feedback"
 const val CHAT_FOLLOW_UP_CLICK = "aws/chat/followUpClick"
+const val AUTH_FOLLOW_UP_CLICKED = "authFollowUpClicked"
 const val CHAT_LIST_CONVERSATIONS = "aws/chat/listConversations"
 const val CHAT_CONVERSATION_CLICK = "aws/chat/conversationClick"
 const val CHAT_FILE_CLICK = "aws/chat/fileClick"