feedback

Author: leigaol

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

@@ -9,7 +9,6 @@ import com.intellij.ide.BrowserUtil
 import com.intellij.ide.util.RunOnceUtil
 import com.intellij.openapi.project.Project
 import com.intellij.ui.jcef.JBCefJSQuery.Response
-import fleet.multiplatform.shims.ConcurrentHashMap
 import kotlinx.coroutines.CompletableDeferred
 import kotlinx.coroutines.channels.awaitClose
 import kotlinx.coroutines.coroutineScope
@@ -119,7 +118,6 @@ class BrowserConnector(
 ) {
     var uiReady = CompletableDeferred<Boolean>()
     private val chatCommunicationManager = ChatCommunicationManager.getInstance(project)
-    private val inflightRequestByTabId = ConcurrentHashMap<String, CompletableFuture<String>>() // tab id -> cancellation token
 
     suspend fun connect(
         browser: Browser,
@@ -254,7 +252,7 @@ class BrowserConnector(
 
                 // We assume there is only one outgoing request per tab because the input is
                 // blocked when there is an outgoing request
-                inflightRequestByTabId[tabId] = result
+                chatCommunicationManager.setInflightRequestForTab(tabId, result)
                 showResult(result, partialResultToken, tabId, encryptionManager, browser)
             }
             CHAT_QUICK_ACTION -> {
@@ -274,7 +272,7 @@ class BrowserConnector(
 
                 // We assume there is only one outgoing request per tab because the input is
                 // blocked when there is an outgoing request
-                inflightRequestByTabId[tabId] = result
+                chatCommunicationManager.setInflightRequestForTab(tabId, result)
 
                 showResult(result, partialResultToken, tabId, encryptionManager, browser)
             }
@@ -431,7 +429,7 @@ class BrowserConnector(
             }
             STOP_CHAT_RESPONSE -> {
                 val stopResponseRequest = serializer.deserializeChatMessages<StopResponseMessage>(node)
-                if (inflightRequestByTabId[stopResponseRequest.params.tabId] == null) {
+                if (!chatCommunicationManager.hasInflightRequest(stopResponseRequest.params.tabId)) {
                     return
                 }
                 cancelInflightRequests(stopResponseRequest.params.tabId)
@@ -473,14 +471,14 @@ class BrowserConnector(
                 isPartialResult = false
             )
             browser.postChat(messageToChat)
-            inflightRequestByTabId.remove(tabId)
+            chatCommunicationManager.removeInflightRequestForTab(tabId)
         }
     }
 
     private fun cancelInflightRequests(tabId: String) {
-        inflightRequestByTabId[tabId]?.let { request ->
+        chatCommunicationManager.getInflightRequestForTab(tabId)?.let { request ->
             request.cancel(true)
-            inflightRequestByTabId.remove(tabId)
+            chatCommunicationManager.removeInflightRequestForTab(tabId)
         }
     }
 

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

@@ -7,6 +7,7 @@ import com.google.gson.Gson
 import com.intellij.openapi.components.Service
 import com.intellij.openapi.components.service
 import com.intellij.openapi.project.Project
+import fleet.multiplatform.shims.ConcurrentHashMap
 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
@@ -23,6 +24,23 @@ class ChatCommunicationManager {
     private fun getPartialChatMessage(partialResultToken: String): String? =
         chatPartialResultMap.getOrDefault(partialResultToken, null)
 
+    private val inflightRequestByTabId = ConcurrentHashMap<String, CompletableFuture<String>>()
+
+    fun setInflightRequestForTab(tabId: String, result: CompletableFuture<String>) {
+        inflightRequestByTabId[tabId] = result
+    }
+    fun removeInflightRequestForTab(tabId: String) {
+        inflightRequestByTabId.remove(tabId)
+    }
+
+    fun getInflightRequestForTab(tabId: String): CompletableFuture<String>? {
+        return inflightRequestByTabId[tabId]
+    }
+
+    fun hasInflightRequest(tabId: String): Boolean {
+        return inflightRequestByTabId.containsKey(tabId)
+    }
+
     fun addPartialChatMessage(tabId: String): String {
         val partialResultToken: String = UUID.randomUUID().toString()
         chatPartialResultMap[partialResultToken] = tabId