fix(amazonq): Enable stop button of chat

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

@@ -18,11 +18,11 @@
 import kotlinx.coroutines.flow.launchIn
 import kotlinx.coroutines.flow.merge
 import kotlinx.coroutines.flow.onEach
-import kotlinx.coroutines.future.await
 import kotlinx.coroutines.launch
 import org.cef.browser.CefBrowser
 import org.eclipse.lsp4j.Position
 import org.eclipse.lsp4j.Range
+import org.json.JSONObject
 import software.aws.toolkits.core.utils.getLogger
 import software.aws.toolkits.core.utils.warn
 import software.aws.toolkits.jetbrains.services.amazonq.apps.AppConnection
@@ -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.ChatCommunicationManager.Companion.convertToJsonToSendToChat
 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
@@ -90,9 +91,11 @@
 import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.PromptInputOptionChangeParams
 import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.QuickChatActionRequest
 import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.SEND_CHAT_COMMAND_PROMPT
+import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.STOP_CHAT_RESPONSE
 import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.SendChatPromptRequest
 import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.SourceLinkClickNotification
 import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.SourceLinkClickParams
+import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.StopResponseMessage
 import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.TabBarActionParams
 import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.TabBarActionRequest
 import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.TabEventParams
@@ -104,6 +107,7 @@
 import software.aws.toolkits.jetbrains.settings.MeetQSettings
 import software.aws.toolkits.telemetry.MetricResult
 import software.aws.toolkits.telemetry.Telemetry
+import java.util.UUID
 import java.util.concurrent.CompletableFuture
 import java.util.function.Function
 
@@ -245,6 +249,10 @@
                     encryptionManager = this.encryptionManager
                     encryptionManager?.encrypt(chatParams)?.let { EncryptedChatParams(it, partialResultToken) }?.let { server.sendChatPrompt(it) }
                 } ?: (CompletableFuture.failedFuture(IllegalStateException("LSP Server not running")))
+
+                // We assume there is only one outgoing request per tab because the input is
+                // blocked when there is an outgoing request
+                chatCommunicationManager.setInflightRequestForTab(tabId, result)
                 showResult(result, partialResultToken, tabId, encryptionManager, browser)
             }
             CHAT_QUICK_ACTION -> {
@@ -262,6 +270,10 @@
                     }
                 } ?: (CompletableFuture.failedFuture(IllegalStateException("LSP Server not running")))
 
+                // We assume there is only one outgoing request per tab because the input is
+                // blocked when there is an outgoing request
+                chatCommunicationManager.setInflightRequestForTab(tabId, result)
+
                 showResult(result, partialResultToken, tabId, encryptionManager, browser)
             }
             CHAT_LIST_CONVERSATIONS -> {
@@ -318,6 +330,7 @@
             CHAT_TAB_REMOVE -> {
                 handleChatNotification<TabEventRequest, TabEventParams>(node) { server, params ->
                     chatCommunicationManager.removePartialChatMessage(params.tabId)
+                    cancelInflightRequests(params.tabId)
                     server.tabRemove(params)
                 }
             }
@@ -414,6 +427,31 @@
                     server.createPrompt(params)
                 }
             }
+            STOP_CHAT_RESPONSE -> {
+                val stopResponseRequest = serializer.deserializeChatMessages<StopResponseMessage>(node)
+                if (!chatCommunicationManager.hasInflightRequest(stopResponseRequest.params.tabId)) {
+                    return
+                }
+                cancelInflightRequests(stopResponseRequest.params.tabId)
+                chatCommunicationManager.removePartialChatMessage(stopResponseRequest.params.tabId)
+
+                val paramsJson = JSONObject().apply {
+                    put("title", "You stopped your current work, please provide additional examples or ask another question.")
+                    put("additionalMessages", arrayOf<String>())
+                    put("messageId", UUID.randomUUID().toString())
+                    put("buttons", arrayOf<String>())
+                    put("codeReference", arrayOf<String>())
+                    put("body", "")
+                }
+
+                val uiMessage = convertToJsonToSendToChat(
+                    command = SEND_CHAT_COMMAND_PROMPT,
+                    tabId = stopResponseRequest.params.tabId,
+                    params = paramsJson.toString(),
+                    isPartialResult = false
+                )
+                browser.postChat(uiMessage)
+            }
         }
     }
 
@@ -426,13 +464,21 @@
     ) {
         result.whenComplete { value, error ->
             chatCommunicationManager.removePartialChatMessage(partialResultToken)
             val messageToChat = ChatCommunicationManager.convertToJsonToSendToChat(
                 SEND_CHAT_COMMAND_PROMPT,
                 tabId,
                 encryptionManager?.decrypt(value).orEmpty(),
                 isPartialResult = false
             )
             browser.postChat(messageToChat)
+            chatCommunicationManager.removeInflightRequestForTab(tabId)
+        }
+    }
+
+    private fun cancelInflightRequests(tabId: String) {
+        chatCommunicationManager.getInflightRequestForTab(tabId)?.let { request ->
+            request.cancel(true)
+            chatCommunicationManager.removeInflightRequestForTab(tabId)
         }
     }
 

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

@@ -20,8 +20,21 @@ import java.util.concurrent.ConcurrentHashMap
 @Service(Service.Level.PROJECT)
 class ChatCommunicationManager {
     private val chatPartialResultMap = ConcurrentHashMap<String, String>()
-    private fun getPartialChatMessage(partialResultToken: String): String =
-        chatPartialResultMap.getValue(partialResultToken)
+    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>? = inflightRequestByTabId[tabId]
+
+    fun hasInflightRequest(tabId: String): Boolean = inflightRequestByTabId.containsKey(tabId)
 
     fun addPartialChatMessage(tabId: String): String {
         val partialResultToken: String = UUID.randomUUID().toString()
@@ -35,7 +48,7 @@ class ChatCommunicationManager {
     fun handlePartialResultProgressNotification(project: Project, params: ProgressParams) {
         val token = ProgressNotificationUtils.getToken(params)
         val tabId = getPartialChatMessage(token)
-        if (tabId == null || tabId.isEmpty()) {
+        if (tabId.isNullOrEmpty()) {
             return
         }
         if (params.value.isLeft || params.value.right == null) {
@@ -47,7 +60,6 @@ class ChatCommunicationManager {
         val encryptedPartialChatResult = getObject(params, String::class.java)
         if (encryptedPartialChatResult != null) {
             val partialChatResult = AmazonQLspService.getInstance(project).encryptionManager.decrypt(encryptedPartialChatResult)
-
             val uiMessage = convertToJsonToSendToChat(
                 command = SEND_CHAT_COMMAND_PROMPT,
                 tabId = tabId,

plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/model/aws/chat/FlareChatCommands.kt

@@ -31,3 +31,4 @@ const val CHAT_SEND_UPDATE = "aws/chat/sendChatUpdate"
 const val CHAT_CREATE_PROMPT = "aws/chat/createPrompt"
 const val SHOW_SAVE_FILE_DIALOG_REQUEST_METHOD = "aws/showSaveFileDialog"
 const val GET_SERIALIZED_CHAT_REQUEST_METHOD = "aws/chat/getSerializedChat"
+const val STOP_CHAT_RESPONSE = "stopChatResponse"

plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/model/aws/chat/StopResponseParams.kt

@@ -0,0 +1,12 @@
+// Copyright 2025 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+package software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat
+
+data class StopResponseMessage(
+    val params: StopResponseParams,
+)
+
+data class StopResponseParams(
+    val tabId: String,
+)