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

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

@@ -140,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
                         }
@@ -324,6 +325,7 @@
             CHAT_READY -> {
                 handleChatNotification<ChatReadyNotification, Unit>(node) { server, _ ->
                     uiReady.complete(true)
+                    chatCommunicationManager.setUiReady()
                     RunOnceUtil.runOnceForApp("AmazonQ-UI-Ready") {
                         MeetQSettings.getInstance().reinvent2024OnboardingCount += 1
                     }
@@ -349,7 +351,7 @@
             }
             CHAT_OPEN_TAB -> {
                 val response = serializer.deserializeChatMessages<OpenTabResponse>(node)
-                ChatCommunicationManager.completeTabOpen(
+                chatCommunicationManager.completeTabOpen(
                     response.requestId,
                     response.params.result.tabId
                 )
@@ -420,7 +422,7 @@
 
             GET_SERIALIZED_CHAT_REQUEST_METHOD -> {
                 val response = serializer.deserializeChatMessages<GetSerializedChatResponse>(node)
-                ChatCommunicationManager.completeSerializedChatResponse(
+                chatCommunicationManager.completeSerializedChatResponse(
                     response.requestId,
                     response.params.result.content
                 )

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

@@ -137,9 +137,10 @@
     override fun openTab(params: OpenTabParams): CompletableFuture<OpenTabResult> {
         val requestId = UUID.randomUUID().toString()
         val result = CompletableFuture<OpenTabResult>()
-        ChatCommunicationManager.pendingTabRequests[requestId] = result
+        val chatManager = ChatCommunicationManager.getInstance(project)
+        chatManager.addTabOpenRequest(requestId, result)
 
-        AsyncChatUiListener.notifyPartialMessageUpdate(
+        chatManager.notifyUi(
             FlareUiMessage(
                 command = CHAT_OPEN_TAB,
                 params = params,
@@ -149,7 +150,7 @@
 
         result.orTimeout(30000, TimeUnit.MILLISECONDS)
             .whenComplete { _, error ->
-                ChatCommunicationManager.pendingTabRequests.remove(requestId)
+                chatManager.removeTabOpenRequest(requestId)
             }
 
         return result
@@ -188,10 +189,10 @@
     override fun getSerializedChat(params: GetSerializedChatParams): CompletableFuture<GetSerializedChatResult> {
         val requestId = UUID.randomUUID().toString()
         val result = CompletableFuture<GetSerializedChatResult>()
+        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,
@@ -201,7 +202,7 @@
 
         result.orTimeout(30000, TimeUnit.MILLISECONDS)
             .whenComplete { _, error ->
-                ChatCommunicationManager.pendingSerializedChatRequests.remove(requestId)
+                chatManager.removeSerializedChatRequest(requestId)
             }
 
         return result
@@ -340,13 +341,13 @@
         )
 
     override fun sendContextCommands(params: LSPAny): CompletableFuture<Unit> {
-        AsyncChatUiListener.notifyPartialMessageUpdate(
+        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/ChatCommunicationManager.kt

@@ -7,6 +7,9 @@
 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.CoroutineScope
+import kotlinx.coroutines.launch
 import org.eclipse.lsp4j.ProgressParams
 import software.aws.toolkits.core.utils.getLogger
 import software.aws.toolkits.core.utils.warn
@@ -29,12 +32,23 @@
 import java.util.concurrent.ConcurrentHashMap
 
 @Service(Service.Level.PROJECT)
-class ChatCommunicationManager {
+class ChatCommunicationManager(private val cs: CoroutineScope) {
+    val uiReady = CompletableDeferred<Boolean>()
     private val chatPartialResultMap = ConcurrentHashMap<String, String>()
-    private fun getPartialChatMessage(partialResultToken: String): String? =
-        chatPartialResultMap.getOrDefault(partialResultToken, null)
-
     private val inflightRequestByTabId = ConcurrentHashMap<String, CompletableFuture<String>>()
+    private val pendingSerializedChatRequests = ConcurrentHashMap<String, CompletableFuture<GetSerializedChatResult>>()
+    private val pendingTabRequests = ConcurrentHashMap<String, CompletableFuture<OpenTabResult>>()
+
+    fun setUiReady() {
+        uiReady.complete(true)
+    }
+
+    fun notifyUi(uiMessage: FlareUiMessage) {
+        cs.launch {
+            uiReady.await()
+            AsyncChatUiListener.notifyPartialMessageUpdate(uiMessage)
+        }
+    }
 
     fun setInflightRequestForTab(tabId: String, result: CompletableFuture<String>) {
         inflightRequestByTabId[tabId] = result
@@ -53,9 +67,36 @@
         return partialResultToken
     }
 
+    private fun getPartialChatMessage(partialResultToken: String): String? =
+        chatPartialResultMap.getOrDefault(partialResultToken, null)
+
     fun removePartialChatMessage(partialResultToken: String) =
         chatPartialResultMap.remove(partialResultToken)
 
+    fun addSerializedChatRequest(requestId: String, result: CompletableFuture<GetSerializedChatResult>) {
+        pendingSerializedChatRequests[requestId] = result
+    }
+
+    fun completeSerializedChatResponse(requestId: String, content: String) {
+        pendingSerializedChatRequests.remove(requestId)?.complete(GetSerializedChatResult((content)))
+    }
+
+    fun removeSerializedChatRequest(requestId: String) {
+        pendingSerializedChatRequests.remove(requestId)
+    }
+
+    fun addTabOpenRequest(requestId: String, result: CompletableFuture<OpenTabResult>) {
+        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) {
         val token = ProgressNotificationUtils.getToken(params)
         val tabId = getPartialChatMessage(token)
@@ -134,11 +175,6 @@
 
         private val LOG = getLogger<ChatCommunicationManager>()
 
-        val pendingSerializedChatRequests = ConcurrentHashMap<String, CompletableFuture<GetSerializedChatResult>>()
-        fun completeSerializedChatResponse(requestId: String, content: String) {
-            pendingSerializedChatRequests.remove(requestId)?.complete(GetSerializedChatResult((content)))
-        }
-
         fun convertToJsonToSendToChat(command: String, tabId: String, params: String, isPartialResult: Boolean): String =
             """
                 {
@@ -148,11 +184,5 @@
                 "isPartialResult": $isPartialResult
                 }
             """.trimIndent()
-
-        val pendingTabRequests = ConcurrentHashMap<String, CompletableFuture<OpenTabResult>>()
-
-        fun completeTabOpen(requestId: String, tabId: String) {
-            pendingTabRequests.remove(requestId)?.complete(OpenTabResult(tabId))
-        }
     }
 }