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

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

@@ -1,16 +1,19 @@
 // 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
 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 +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

@@ -32,6 +32,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.ButtonClickNotification
 import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.ButtonClickParams
@@ -131,6 +132,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
                         }
@@ -256,7 +258,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
@@ -272,11 +276,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
@@ -325,6 +327,7 @@
             CHAT_READY -> {
                 handleChatNotification<ChatReadyNotification, Unit>(node) { server, _ ->
                     uiReady.complete(true)
+                    chatCommunicationManager.setUiReady()
                     RunOnceUtil.runOnceForApp("AmazonQ-UI-Ready") {
                         MeetQSettings.getInstance().reinvent2024OnboardingCount += 1
                     }
@@ -350,7 +353,7 @@
             }
             CHAT_OPEN_TAB -> {
                 val response = serializer.deserializeChatMessages<OpenTabResponse>(node)
-                ChatCommunicationManager.completeTabOpen(
+                chatCommunicationManager.completeTabOpen(
                     response.requestId,
                     response.params.result.tabId
                 )
@@ -414,7 +417,7 @@
 
             GET_SERIALIZED_CHAT_REQUEST_METHOD -> {
                 val response = serializer.deserializeChatMessages<GetSerializedChatResponse>(node)
-                ChatCommunicationManager.completeSerializedChatResponse(
+                chatCommunicationManager.completeSerializedChatResponse(
                     response.requestId,
                     response.params.result.content
                 )
@@ -424,13 +427,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/AmazonQLanguageClient.kt

@@ -6,6 +6,8 @@ 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_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
@@ -43,4 +45,7 @@ interface AmazonQLanguageClient : LanguageClient {
 
     @JsonNotification(OPEN_FILE_DIFF)
     fun openFileDiff(params: OpenFileDiffParams): CompletableFuture<Unit>
+
+    @JsonNotification(CHAT_SEND_CONTEXT_COMMANDS)
+    fun sendContextCommands(params: LSPAny): CompletableFuture<Unit>
 }

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,14 +23,18 @@
 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
 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
@@ -131,20 +134,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
@@ -183,21 +186,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
@@ -252,19 +254,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)
     }
@@ -308,6 +308,17 @@
             ApplicationManager.getApplication()::invokeLater
         )
 
+    override fun sendContextCommands(params: LSPAny): CompletableFuture<Unit> {
+        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)
+    }
+
     companion object {
         private val LOG = getLogger<AmazonQLanguageClientImpl>()
     }

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
@@ -282,23 +279,12 @@ private class AmazonQServerInstance(private val project: Project, private val cs
 
         launcher = LSPLauncher.Builder<AmazonQLanguageServer>()
             .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)
                 }
             }

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