From 8da7fce08b2150daa5875c2fc00dc1f33cee9fef Mon Sep 17 00:00:00 2001 From: samgst-amazon Date: Fri, 2 May 2025 16:20:05 -0700 Subject: [PATCH 1/5] handle followupAuth --- .../amazonq/webview/BrowserConnector.kt | 10 ++++- .../lsp/flareChat/ChatCommunicationManager.kt | 45 +++++++++++++++++++ .../lsp/model/aws/chat/AuthFollowUpClicked.kt | 31 +++++++++++++ .../lsp/model/aws/chat/FlareChatCommands.kt | 1 + 4 files changed, 86 insertions(+), 1 deletion(-) create mode 100644 plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/model/aws/chat/AuthFollowUpClicked.kt diff --git a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/webview/BrowserConnector.kt b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/webview/BrowserConnector.kt index b7d76719348..4442a8f6cbd 100644 --- a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/webview/BrowserConnector.kt +++ b/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(node) + chatCommunicationManager.handleAuthFollowUpClicked( + project, + message.params + ) + } CHAT_COPY_CODE_TO_CLIPBOARD -> { handleChatNotification(node) { server, params -> server.copyCodeToClipboard(params) diff --git a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/flareChat/ChatCommunicationManager.kt b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/flareChat/ChatCommunicationManager.kt index f148bc2ffe7..10d26032197 100644 --- a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/flareChat/ChatCommunicationManager.kt +++ b/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() + val LOG = getLogger() + val pendingSerializedChatRequests = ConcurrentHashMap>() fun completeSerializedChatResponse(requestId: String, content: String) { pendingSerializedChatRequests.remove(requestId)?.complete(GetSerializedChatResult((content))) diff --git a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/model/aws/chat/AuthFollowUpClicked.kt b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/model/aws/chat/AuthFollowUpClicked.kt new file mode 100644 index 00000000000..1c73d8aeb8b --- /dev/null +++ b/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 + +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() +} diff --git a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/model/aws/chat/FlareChatCommands.kt b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/model/aws/chat/FlareChatCommands.kt index 02934b87db1..e0f4d8ace9f 100644 --- a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/model/aws/chat/FlareChatCommands.kt +++ b/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" From 7c8a681ce2e8b22f991dea11166c30033e1f50c0 Mon Sep 17 00:00:00 2001 From: samgst-amazon Date: Mon, 5 May 2025 15:41:51 -0700 Subject: [PATCH 2/5] followup capital casing --- .../amazonq/lsp/flareChat/ChatCommunicationManager.kt | 9 ++++----- .../amazonq/lsp/model/aws/chat/AuthFollowUpClicked.kt | 6 +++--- 2 files changed, 7 insertions(+), 8 deletions(-) diff --git a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/flareChat/ChatCommunicationManager.kt b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/flareChat/ChatCommunicationManager.kt index 10d26032197..ca70dfe2ce1 100644 --- a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/flareChat/ChatCommunicationManager.kt +++ b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/flareChat/ChatCommunicationManager.kt @@ -19,7 +19,7 @@ import software.aws.toolkits.jetbrains.core.gettingstarted.editor.checkBearerCon 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.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 @@ -69,17 +69,17 @@ class ChatCommunicationManager { } fun handleAuthFollowUpClicked(project: Project, params: AuthFollowUpClickedParams) { - val incomingType = params.authFollowUpType + val incomingType = params.authFollowupType val connectionManager = ToolkitConnectionManager.getInstance(project) try { when (incomingType) { - AuthFollowUpType.RE_AUTH.value, AuthFollowUpType.MISSING_SCOPES.value -> { + 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 -> { + AuthFollowupType.FULL_AUTH.value, AuthFollowupType.USE_SUPPORTED_AUTH.value -> { // Logout by deleting token credentials val validConnection = checkBearerConnectionValidity(project, BearerTokenFeatureSet.Q) val connection = validConnection.activeConnectionBearer @@ -92,7 +92,6 @@ class ChatCommunicationManager { } 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) { diff --git a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/model/aws/chat/AuthFollowUpClicked.kt b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/model/aws/chat/AuthFollowUpClicked.kt index 1c73d8aeb8b..f3d068947bd 100644 --- a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/model/aws/chat/AuthFollowUpClicked.kt +++ b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/model/aws/chat/AuthFollowUpClicked.kt @@ -11,15 +11,15 @@ data class AuthFollowUpClickNotification( data class AuthFollowUpClickedParams( val tabId: String, val messageId: String, - val authFollowUpType: String, + val authFollowupType: String, ) { companion object { - fun create(tabId: String, messageId: String, authType: AuthFollowUpType): AuthFollowUpClickedParams = + fun create(tabId: String, messageId: String, authType: AuthFollowupType): AuthFollowUpClickedParams = AuthFollowUpClickedParams(tabId, messageId, authType.value) } } -enum class AuthFollowUpType(val value: String) { +enum class AuthFollowupType(val value: String) { FULL_AUTH("full-auth"), RE_AUTH("re-auth"), MISSING_SCOPES("missing_scopes"), From a9fd11764ccede3fc5f35ac5c66bb8cb4523be91 Mon Sep 17 00:00:00 2001 From: samgst-amazon Date: Wed, 7 May 2025 16:28:46 -0700 Subject: [PATCH 3/5] handle in reauth function and sync q profile --- .../lsp/flareChat/ChatCommunicationManager.kt | 28 ++++++++----------- 1 file changed, 12 insertions(+), 16 deletions(-) diff --git a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/flareChat/ChatCommunicationManager.kt b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/flareChat/ChatCommunicationManager.kt index ca70dfe2ce1..3ce274d6302 100644 --- a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/flareChat/ChatCommunicationManager.kt +++ b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/flareChat/ChatCommunicationManager.kt @@ -11,11 +11,8 @@ 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 @@ -23,6 +20,8 @@ import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.AuthF 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 +import software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfileManager +import software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfileSelectedListener import java.util.UUID import java.util.concurrent.CompletableFuture import java.util.concurrent.ConcurrentHashMap @@ -72,22 +71,19 @@ class ChatCommunicationManager { val incomingType = params.authFollowupType val connectionManager = ToolkitConnectionManager.getInstance(project) try { + connectionManager.activeConnectionForFeature(QConnection.getInstance())?.let { + reauthConnectionIfNeeded(project, it, isReAuth = true) + } when (incomingType) { - AuthFollowupType.RE_AUTH.value, AuthFollowupType.MISSING_SCOPES.value -> { - connectionManager.activeConnectionForFeature(QConnection.getInstance())?.let { - reauthConnectionIfNeeded(project, it, isReAuth = true) - } + AuthFollowupType.USE_SUPPORTED_AUTH.value -> { + project.messageBus.syncPublisher(QRegionProfileSelectedListener.TOPIC) + .onProfileSelected(project, QRegionProfileManager.getInstance().activeProfile(project)) 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" } - } + AuthFollowupType.RE_AUTH.value, + AuthFollowupType.MISSING_SCOPES.value, + AuthFollowupType.FULL_AUTH.value, + -> { return } else -> { From 5892ee39d872dc052172f8b38158d10dd3f0e951 Mon Sep 17 00:00:00 2001 From: Richard Li Date: Thu, 8 May 2025 12:21:02 -0700 Subject: [PATCH 4/5] lint --- .../amazonq/lsp/flareChat/ChatCommunicationManager.kt | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/flareChat/ChatCommunicationManager.kt b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/flareChat/ChatCommunicationManager.kt index 089fb9ea53b..249e7c65b19 100644 --- a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/flareChat/ChatCommunicationManager.kt +++ b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/flareChat/ChatCommunicationManager.kt @@ -15,10 +15,10 @@ import software.aws.toolkits.jetbrains.core.credentials.pinning.QConnection import software.aws.toolkits.jetbrains.core.credentials.reauthConnectionIfNeeded 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.CHAT_ERROR_PARAMS -import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.ErrorParams 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.CHAT_ERROR_PARAMS +import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.ErrorParams 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 From a44fc7068a56ac6b09f8629684f56c163f9f3527 Mon Sep 17 00:00:00 2001 From: Richard Li Date: Thu, 8 May 2025 12:34:19 -0700 Subject: [PATCH 5/5] lint --- .../jetbrains/services/amazonq/lsp/AmazonQLspService.kt | 2 -- 1 file changed, 2 deletions(-) diff --git a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/AmazonQLspService.kt b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/AmazonQLspService.kt index 8b38823145c..54a9d97573b 100644 --- a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/AmazonQLspService.kt +++ b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/AmazonQLspService.kt @@ -19,7 +19,6 @@ import com.intellij.openapi.project.Project import com.intellij.openapi.util.Disposer import com.intellij.openapi.util.Key import com.intellij.openapi.util.SystemInfo -import com.intellij.util.animation.consumer import com.intellij.util.io.await import kotlinx.coroutines.CoroutineScope import kotlinx.coroutines.Deferred @@ -40,7 +39,6 @@ import org.eclipse.lsp4j.SynchronizationCapabilities import org.eclipse.lsp4j.TextDocumentClientCapabilities import org.eclipse.lsp4j.WorkspaceClientCapabilities import org.eclipse.lsp4j.jsonrpc.Launcher -import org.eclipse.lsp4j.jsonrpc.Launcher.Builder import org.eclipse.lsp4j.jsonrpc.MessageConsumer import org.eclipse.lsp4j.jsonrpc.messages.ResponseMessage import org.eclipse.lsp4j.launch.LSPLauncher