fix(amazonq): ChatCommunicationManager blocks messages to UI until ready

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,10 +56,12 @@ class AmazonQPanel(val project: Project, private val scope: CoroutineScope) : Di
         project.messageBus.connect().subscribe(
             AsyncChatUiListener.TOPIC,
             object : AsyncChatUiListener {
-                override fun onChange(message: String) {
-                    runInEdt {
-                        browser.get()?.postChat(message)
-                    }
+                override fun onChange(command: String) {
+                    browser.get()?.postChat(command)
+                }
+
+                override fun onChange(command: FlareUiMessage) {
+                    browser.get()?.postChat(command)
                 }
             }
         )

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

@@ -1,9 +1,10 @@
 // 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
+import com.google.gson.Gson
 import com.intellij.openapi.Disposable
 import com.intellij.openapi.project.Project
 import com.intellij.openapi.util.Disposer
@@ -12,6 +13,7 @@ import org.cef.CefApp
 import software.aws.toolkits.jetbrains.services.amazonq.CodeWhispererFeatureConfigService
 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
@@ -63,6 +65,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

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

@@ -34,6 +34,7 @@
 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.AUTH_FOLLOW_UP_CLICKED
 import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.AuthFollowUpClickNotification
@@ -139,6 +140,7 @@
                     // 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
                         }
@@ -264,7 +266,9 @@
                 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
@@ -280,11 +284,9 @@
                 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
@@ -333,6 +335,7 @@
             CHAT_READY -> {
                 handleChatNotification<ChatReadyNotification, Unit>(node) { server, _ ->
                     uiReady.complete(true)
+                    chatCommunicationManager.setUiReady()
                     RunOnceUtil.runOnceForApp("AmazonQ-UI-Ready") {
                         MeetQSettings.getInstance().reinvent2024OnboardingCount += 1
                     }
@@ -358,7 +361,7 @@
             }
             CHAT_OPEN_TAB -> {
                 val response = serializer.deserializeChatMessages<OpenTabResponse>(node)
-                ChatCommunicationManager.completeTabOpen(
+                chatCommunicationManager.completeTabOpen(
                     response.requestId,
                     response.params.result.tabId
                 )
@@ -429,7 +432,7 @@
 
             GET_SERIALIZED_CHAT_REQUEST_METHOD -> {
                 val response = serializer.deserializeChatMessages<GetSerializedChatResponse>(node)
-                ChatCommunicationManager.completeSerializedChatResponse(
+                chatCommunicationManager.completeSerializedChatResponse(
                     response.requestId,
                     response.params.result.content
                 )
@@ -439,13 +442,18 @@
                 handleChatNotification<TabBarActionRequest, TabBarActionParams>(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 ?: emptyMap<Any, Any>()
+                                )
+                            )
                         } catch (e: Exception) {
                             LOG.error { "Failed to perform chat tab bar action $e" }
                             params.tabId?.let {

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.DiffManagerEx
@@ -26,13 +25,15 @@
 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
 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
@@ -136,20 +137,20 @@
     override fun openTab(params: OpenTabParams): CompletableFuture<OpenTabResult> {
         val requestId = UUID.randomUUID().toString()
         val result = CompletableFuture<OpenTabResult>()
-        ChatCommunicationManager.pendingTabRequests[requestId] = result
-
-        val uiMessage = """
-                {
-                "command": "$CHAT_OPEN_TAB",
-                "params": ${Gson().toJson(params)},
-                "requestId": "$requestId"
-                }
-        """.trimIndent()
-        AsyncChatUiListener.notifyPartialMessageUpdate(uiMessage)
+        val chatManager = ChatCommunicationManager.getInstance(project)
+        chatManager.addTabOpenRequest(requestId, result)
+
+        chatManager.notifyUi(
+            FlareUiMessage(
+                command = CHAT_OPEN_TAB,
+                params = params,
+                requestId = requestId,
+            )
+        )
 
         result.orTimeout(30000, TimeUnit.MILLISECONDS)
             .whenComplete { _, error ->
-                ChatCommunicationManager.pendingTabRequests.remove(requestId)
+                chatManager.removeTabOpenRequest(requestId)
             }
 
         return result
@@ -188,21 +189,20 @@
     override fun getSerializedChat(params: GetSerializedChatParams): CompletableFuture<GetSerializedChatResult> {
         val requestId = UUID.randomUUID().toString()
         val result = CompletableFuture<GetSerializedChatResult>()
-
-        ChatCommunicationManager.pendingSerializedChatRequests[requestId] = result
-
-        val uiMessage = """
-                {
-                "command": "$GET_SERIALIZED_CHAT_REQUEST_METHOD",
-                "params": ${Gson().toJson(params)},
-                "requestId": "$requestId"
-                }
-        """.trimIndent()
-        AsyncChatUiListener.notifyPartialMessageUpdate(uiMessage)
+        val chatManager = ChatCommunicationManager.getInstance(project)
+        chatManager.addSerializedChatRequest(requestId, result)
+
+        chatManager.notifyUi(
+            FlareUiMessage(
+                command = GET_SERIALIZED_CHAT_REQUEST_METHOD,
+                params = params,
+                requestId = requestId,
+            )
+        )
 
         result.orTimeout(30000, TimeUnit.MILLISECONDS)
             .whenComplete { _, error ->
-                ChatCommunicationManager.pendingSerializedChatRequests.remove(requestId)
+                chatManager.removeSerializedChatRequest(requestId)
             }
 
         return result
@@ -257,19 +257,17 @@
         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<Unit> {
-        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)
     }
@@ -343,15 +341,13 @@
         )
 
     override fun sendContextCommands(params: LSPAny): CompletableFuture<Unit> {
-        val showContextCommands = """
-            {
-            "command":"$CHAT_SEND_CONTEXT_COMMANDS",
-            "params": ${Gson().toJson(params)}
-            }
-        """.trimIndent()
-
-        AsyncChatUiListener.notifyPartialMessageUpdate(showContextCommands)
-
+        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)
     }
 

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)
         }
     }
 }