refactor for consistency

Author: samgst-amazon

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

@@ -286,7 +286,7 @@ class BrowserConnector(
             }
             CHAT_OPEN_TAB -> {
                 val response = serializer.deserializeChatMessages<OpenTabResponse>(node)
-                chatCommunicationManager.completeTabOpen(
+                ChatCommunicationManager.completeTabOpen(
                     response.requestId,
                     response.params.result.tabId
                 )

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

@@ -38,6 +38,7 @@ import software.aws.toolkits.jetbrains.services.codewhisperer.customization.Code
 import software.aws.toolkits.jetbrains.settings.CodeWhispererSettings
 import java.util.UUID
 import java.util.concurrent.CompletableFuture
+import java.util.concurrent.TimeUnit
 
 /**
  * Concrete implementation of [AmazonQLanguageClient] to handle messages sent from server
@@ -116,7 +117,8 @@ class AmazonQLanguageClientImpl(private val project: Project) : AmazonQLanguageC
 
     override fun openTab(params: OpenTabParams): CompletableFuture<OpenTabResult> {
         val requestId = UUID.randomUUID().toString()
-        val result = ChatCommunicationManager.getInstance(project).addPendingOpenTabRequest(requestId)
+        val result = CompletableFuture<OpenTabResult>()
+        ChatCommunicationManager.pendingTabRequests[requestId] = result
 
         val uiMessage = """
                 {
@@ -127,6 +129,11 @@ class AmazonQLanguageClientImpl(private val project: Project) : AmazonQLanguageC
         """.trimIndent()
         AsyncChatUiListener.notifyPartialMessageUpdate(uiMessage)
 
+        result.orTimeout(30000, TimeUnit.MILLISECONDS)
+            .whenComplete { _, error ->
+                ChatCommunicationManager.pendingTabRequests.remove(requestId)
+            }
+
         return result
     }
 

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

@@ -14,25 +14,10 @@ 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.TimeUnit
 
 @Service(Service.Level.PROJECT)
 class ChatCommunicationManager {
     private val chatPartialResultMap = ConcurrentHashMap<String, String>()
-    private val pendingTabRequests = ConcurrentHashMap<String, CompletableFuture<OpenTabResult>>()
-
-    fun addPendingOpenTabRequest(requestId: String): CompletableFuture<OpenTabResult> {
-        return CompletableFuture<OpenTabResult>().also { future ->
-            pendingTabRequests[requestId] = future
-            future.orTimeout(30000, TimeUnit.MILLISECONDS)
-                .whenComplete { _, error ->
-                    if (error != null) {
-                        pendingTabRequests.remove(requestId)
-                    }
-                }
-        }
-    }
-
     private fun getPartialChatMessage(partialResultToken: String): String =
         chatPartialResultMap.getValue(partialResultToken)
 
@@ -71,10 +56,6 @@ class ChatCommunicationManager {
         }
     }
 
-    fun completeTabOpen(requestId: String, tabId: String) {
-        pendingTabRequests.remove(requestId)?.complete(OpenTabResult(tabId))
-    }
-
     companion object {
         fun getInstance(project: Project) = project.service<ChatCommunicationManager>()
 
@@ -87,5 +68,17 @@ class ChatCommunicationManager {
                 "isPartialResult": $isPartialResult
                 }
             """.trimIndent()
+
+        val pendingTabRequests = ConcurrentHashMap<String, CompletableFuture<OpenTabResult>>()
+
+        fun addPendingOpenTabRequest(requestId: String): CompletableFuture<OpenTabResult> {
+            return CompletableFuture<OpenTabResult>().also { future ->
+                pendingTabRequests[requestId] = future
+            }
+        }
+
+        fun completeTabOpen(requestId: String, tabId: String) {
+            pendingTabRequests.remove(requestId)?.complete(OpenTabResult(tabId))
+        }
     }
 }