move tab request map into ChatCommunicationManager

Author: samgst-amazon

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

@@ -26,7 +26,6 @@ import software.aws.toolkits.core.utils.getLogger
 import software.aws.toolkits.core.utils.warn
 import software.aws.toolkits.jetbrains.services.amazonq.apps.AppConnection
 import software.aws.toolkits.jetbrains.services.amazonq.commands.MessageSerializer
-import software.aws.toolkits.jetbrains.services.amazonq.lsp.AmazonQLanguageClientImpl
 import software.aws.toolkits.jetbrains.services.amazonq.lsp.AmazonQLanguageServer
 import software.aws.toolkits.jetbrains.services.amazonq.lsp.AmazonQLspService
 import software.aws.toolkits.jetbrains.services.amazonq.lsp.encryption.JwtEncryptionManager
@@ -288,7 +287,7 @@ class BrowserConnector(
             }
             CHAT_OPEN_TAB -> {
                 val response = serializer.deserializeChatMessages<OpenTabResponse>(node)
-                AmazonQLanguageClientImpl.completeTabOpen(
+                chatCommunicationManager.completeTabOpen(
                     response.requestId,
                     response.params.result.tabId
                 )

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

@@ -36,8 +36,6 @@ 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.ConcurrentHashMap
-import java.util.concurrent.TimeUnit
 
 /**
  * Concrete implementation of [AmazonQLanguageClient] to handle messages sent from server
@@ -116,9 +114,7 @@ class AmazonQLanguageClientImpl(private val project: Project) : AmazonQLanguageC
 
     override fun openTab(params: OpenTabParams): CompletableFuture<OpenTabResult> {
         val requestId = UUID.randomUUID().toString()
-        val result = CompletableFuture<OpenTabResult>()
-
-        pendingTabRequests[requestId] = result
+        val result = ChatCommunicationManager.getInstance(project).addPendingOpenTabRequest(requestId)
 
         val uiMessage = """
                 {
@@ -129,13 +125,6 @@ class AmazonQLanguageClientImpl(private val project: Project) : AmazonQLanguageC
         """.trimIndent()
         AsyncChatUiListener.notifyPartialMessageUpdate(uiMessage)
 
-        result.orTimeout(30000, TimeUnit.MILLISECONDS)
-            .whenComplete { _, error ->
-                if (error != null) {
-                    pendingTabRequests.remove(requestId)
-                }
-            }
-
         return result
     }
 
@@ -194,11 +183,5 @@ class AmazonQLanguageClientImpl(private val project: Project) : AmazonQLanguageC
 
     companion object {
         private val LOG = getLogger<AmazonQLanguageClientImpl>()
-
-        private val pendingTabRequests = ConcurrentHashMap<String, CompletableFuture<OpenTabResult>>()
-
-        fun completeTabOpen(requestId: String, tabId: String) {
-            pendingTabRequests.remove(requestId)?.complete(OpenTabResult(tabId))
-        }
     }
 }

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

@@ -9,13 +9,30 @@ import com.intellij.openapi.project.Project
 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
+import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.OpenTabResult
 import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.SEND_CHAT_COMMAND_PROMPT
 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)
 
@@ -53,6 +70,11 @@ class ChatCommunicationManager {
             AsyncChatUiListener.notifyPartialMessageUpdate(uiMessage)
         }
     }
+
+    fun completeTabOpen(requestId: String, tabId: String) {
+        pendingTabRequests.remove(requestId)?.complete(OpenTabResult(tabId))
+    }
+
     companion object {
         fun getInstance(project: Project) = project.service<ChatCommunicationManager>()