From a896c16e4927b8f6ed9fee9d1c56d7235a6c2bf6 Mon Sep 17 00:00:00 2001 From: Richard Li Date: Wed, 7 May 2025 19:11:01 -0700 Subject: [PATCH 01/12] fix(amazonq): fix lsp warning about unsupported 'aws/chat/sendContextCommands' --- .../services/amazonq/lsp/AmazonQLanguageClient.kt | 4 ++++ .../amazonq/lsp/AmazonQLanguageClientImpl.kt | 14 ++++++++++++++ .../services/amazonq/lsp/AmazonQLspService.kt | 13 +------------ 3 files changed, 19 insertions(+), 12 deletions(-) diff --git a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/AmazonQLanguageClient.kt b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/AmazonQLanguageClient.kt index 5cb296926f4..4a63ad7c0c7 100644 --- a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/AmazonQLanguageClient.kt +++ b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/AmazonQLanguageClient.kt @@ -6,6 +6,7 @@ package software.aws.toolkits.jetbrains.services.amazonq.lsp import org.eclipse.lsp4j.jsonrpc.services.JsonNotification import org.eclipse.lsp4j.jsonrpc.services.JsonRequest import org.eclipse.lsp4j.services.LanguageClient +import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.LSPAny import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.CHAT_SEND_UPDATE import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.ChatUpdateParams import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.GET_SERIALIZED_CHAT_REQUEST_METHOD @@ -43,4 +44,7 @@ interface AmazonQLanguageClient : LanguageClient { @JsonNotification(OPEN_FILE_DIFF) fun openFileDiff(params: OpenFileDiffParams): CompletableFuture + + @JsonNotification("aws/chat/sendContextCommands") + fun sendContextCommands(params: LSPAny): CompletableFuture } diff --git a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/AmazonQLanguageClientImpl.kt b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/AmazonQLanguageClientImpl.kt index f9ba4b44bd1..dd54010267a 100644 --- a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/AmazonQLanguageClientImpl.kt +++ b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/AmazonQLanguageClientImpl.kt @@ -31,6 +31,7 @@ import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManager import software.aws.toolkits.jetbrains.core.credentials.pinning.QConnection import software.aws.toolkits.jetbrains.services.amazonq.lsp.flareChat.AsyncChatUiListener import software.aws.toolkits.jetbrains.services.amazonq.lsp.flareChat.ChatCommunicationManager +import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.LSPAny import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.CHAT_OPEN_TAB import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.CHAT_SEND_UPDATE import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.ChatUpdateParams @@ -308,6 +309,19 @@ class AmazonQLanguageClientImpl(private val project: Project) : AmazonQLanguageC ApplicationManager.getApplication()::invokeLater ) + override fun sendContextCommands(params: LSPAny): CompletableFuture { + val showContextCommands = """ + { + "command":"aws/chat/sendContextCommands", + "params": ${Gson().toJson(params)} + } + """.trimIndent() + + AsyncChatUiListener.notifyPartialMessageUpdate(showContextCommands) + + return CompletableFuture.completedFuture(Unit) + } + companion object { private val LOG = getLogger() } 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 7de94a494f7..b8e29cac64b 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 @@ -282,23 +282,12 @@ private class AmazonQServerInstance(private val project: Project, private val cs launcher = LSPLauncher.Builder() .wrapMessages { consumer -> - MessageConsumer { - message -> + MessageConsumer { message -> if (message is ResponseMessage && message.result is AwsExtendedInitializeResult) { val result = message.result as AwsExtendedInitializeResult AwsServerCapabilitiesProvider.getInstance(project).setAwsServerCapabilities(result.getAwsServerCapabilities()) AmazonQLspService.getInstance(project).notifyInitializeMessageReceived() } - if (message is NotificationMessage && message.method == "aws/chat/sendContextCommands") { - val showContextCommands = """ - { - "command":"aws/chat/sendContextCommands", - "params": ${Gson().toJson(message.params)} - } - """.trimIndent() - - AsyncChatUiListener.notifyPartialMessageUpdate(showContextCommands) - } consumer?.consume(message) } } From 5ef38f7c4ab9ed7ce2f00242a65d71771e6de47e Mon Sep 17 00:00:00 2001 From: Richard Li Date: Wed, 7 May 2025 19:16:12 -0700 Subject: [PATCH 02/12] sort --- .../amazonq/lsp/AmazonQLanguageClient.kt | 3 +- .../amazonq/lsp/AmazonQLanguageClientImpl.kt | 3 +- .../lsp/model/aws/chat/FlareChatCommands.kt | 49 ++++++++++--------- 3 files changed, 31 insertions(+), 24 deletions(-) diff --git a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/AmazonQLanguageClient.kt b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/AmazonQLanguageClient.kt index 4a63ad7c0c7..fe00aa835d6 100644 --- a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/AmazonQLanguageClient.kt +++ b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/AmazonQLanguageClient.kt @@ -7,6 +7,7 @@ import org.eclipse.lsp4j.jsonrpc.services.JsonNotification import org.eclipse.lsp4j.jsonrpc.services.JsonRequest import org.eclipse.lsp4j.services.LanguageClient import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.LSPAny +import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.CHAT_SEND_CONTEXT_COMMANDS import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.CHAT_SEND_UPDATE import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.ChatUpdateParams import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.GET_SERIALIZED_CHAT_REQUEST_METHOD @@ -45,6 +46,6 @@ interface AmazonQLanguageClient : LanguageClient { @JsonNotification(OPEN_FILE_DIFF) fun openFileDiff(params: OpenFileDiffParams): CompletableFuture - @JsonNotification("aws/chat/sendContextCommands") + @JsonNotification(CHAT_SEND_CONTEXT_COMMANDS) fun sendContextCommands(params: LSPAny): CompletableFuture } diff --git a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/AmazonQLanguageClientImpl.kt b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/AmazonQLanguageClientImpl.kt index dd54010267a..61ec067ab87 100644 --- a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/AmazonQLanguageClientImpl.kt +++ b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/AmazonQLanguageClientImpl.kt @@ -33,6 +33,7 @@ import software.aws.toolkits.jetbrains.services.amazonq.lsp.flareChat.AsyncChatU import software.aws.toolkits.jetbrains.services.amazonq.lsp.flareChat.ChatCommunicationManager import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.LSPAny import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.CHAT_OPEN_TAB +import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.CHAT_SEND_CONTEXT_COMMANDS import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.CHAT_SEND_UPDATE import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.ChatUpdateParams import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.GET_SERIALIZED_CHAT_REQUEST_METHOD @@ -312,7 +313,7 @@ class AmazonQLanguageClientImpl(private val project: Project) : AmazonQLanguageC override fun sendContextCommands(params: LSPAny): CompletableFuture { val showContextCommands = """ { - "command":"aws/chat/sendContextCommands", + "command":"$CHAT_SEND_CONTEXT_COMMANDS", "params": ${Gson().toJson(params)} } """.trimIndent() 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 6b7d5a1f3e5..497c6db1430 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 @@ -3,34 +3,39 @@ package software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat -const val SEND_CHAT_COMMAND_PROMPT = "aws/chat/sendChatPrompt" -const val CHAT_QUICK_ACTION = "aws/chat/sendChatQuickAction" -const val CHAT_READY = "aws/chat/ready" -const val CHAT_LINK_CLICK = "aws/chat/linkClick" -const val CHAT_INFO_LINK_CLICK = "aws/chat/infoLinkClick" -const val CHAT_SOURCE_LINK_CLICK = "aws/chat/sourceLinkClick" -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 CHAT_LIST_CONVERSATIONS = "aws/chat/listConversations" -const val CHAT_CONVERSATION_CLICK = "aws/chat/conversationClick" -const val CHAT_FILE_CLICK = "aws/chat/fileClick" -const val CHAT_TAB_ADD = "aws/chat/tabAdd" -const val CHAT_TAB_CHANGE = "aws/chat/tabChange" -const val CHAT_TAB_REMOVE = "aws/chat/tabRemove" -const val CHAT_OPEN_TAB = "aws/chat/openTab" const val CHAT_BUTTON_CLICK = "aws/chat/buttonClick" -const val CHAT_DISCLAIMER_ACKNOWLEDGED = "disclaimerAcknowledged" +const val CHAT_CONVERSATION_CLICK = "aws/chat/conversationClick" const val CHAT_COPY_CODE_TO_CLIPBOARD = "copyToClipboard" const val CHAT_COPY_CODE_TO_CLIPBOARD_NOTIFICATION = "aws/chat/copyCodeToClipboard" +const val CHAT_CREATE_PROMPT = "aws/chat/createPrompt" +const val CHAT_DISCLAIMER_ACKNOWLEDGED = "disclaimerAcknowledged" +const val CHAT_ERROR_PARAMS = "errorMessage" +const val CHAT_FEEDBACK = "aws/chat/feedback" +const val CHAT_FILE_CLICK = "aws/chat/fileClick" +const val CHAT_FOLLOW_UP_CLICK = "aws/chat/followUpClick" +const val CHAT_INFO_LINK_CLICK = "aws/chat/infoLinkClick" const val CHAT_INSERT_TO_CURSOR = "insertToCursorPosition" const val CHAT_INSERT_TO_CURSOR_NOTIFICATION = "aws/chat/insertToCursorPosition" -const val CHAT_TAB_BAR_ACTIONS = "aws/chat/tabBarAction" +const val CHAT_LINK_CLICK = "aws/chat/linkClick" +const val CHAT_LIST_CONVERSATIONS = "aws/chat/listConversations" +const val CHAT_OPEN_TAB = "aws/chat/openTab" +const val CHAT_PROMPT_OPTION_ACKNOWLEDGED = "chatPromptOptionAcknowledged" +const val CHAT_QUICK_ACTION = "aws/chat/sendChatQuickAction" +const val CHAT_READY = "aws/chat/ready" +const val CHAT_SEND_CONTEXT_COMMANDS = "aws/chat/sendContextCommands" const val CHAT_SEND_UPDATE = "aws/chat/sendChatUpdate" -const val CHAT_CREATE_PROMPT = "aws/chat/createPrompt" -const val SHOW_SAVE_FILE_DIALOG_REQUEST_METHOD = "aws/showSaveFileDialog" +const val CHAT_SOURCE_LINK_CLICK = "aws/chat/sourceLinkClick" +const val CHAT_TAB_ADD = "aws/chat/tabAdd" +const val CHAT_TAB_BAR_ACTIONS = "aws/chat/tabBarAction" +const val CHAT_TAB_CHANGE = "aws/chat/tabChange" +const val CHAT_TAB_REMOVE = "aws/chat/tabRemove" + const val GET_SERIALIZED_CHAT_REQUEST_METHOD = "aws/chat/getSerializedChat" + const val OPEN_FILE_DIFF = "aws/openFileDiff" -const val CHAT_ERROR_PARAMS = "errorMessage" + +const val PROMPT_INPUT_OPTIONS_CHANGE = "aws/chat/promptInputOptionChange" + +const val SEND_CHAT_COMMAND_PROMPT = "aws/chat/sendChatPrompt" +const val SHOW_SAVE_FILE_DIALOG_REQUEST_METHOD = "aws/showSaveFileDialog" const val STOP_CHAT_RESPONSE = "stopChatResponse" From 916eabe3456312cf9e3bd4320ccd3b6d986eb240 Mon Sep 17 00:00:00 2001 From: Richard Li Date: Wed, 7 May 2025 19:46:02 -0700 Subject: [PATCH 03/12] refactor: avoid manually crafting JSON blobs for postMessage(...) --- .../amazonq/toolwindow/AmazonQPanel.kt | 11 +++- .../services/amazonq/webview/Browser.kt | 5 ++ .../amazonq/webview/BrowserConnector.kt | 24 ++++--- .../amazonq/lsp/AmazonQLanguageClientImpl.kt | 63 +++++++++---------- .../lsp/flareChat/AsyncChatUiListener.kt | 14 ++++- .../lsp/flareChat/ChatCommunicationManager.kt | 8 --- .../amazonq/lsp/flareChat/FlareUiMessage.kt | 10 +++ 7 files changed, 79 insertions(+), 56 deletions(-) create mode 100644 plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/flareChat/FlareUiMessage.kt diff --git a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/toolwindow/AmazonQPanel.kt b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/toolwindow/AmazonQPanel.kt index 41487d0a658..2dead56fa7f 100644 --- a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/toolwindow/AmazonQPanel.kt +++ b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/toolwindow/AmazonQPanel.kt @@ -27,6 +27,7 @@ import software.aws.toolkits.jetbrains.services.amazonq.apps.AppConnection import software.aws.toolkits.jetbrains.services.amazonq.commands.MessageTypeRegistry import software.aws.toolkits.jetbrains.services.amazonq.lsp.artifacts.ArtifactManager import software.aws.toolkits.jetbrains.services.amazonq.lsp.flareChat.AsyncChatUiListener +import software.aws.toolkits.jetbrains.services.amazonq.lsp.flareChat.FlareUiMessage import software.aws.toolkits.jetbrains.services.amazonq.messages.AmazonQMessage import software.aws.toolkits.jetbrains.services.amazonq.messages.MessageConnector import software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfileManager @@ -55,9 +56,15 @@ class AmazonQPanel(val project: Project, private val scope: CoroutineScope) : Di project.messageBus.connect().subscribe( AsyncChatUiListener.TOPIC, object : AsyncChatUiListener { - override fun onChange(message: String) { + override fun onChange(command: String) { runInEdt { - browser.get()?.postChat(message) + browser.get()?.postChat(command) + } + } + + override fun onChange(command: FlareUiMessage) { + runInEdt { + browser.get()?.postChat(command) } } } diff --git a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/webview/Browser.kt b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/webview/Browser.kt index 6490a145721..bc686077a0f 100644 --- a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/webview/Browser.kt +++ b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/webview/Browser.kt @@ -4,6 +4,7 @@ package software.aws.toolkits.jetbrains.services.amazonq.webview import com.fasterxml.jackson.module.kotlin.jacksonObjectMapper +import com.google.gson.Gson import com.intellij.openapi.Disposable import com.intellij.openapi.project.Project import com.intellij.openapi.util.Disposer @@ -11,6 +12,7 @@ import com.intellij.ui.jcef.JBCefJSQuery import org.cef.CefApp import software.aws.toolkits.jetbrains.services.amazonq.lsp.AmazonQLspService import software.aws.toolkits.jetbrains.services.amazonq.lsp.flareChat.AwsServerCapabilitiesProvider +import software.aws.toolkits.jetbrains.services.amazonq.lsp.flareChat.FlareUiMessage import software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfile import software.aws.toolkits.jetbrains.services.amazonq.util.HighlightCommand import software.aws.toolkits.jetbrains.services.amazonq.util.createBrowser @@ -62,6 +64,9 @@ class Browser(parent: Disposable, private val webUri: URI, val project: Project) fun component() = jcefBrowser.component + fun postChat(command: FlareUiMessage) = postChat(Gson().toJson(command)) + + @Deprecated("shouldn't need this version") fun postChat(message: String) { jcefBrowser .cefBrowser 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 f9b19925178..d22a42c76ce 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 @@ -32,6 +32,7 @@ import software.aws.toolkits.jetbrains.services.amazonq.lsp.AmazonQLspService import software.aws.toolkits.jetbrains.services.amazonq.lsp.encryption.JwtEncryptionManager 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.FlareUiMessage import software.aws.toolkits.jetbrains.services.amazonq.lsp.flareChat.getTextDocumentIdentifier import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.ButtonClickNotification import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.ButtonClickParams @@ -256,7 +257,9 @@ class BrowserConnector( var encryptionManager: JwtEncryptionManager? = null val result = AmazonQLspService.executeIfRunning(project) { server -> encryptionManager = this.encryptionManager - encryptionManager?.encrypt(chatParams)?.let { EncryptedChatParams(it, partialResultToken) }?.let { server.sendChatPrompt(it) } + + val encryptedParams = EncryptedChatParams(this.encryptionManager.encrypt(chatParams), partialResultToken) + server.sendChatPrompt(encryptedParams) } ?: (CompletableFuture.failedFuture(IllegalStateException("LSP Server not running"))) // We assume there is only one outgoing request per tab because the input is @@ -272,11 +275,9 @@ class BrowserConnector( var encryptionManager: JwtEncryptionManager? = null val result = AmazonQLspService.executeIfRunning(project) { server -> encryptionManager = this.encryptionManager - encryptionManager?.encrypt(quickActionParams)?.let { - EncryptedQuickActionChatParams(it, partialResultToken) - }?.let { - server.sendQuickAction(it) - } + + val encryptedParams = EncryptedQuickActionChatParams(this.encryptionManager.encrypt(quickActionParams), partialResultToken) + server.sendQuickAction(encryptedParams) } ?: (CompletableFuture.failedFuture(IllegalStateException("LSP Server not running"))) // We assume there is only one outgoing request per tab because the input is @@ -424,13 +425,18 @@ class BrowserConnector( handleChatNotification(node) { server, params -> val result = server.tabBarActions(params) - result.whenComplete { params1, error -> + result.whenComplete { actions, error -> try { if (error != null) { throw error } - val res = ChatCommunicationManager.convertNotificationToJsonForChat(CHAT_TAB_BAR_ACTIONS, params1) - browser.postChat(res) + + browser.postChat( + FlareUiMessage( + command = CHAT_TAB_BAR_ACTIONS, + params = actions ?: "{}" + ) + ) } catch (e: Exception) { LOG.error { "Failed to perform chat tab bar action $e" } params.tabId?.let { diff --git a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/AmazonQLanguageClientImpl.kt b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/AmazonQLanguageClientImpl.kt index 61ec067ab87..515132d8adf 100644 --- a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/AmazonQLanguageClientImpl.kt +++ b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/AmazonQLanguageClientImpl.kt @@ -3,7 +3,6 @@ @file:Suppress("BannedImports") package software.aws.toolkits.jetbrains.services.amazonq.lsp -import com.google.gson.Gson import com.intellij.diff.DiffContentFactory import com.intellij.diff.DiffManager import com.intellij.diff.requests.SimpleDiffRequest @@ -24,6 +23,7 @@ import org.eclipse.lsp4j.PublishDiagnosticsParams import org.eclipse.lsp4j.ShowDocumentParams import org.eclipse.lsp4j.ShowDocumentResult import org.eclipse.lsp4j.ShowMessageRequestParams +import software.aws.toolkits.core.utils.error import software.aws.toolkits.core.utils.getLogger import software.aws.toolkits.core.utils.warn import software.aws.toolkits.jetbrains.core.credentials.AwsBearerTokenConnection @@ -31,6 +31,7 @@ import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManager import software.aws.toolkits.jetbrains.core.credentials.pinning.QConnection import software.aws.toolkits.jetbrains.services.amazonq.lsp.flareChat.AsyncChatUiListener import software.aws.toolkits.jetbrains.services.amazonq.lsp.flareChat.ChatCommunicationManager +import software.aws.toolkits.jetbrains.services.amazonq.lsp.flareChat.FlareUiMessage import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.LSPAny import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.CHAT_OPEN_TAB import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.CHAT_SEND_CONTEXT_COMMANDS @@ -135,14 +136,13 @@ class AmazonQLanguageClientImpl(private val project: Project) : AmazonQLanguageC val result = CompletableFuture() ChatCommunicationManager.pendingTabRequests[requestId] = result - val uiMessage = """ - { - "command": "$CHAT_OPEN_TAB", - "params": ${Gson().toJson(params)}, - "requestId": "$requestId" - } - """.trimIndent() - AsyncChatUiListener.notifyPartialMessageUpdate(uiMessage) + AsyncChatUiListener.notifyPartialMessageUpdate( + FlareUiMessage( + command = CHAT_OPEN_TAB, + params = params, + requestId = requestId, + ) + ) result.orTimeout(30000, TimeUnit.MILLISECONDS) .whenComplete { _, error -> @@ -188,14 +188,13 @@ class AmazonQLanguageClientImpl(private val project: Project) : AmazonQLanguageC ChatCommunicationManager.pendingSerializedChatRequests[requestId] = result - val uiMessage = """ - { - "command": "$GET_SERIALIZED_CHAT_REQUEST_METHOD", - "params": ${Gson().toJson(params)}, - "requestId": "$requestId" - } - """.trimIndent() - AsyncChatUiListener.notifyPartialMessageUpdate(uiMessage) + AsyncChatUiListener.notifyPartialMessageUpdate( + FlareUiMessage( + command = GET_SERIALIZED_CHAT_REQUEST_METHOD, + params = params, + requestId = requestId, + ) + ) result.orTimeout(30000, TimeUnit.MILLISECONDS) .whenComplete { _, error -> @@ -254,19 +253,17 @@ class AmazonQLanguageClientImpl(private val project: Project) : AmazonQLanguageC try { chatCommunicationManager.handlePartialResultProgressNotification(project, params) } catch (e: Exception) { - error("Cannot handle partial chat") + LOG.error(e) { "Cannot handle partial chat" } } } override fun sendChatUpdate(params: ChatUpdateParams): CompletableFuture { - val uiMessage = """ - { - "command":"$CHAT_SEND_UPDATE", - "params":${Gson().toJson(params)} - } - """.trimIndent() - - AsyncChatUiListener.notifyPartialMessageUpdate(uiMessage) + AsyncChatUiListener.notifyPartialMessageUpdate( + FlareUiMessage( + command = CHAT_SEND_UPDATE, + params = params, + ) + ) return CompletableFuture.completedFuture(Unit) } @@ -311,14 +308,12 @@ class AmazonQLanguageClientImpl(private val project: Project) : AmazonQLanguageC ) override fun sendContextCommands(params: LSPAny): CompletableFuture { - val showContextCommands = """ - { - "command":"$CHAT_SEND_CONTEXT_COMMANDS", - "params": ${Gson().toJson(params)} - } - """.trimIndent() - - AsyncChatUiListener.notifyPartialMessageUpdate(showContextCommands) + AsyncChatUiListener.notifyPartialMessageUpdate( + FlareUiMessage( + command = CHAT_SEND_CONTEXT_COMMANDS, + params = params ?: error("received empty payload for $CHAT_SEND_CONTEXT_COMMANDS"), + ) + ) return CompletableFuture.completedFuture(Unit) } diff --git a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/flareChat/AsyncChatUiListener.kt b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/flareChat/AsyncChatUiListener.kt index 9aa538be924..c6992526ec9 100644 --- a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/flareChat/AsyncChatUiListener.kt +++ b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/flareChat/AsyncChatUiListener.kt @@ -8,14 +8,22 @@ import com.intellij.util.messages.Topic import java.util.EventListener interface AsyncChatUiListener : EventListener { - fun onChange(message: String) {} + @Deprecated("shouldn't need this version") + fun onChange(command: String) {} + + fun onChange(command: FlareUiMessage) {} companion object { @Topic.AppLevel val TOPIC = Topic.create("Partial chat message provider", AsyncChatUiListener::class.java) - fun notifyPartialMessageUpdate(message: String) { - ApplicationManager.getApplication().messageBus.syncPublisher(TOPIC).onChange(message) + fun notifyPartialMessageUpdate(command: FlareUiMessage) { + ApplicationManager.getApplication().messageBus.syncPublisher(TOPIC).onChange(command) + } + + @Deprecated("shouldn't need this version") + fun notifyPartialMessageUpdate(command: String) { + ApplicationManager.getApplication().messageBus.syncPublisher(TOPIC).onChange(command) } } } 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 67116573e1f..223bf0235ee 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 @@ -114,13 +114,5 @@ class ChatCommunicationManager { fun completeTabOpen(requestId: String, tabId: String) { pendingTabRequests.remove(requestId)?.complete(OpenTabResult(tabId)) } - - inline fun convertNotificationToJsonForChat(command: String, params: T? = null) = - """ - { - "command":"$command", - "params": ${if (params != null) Gson().toJson(params) else "{}"} - } - """.trimIndent() } } diff --git a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/flareChat/FlareUiMessage.kt b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/flareChat/FlareUiMessage.kt new file mode 100644 index 00000000000..a9ede9941e8 --- /dev/null +++ b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/flareChat/FlareUiMessage.kt @@ -0,0 +1,10 @@ +// 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.flareChat + +data class FlareUiMessage( + val command: String, + val params: Any, + val requestId: String? = null +) From 04daaaecd435169adfe255bb77aaea18b2bdc120 Mon Sep 17 00:00:00 2001 From: Richard Li <742829+rli@users.noreply.github.com> Date: Wed, 7 May 2025 19:49:08 -0700 Subject: [PATCH 04/12] Update BrowserConnector.kt --- .../jetbrains/services/amazonq/webview/BrowserConnector.kt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 d22a42c76ce..157c0c2d01a 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 @@ -434,7 +434,7 @@ class BrowserConnector( browser.postChat( FlareUiMessage( command = CHAT_TAB_BAR_ACTIONS, - params = actions ?: "{}" + params = actions ?: emptyMap() ) ) } catch (e: Exception) { From 0b2aaaffd8aae760a990ea7eee69a1e26d6178e1 Mon Sep 17 00:00:00 2001 From: samgst-amazon Date: Thu, 8 May 2025 12:22:12 -0700 Subject: [PATCH 05/12] queue up contextCommands --- .../amazonq/webview/BrowserConnector.kt | 1 + .../amazonq/lsp/AmazonQLanguageClientImpl.kt | 9 +--- .../services/amazonq/lsp/AmazonQLspService.kt | 3 -- .../lsp/flareChat/ChatCommunicationManager.kt | 47 ++++++++++++++++--- .../amazonq/lsp/flareChat/FlareUiMessage.kt | 2 +- 5 files changed, 44 insertions(+), 18 deletions(-) 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 157c0c2d01a..ae2d844da17 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 @@ -326,6 +326,7 @@ class BrowserConnector( CHAT_READY -> { handleChatNotification(node) { server, _ -> uiReady.complete(true) + chatCommunicationManager.setUiReady(true) RunOnceUtil.runOnceForApp("AmazonQ-UI-Ready") { MeetQSettings.getInstance().reinvent2024OnboardingCount += 1 } diff --git a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/AmazonQLanguageClientImpl.kt b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/AmazonQLanguageClientImpl.kt index 515132d8adf..ebef5f3ddd4 100644 --- a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/AmazonQLanguageClientImpl.kt +++ b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/AmazonQLanguageClientImpl.kt @@ -34,7 +34,6 @@ import software.aws.toolkits.jetbrains.services.amazonq.lsp.flareChat.ChatCommun import software.aws.toolkits.jetbrains.services.amazonq.lsp.flareChat.FlareUiMessage import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.LSPAny import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.CHAT_OPEN_TAB -import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.CHAT_SEND_CONTEXT_COMMANDS import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.CHAT_SEND_UPDATE import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.ChatUpdateParams import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.GET_SERIALIZED_CHAT_REQUEST_METHOD @@ -308,13 +307,7 @@ class AmazonQLanguageClientImpl(private val project: Project) : AmazonQLanguageC ) override fun sendContextCommands(params: LSPAny): CompletableFuture { - AsyncChatUiListener.notifyPartialMessageUpdate( - FlareUiMessage( - command = CHAT_SEND_CONTEXT_COMMANDS, - params = params ?: error("received empty payload for $CHAT_SEND_CONTEXT_COMMANDS"), - ) - ) - + ChatCommunicationManager.getInstance(project).queueContextCommands(params) return CompletableFuture.completedFuture(Unit) } 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 b8e29cac64b..97c56e80cdb 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 @@ -3,7 +3,6 @@ @file:Suppress("BannedImports") package software.aws.toolkits.jetbrains.services.amazonq.lsp -import com.google.gson.Gson import com.google.gson.ToNumberPolicy import com.intellij.execution.configurations.GeneralCommandLine import com.intellij.execution.impl.ExecutionManagerImpl @@ -44,7 +43,6 @@ 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.NotificationMessage import org.eclipse.lsp4j.jsonrpc.messages.ResponseMessage import org.eclipse.lsp4j.launch.LSPLauncher import org.slf4j.event.Level @@ -57,7 +55,6 @@ import software.aws.toolkits.jetbrains.services.amazonq.lsp.auth.DefaultAuthCred import software.aws.toolkits.jetbrains.services.amazonq.lsp.dependencies.DefaultModuleDependenciesService import software.aws.toolkits.jetbrains.services.amazonq.lsp.encryption.JwtEncryptionManager import software.aws.toolkits.jetbrains.services.amazonq.lsp.flareChat.AmazonQLspTypeAdapterFactory -import software.aws.toolkits.jetbrains.services.amazonq.lsp.flareChat.AsyncChatUiListener import software.aws.toolkits.jetbrains.services.amazonq.lsp.flareChat.AwsExtendedInitializeResult import software.aws.toolkits.jetbrains.services.amazonq.lsp.flareChat.AwsServerCapabilitiesProvider import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.createExtendedClientMetadata 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 223bf0235ee..c8e0f3b978a 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 @@ -10,7 +10,9 @@ import com.intellij.openapi.project.Project import org.eclipse.lsp4j.ProgressParams 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.LSPAny 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.CHAT_SEND_CONTEXT_COMMANDS 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 @@ -18,10 +20,14 @@ import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.SEND_ import java.util.UUID import java.util.concurrent.CompletableFuture import java.util.concurrent.ConcurrentHashMap +import java.util.concurrent.ConcurrentLinkedQueue @Service(Service.Level.PROJECT) class ChatCommunicationManager { + private val contextCommandsQueue = ConcurrentLinkedQueue() + private var uiReady = false private val chatPartialResultMap = ConcurrentHashMap() + private fun getPartialChatMessage(partialResultToken: String): String? = chatPartialResultMap.getOrDefault(partialResultToken, null) @@ -47,6 +53,36 @@ class ChatCommunicationManager { fun removePartialChatMessage(partialResultToken: String) = chatPartialResultMap.remove(partialResultToken) + fun setUiReady(ready: Boolean) { + uiReady = ready + if (ready) { + sendQueuedContextCommands() + } + } + + fun queueContextCommands(params: LSPAny) { + if (uiReady) { + sendContextCommandsToUi(params) + } else { + contextCommandsQueue.offer(params) + } + } + + private fun sendQueuedContextCommands() { + while (contextCommandsQueue.isNotEmpty()) { + sendContextCommandsToUi(contextCommandsQueue.poll()) + } + } + + private fun sendContextCommandsToUi(params: LSPAny) { + AsyncChatUiListener.notifyPartialMessageUpdate( + FlareUiMessage( + command = CHAT_SEND_CONTEXT_COMMANDS, + params = params ?: error("received empty payload for $CHAT_SEND_CONTEXT_COMMANDS"), + ) + ) + } + fun handlePartialResultProgressNotification(project: Project, params: ProgressParams) { val token = ProgressNotificationUtils.getToken(params) val tabId = getPartialChatMessage(token) @@ -94,11 +130,6 @@ class ChatCommunicationManager { companion object { fun getInstance(project: Project) = project.service() - val pendingSerializedChatRequests = ConcurrentHashMap>() - fun completeSerializedChatResponse(requestId: String, content: String) { - pendingSerializedChatRequests.remove(requestId)?.complete(GetSerializedChatResult((content))) - } - fun convertToJsonToSendToChat(command: String, tabId: String, params: String, isPartialResult: Boolean): String = """ { @@ -109,8 +140,12 @@ class ChatCommunicationManager { } """.trimIndent() - val pendingTabRequests = ConcurrentHashMap>() + val pendingSerializedChatRequests = ConcurrentHashMap>() + fun completeSerializedChatResponse(requestId: String, content: String) { + pendingSerializedChatRequests.remove(requestId)?.complete(GetSerializedChatResult((content))) + } + val pendingTabRequests = ConcurrentHashMap>() fun completeTabOpen(requestId: String, tabId: String) { pendingTabRequests.remove(requestId)?.complete(OpenTabResult(tabId)) } diff --git a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/flareChat/FlareUiMessage.kt b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/flareChat/FlareUiMessage.kt index a9ede9941e8..7bc5fe2461f 100644 --- a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/flareChat/FlareUiMessage.kt +++ b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/flareChat/FlareUiMessage.kt @@ -6,5 +6,5 @@ package software.aws.toolkits.jetbrains.services.amazonq.lsp.flareChat data class FlareUiMessage( val command: String, val params: Any, - val requestId: String? = null + val requestId: String? = null, ) From 25cdb75a8893c69a9e77c9e6ad0dfaaa266c2019 Mon Sep 17 00:00:00 2001 From: samgst-amazon Date: Thu, 8 May 2025 12:33:25 -0700 Subject: [PATCH 06/12] detekt --- .../aws/toolkits/jetbrains/services/amazonq/webview/Browser.kt | 1 + 1 file changed, 1 insertion(+) diff --git a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/webview/Browser.kt b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/webview/Browser.kt index bc686077a0f..3d56168d187 100644 --- a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/webview/Browser.kt +++ b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/webview/Browser.kt @@ -1,6 +1,7 @@ // Copyright 2023 Amazon.com, Inc. or its affiliates. All Rights Reserved. // SPDX-License-Identifier: Apache-2.0 +@file:Suppress("BannedImports") package software.aws.toolkits.jetbrains.services.amazonq.webview import com.fasterxml.jackson.module.kotlin.jacksonObjectMapper From 195e0fc440727746df5543e3dd63183827524074 Mon Sep 17 00:00:00 2001 From: Richard Li <742829+rli@users.noreply.github.com> Date: Thu, 8 May 2025 14:45:49 -0700 Subject: [PATCH 07/12] Update AmazonQPanel.kt --- .../jetbrains/services/amazonq/toolwindow/AmazonQPanel.kt | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/toolwindow/AmazonQPanel.kt b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/toolwindow/AmazonQPanel.kt index 2dead56fa7f..92e81d376b1 100644 --- a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/toolwindow/AmazonQPanel.kt +++ b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/toolwindow/AmazonQPanel.kt @@ -57,15 +57,11 @@ class AmazonQPanel(val project: Project, private val scope: CoroutineScope) : Di AsyncChatUiListener.TOPIC, object : AsyncChatUiListener { override fun onChange(command: String) { - runInEdt { - browser.get()?.postChat(command) - } + browser.get()?.postChat(command) } override fun onChange(command: FlareUiMessage) { - runInEdt { - browser.get()?.postChat(command) - } + browser.get()?.postChat(command) } } ) From 441f516437449726b990e81753805e5fc4d3f2e5 Mon Sep 17 00:00:00 2001 From: samgst-amazon Date: Thu, 8 May 2025 14:44:04 -0700 Subject: [PATCH 08/12] block messages until UI ready --- .../amazonq/webview/BrowserConnector.kt | 7 +- .../amazonq/lsp/AmazonQLanguageClientImpl.kt | 24 ++++-- .../lsp/flareChat/ChatCommunicationManager.kt | 75 +++++++++---------- 3 files changed, 55 insertions(+), 51 deletions(-) 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 ae2d844da17..f4500a79d64 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 @@ -132,6 +132,7 @@ class BrowserConnector( // this is sent when the named agents UI is ready "ui-is-ready" -> { uiReady.complete(true) + chatCommunicationManager.setUiReady() RunOnceUtil.runOnceForApp("AmazonQ-UI-Ready") { MeetQSettings.getInstance().reinvent2024OnboardingCount += 1 } @@ -326,7 +327,7 @@ class BrowserConnector( CHAT_READY -> { handleChatNotification(node) { server, _ -> uiReady.complete(true) - chatCommunicationManager.setUiReady(true) + chatCommunicationManager.setUiReady() RunOnceUtil.runOnceForApp("AmazonQ-UI-Ready") { MeetQSettings.getInstance().reinvent2024OnboardingCount += 1 } @@ -352,7 +353,7 @@ class BrowserConnector( } CHAT_OPEN_TAB -> { val response = serializer.deserializeChatMessages(node) - ChatCommunicationManager.completeTabOpen( + chatCommunicationManager.completeTabOpen( response.requestId, response.params.result.tabId ) @@ -416,7 +417,7 @@ class BrowserConnector( GET_SERIALIZED_CHAT_REQUEST_METHOD -> { val response = serializer.deserializeChatMessages(node) - ChatCommunicationManager.completeSerializedChatResponse( + chatCommunicationManager.completeSerializedChatResponse( response.requestId, response.params.result.content ) diff --git a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/AmazonQLanguageClientImpl.kt b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/AmazonQLanguageClientImpl.kt index ebef5f3ddd4..c777a94fa0c 100644 --- a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/AmazonQLanguageClientImpl.kt +++ b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/AmazonQLanguageClientImpl.kt @@ -34,6 +34,7 @@ import software.aws.toolkits.jetbrains.services.amazonq.lsp.flareChat.ChatCommun import software.aws.toolkits.jetbrains.services.amazonq.lsp.flareChat.FlareUiMessage import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.LSPAny import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.CHAT_OPEN_TAB +import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.CHAT_SEND_CONTEXT_COMMANDS import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.CHAT_SEND_UPDATE import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.ChatUpdateParams import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.GET_SERIALIZED_CHAT_REQUEST_METHOD @@ -133,9 +134,10 @@ class AmazonQLanguageClientImpl(private val project: Project) : AmazonQLanguageC override fun openTab(params: OpenTabParams): CompletableFuture { val requestId = UUID.randomUUID().toString() val result = CompletableFuture() - ChatCommunicationManager.pendingTabRequests[requestId] = result + val chatManager = ChatCommunicationManager.getInstance(project) + chatManager.addTabOpenRequest(requestId, result) - AsyncChatUiListener.notifyPartialMessageUpdate( + chatManager.notifyUi( FlareUiMessage( command = CHAT_OPEN_TAB, params = params, @@ -145,7 +147,7 @@ class AmazonQLanguageClientImpl(private val project: Project) : AmazonQLanguageC result.orTimeout(30000, TimeUnit.MILLISECONDS) .whenComplete { _, error -> - ChatCommunicationManager.pendingTabRequests.remove(requestId) + chatManager.removeTabOpenRequest(requestId) } return result @@ -184,10 +186,10 @@ class AmazonQLanguageClientImpl(private val project: Project) : AmazonQLanguageC override fun getSerializedChat(params: GetSerializedChatParams): CompletableFuture { val requestId = UUID.randomUUID().toString() val result = CompletableFuture() + val chatManager = ChatCommunicationManager.getInstance(project) + chatManager.addSerializedChatRequest(requestId, result) - ChatCommunicationManager.pendingSerializedChatRequests[requestId] = result - - AsyncChatUiListener.notifyPartialMessageUpdate( + chatManager.notifyUi( FlareUiMessage( command = GET_SERIALIZED_CHAT_REQUEST_METHOD, params = params, @@ -197,7 +199,7 @@ class AmazonQLanguageClientImpl(private val project: Project) : AmazonQLanguageC result.orTimeout(30000, TimeUnit.MILLISECONDS) .whenComplete { _, error -> - ChatCommunicationManager.pendingSerializedChatRequests.remove(requestId) + chatManager.removeSerializedChatRequest(requestId) } return result @@ -307,7 +309,13 @@ class AmazonQLanguageClientImpl(private val project: Project) : AmazonQLanguageC ) override fun sendContextCommands(params: LSPAny): CompletableFuture { - ChatCommunicationManager.getInstance(project).queueContextCommands(params) + val chatManager = ChatCommunicationManager.getInstance(project) + chatManager.notifyUi( + FlareUiMessage( + command = CHAT_SEND_CONTEXT_COMMANDS, + params = params ?: error("received empty payload for $CHAT_SEND_CONTEXT_COMMANDS"), + ) + ) return CompletableFuture.completedFuture(Unit) } 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 c8e0f3b978a..9f68ca92c88 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 @@ -7,12 +7,12 @@ import com.google.gson.Gson import com.intellij.openapi.components.Service import com.intellij.openapi.components.service import com.intellij.openapi.project.Project +import kotlinx.coroutines.CompletableDeferred +import kotlinx.coroutines.runBlocking import org.eclipse.lsp4j.ProgressParams 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.LSPAny 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.CHAT_SEND_CONTEXT_COMMANDS 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 @@ -20,18 +20,26 @@ import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.SEND_ import java.util.UUID import java.util.concurrent.CompletableFuture import java.util.concurrent.ConcurrentHashMap -import java.util.concurrent.ConcurrentLinkedQueue @Service(Service.Level.PROJECT) class ChatCommunicationManager { - private val contextCommandsQueue = ConcurrentLinkedQueue() - private var uiReady = false + val uiReady = CompletableDeferred() private val chatPartialResultMap = ConcurrentHashMap() + private val inflightRequestByTabId = ConcurrentHashMap>() + private val pendingSerializedChatRequests = ConcurrentHashMap>() + private val pendingTabRequests = ConcurrentHashMap>() - private fun getPartialChatMessage(partialResultToken: String): String? = - chatPartialResultMap.getOrDefault(partialResultToken, null) + fun setUiReady() { + uiReady.complete(true) + } - private val inflightRequestByTabId = ConcurrentHashMap>() + fun notifyUi(uiMessage: FlareUiMessage): CompletableFuture = + CompletableFuture.runAsync { + runBlocking { + uiReady.await() + AsyncChatUiListener.notifyPartialMessageUpdate(uiMessage) + } + }.thenApply { Unit } fun setInflightRequestForTab(tabId: String, result: CompletableFuture) { inflightRequestByTabId[tabId] = result @@ -50,37 +58,34 @@ class ChatCommunicationManager { return partialResultToken } + private fun getPartialChatMessage(partialResultToken: String): String? = + chatPartialResultMap.getOrDefault(partialResultToken, null) + fun removePartialChatMessage(partialResultToken: String) = chatPartialResultMap.remove(partialResultToken) - fun setUiReady(ready: Boolean) { - uiReady = ready - if (ready) { - sendQueuedContextCommands() - } + fun addSerializedChatRequest(requestId: String, result: CompletableFuture) { + pendingSerializedChatRequests[requestId] = result } - fun queueContextCommands(params: LSPAny) { - if (uiReady) { - sendContextCommandsToUi(params) - } else { - contextCommandsQueue.offer(params) - } + fun completeSerializedChatResponse(requestId: String, content: String) { + pendingSerializedChatRequests.remove(requestId)?.complete(GetSerializedChatResult((content))) } - private fun sendQueuedContextCommands() { - while (contextCommandsQueue.isNotEmpty()) { - sendContextCommandsToUi(contextCommandsQueue.poll()) - } + fun removeSerializedChatRequest(requestId: String) { + pendingSerializedChatRequests.remove(requestId) } - private fun sendContextCommandsToUi(params: LSPAny) { - AsyncChatUiListener.notifyPartialMessageUpdate( - FlareUiMessage( - command = CHAT_SEND_CONTEXT_COMMANDS, - params = params ?: error("received empty payload for $CHAT_SEND_CONTEXT_COMMANDS"), - ) - ) + fun addTabOpenRequest(requestId: String, result: CompletableFuture) { + pendingTabRequests[requestId] = result + } + + fun completeTabOpen(requestId: String, tabId: String) { + pendingTabRequests.remove(requestId)?.complete(OpenTabResult(tabId)) + } + + fun removeTabOpenRequest(requestId: String) { + pendingTabRequests.remove(requestId) } fun handlePartialResultProgressNotification(project: Project, params: ProgressParams) { @@ -139,15 +144,5 @@ class ChatCommunicationManager { "isPartialResult": $isPartialResult } """.trimIndent() - - val pendingSerializedChatRequests = ConcurrentHashMap>() - fun completeSerializedChatResponse(requestId: String, content: String) { - pendingSerializedChatRequests.remove(requestId)?.complete(GetSerializedChatResult((content))) - } - - val pendingTabRequests = ConcurrentHashMap>() - fun completeTabOpen(requestId: String, tabId: String) { - pendingTabRequests.remove(requestId)?.complete(OpenTabResult(tabId)) - } } } From b10aaaab460ba4f59534798733cc761877649cc0 Mon Sep 17 00:00:00 2001 From: Richard Li Date: Thu, 8 May 2025 15:09:38 -0700 Subject: [PATCH 09/12] banned --- .../aws/toolkits/jetbrains/services/amazonq/webview/Browser.kt | 2 +- .../jetbrains/services/amazonq/lsp/flareChat/FlareUiMessage.kt | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/webview/Browser.kt b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/webview/Browser.kt index 029d4fde4f8..91975a72284 100644 --- a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/webview/Browser.kt +++ b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/webview/Browser.kt @@ -1,6 +1,6 @@ // Copyright 2023 Amazon.com, Inc. or its affiliates. All Rights Reserved. // SPDX-License-Identifier: Apache-2.0 - +@file:Suppress("BannedImports") package software.aws.toolkits.jetbrains.services.amazonq.webview import com.fasterxml.jackson.module.kotlin.jacksonObjectMapper diff --git a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/flareChat/FlareUiMessage.kt b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/flareChat/FlareUiMessage.kt index a9ede9941e8..7bc5fe2461f 100644 --- a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/flareChat/FlareUiMessage.kt +++ b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/flareChat/FlareUiMessage.kt @@ -6,5 +6,5 @@ package software.aws.toolkits.jetbrains.services.amazonq.lsp.flareChat data class FlareUiMessage( val command: String, val params: Any, - val requestId: String? = null + val requestId: String? = null, ) From 9f8acbf978aa10d0ee422d7ad840ff36763d5750 Mon Sep 17 00:00:00 2001 From: samgst-amazon Date: Thu, 8 May 2025 16:11:29 -0700 Subject: [PATCH 10/12] coroutine --- .../lsp/flareChat/ChatCommunicationManager.kt | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 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 7d6853bf37a..95255ac02bd 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,7 +8,8 @@ import com.intellij.openapi.components.Service import com.intellij.openapi.components.service import com.intellij.openapi.project.Project import kotlinx.coroutines.CompletableDeferred -import kotlinx.coroutines.runBlocking +import kotlinx.coroutines.CoroutineScope +import kotlinx.coroutines.launch import org.eclipse.lsp4j.ProgressParams import software.aws.toolkits.core.utils.getLogger import software.aws.toolkits.core.utils.warn @@ -31,7 +32,7 @@ import java.util.concurrent.CompletableFuture import java.util.concurrent.ConcurrentHashMap @Service(Service.Level.PROJECT) -class ChatCommunicationManager { +class ChatCommunicationManager(private val cs: CoroutineScope) { val uiReady = CompletableDeferred() private val chatPartialResultMap = ConcurrentHashMap() private val inflightRequestByTabId = ConcurrentHashMap>() @@ -42,13 +43,12 @@ class ChatCommunicationManager { uiReady.complete(true) } - fun notifyUi(uiMessage: FlareUiMessage): CompletableFuture = - CompletableFuture.runAsync { - runBlocking { - uiReady.await() - AsyncChatUiListener.notifyPartialMessageUpdate(uiMessage) - } - }.thenApply { Unit } + fun notifyUi(uiMessage: FlareUiMessage) { + cs.launch { + uiReady.await() + AsyncChatUiListener.notifyPartialMessageUpdate(uiMessage) + } + } fun setInflightRequestForTab(tabId: String, result: CompletableFuture) { inflightRequestByTabId[tabId] = result From 526556a3fc08361c4634e775ce315ca30dc64060 Mon Sep 17 00:00:00 2001 From: Richard Li Date: Thu, 8 May 2025 18:13:16 -0700 Subject: [PATCH 11/12] resolve conflict --- .../amazonq/webview/BrowserConnector.kt | 23 ++++++------------- 1 file changed, 7 insertions(+), 16 deletions(-) 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 042609f6341..cd2926e5067 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 @@ -8,6 +8,7 @@ import com.google.gson.Gson import com.intellij.ide.BrowserUtil import com.intellij.ide.util.RunOnceUtil import com.intellij.openapi.application.runInEdt +import com.intellij.openapi.fileEditor.FileEditorManager import com.intellij.openapi.options.ShowSettingsUtil import com.intellij.openapi.project.Project import com.intellij.ui.jcef.JBCefJSQuery.Response @@ -22,8 +23,7 @@ import kotlinx.coroutines.flow.merge import kotlinx.coroutines.flow.onEach import kotlinx.coroutines.launch import org.cef.browser.CefBrowser -import org.eclipse.lsp4j.Position -import org.eclipse.lsp4j.Range +import org.eclipse.lsp4j.TextDocumentIdentifier import software.aws.toolkits.core.utils.error import software.aws.toolkits.core.utils.getLogger import software.aws.toolkits.core.utils.warn @@ -35,7 +35,6 @@ 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.FlareUiMessage -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 @@ -108,6 +107,8 @@ import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.TabBa import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.TabBarActionRequest import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.TabEventParams import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.TabEventRequest +import software.aws.toolkits.jetbrains.services.amazonq.lsp.util.LspEditorUtil +import software.aws.toolkits.jetbrains.services.amazonq.lsp.util.LspEditorUtil.toUriString import software.aws.toolkits.jetbrains.services.amazonq.util.command import software.aws.toolkits.jetbrains.services.amazonq.util.tabType import software.aws.toolkits.jetbrains.services.amazonq.webview.theme.AmazonQTheme @@ -238,19 +239,9 @@ class BrowserConnector( requestFromUi.params.prompt.escapedPrompt, node.command ) - val textDocumentIdentifier = getTextDocumentIdentifier(project) - val cursorState = CursorState( - Range( - Position( - 0, - 0 - ), - Position( - 1, - 1 - ) - ) - ) + val editor = FileEditorManager.getInstance(project).selectedTextEditor + val textDocumentIdentifier = editor?.let { TextDocumentIdentifier(toUriString(it.virtualFile)) } + val cursorState = editor?.let { LspEditorUtil.getCursorState(it) } val chatParams = ChatParams( requestFromUi.params.tabId, From ea3fb40f126ac5e1ca3108b783bf03a97412b55c Mon Sep 17 00:00:00 2001 From: Richard Li <742829+rli@users.noreply.github.com> Date: Thu, 8 May 2025 18:13:56 -0700 Subject: [PATCH 12/12] Update BrowserConnector.kt --- .../jetbrains/services/amazonq/webview/BrowserConnector.kt | 1 - 1 file changed, 1 deletion(-) 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 cd2926e5067..8db334596df 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 @@ -71,7 +71,6 @@ import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.CopyC import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.CopyCodeToClipboardParams import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.CreatePromptNotification import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.CreatePromptParams -import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.CursorState import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.EncryptedChatParams import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.EncryptedQuickActionChatParams import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.FeedbackNotification