fix(amazonq): Enable stop button of chat (#5671)

* stop response cancel

* avoid get throw

* code

* endpoint

* detekt

* feedback

* update

* update

* msg

* fix ci

Author: leigaol

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

@@ -18,7 +18,6 @@ import kotlinx.coroutines.flow.distinctUntilChanged
 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
@@ -32,6 +31,7 @@ import software.aws.toolkits.jetbrains.services.amazonq.lsp.AmazonQLspService
 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
@@ -61,6 +61,7 @@ import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.ChatN
 import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.ChatParams
 import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.ChatPrompt
 import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.ChatReadyNotification
+import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.ChatUiMessageParams
 import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.ConversationClickRequest
 import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.CopyCodeToClipboardNotification
 import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.CopyCodeToClipboardParams
@@ -90,9 +91,11 @@ import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.Promp
 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
@@ -102,6 +105,7 @@ import software.aws.toolkits.jetbrains.services.amazonq.util.tabType
 import software.aws.toolkits.jetbrains.services.amazonq.webview.theme.AmazonQTheme
 import software.aws.toolkits.jetbrains.services.amazonq.webview.theme.ThemeBrowserAdapter
 import software.aws.toolkits.jetbrains.settings.MeetQSettings
+import software.aws.toolkits.resources.AwsCoreBundle
 import software.aws.toolkits.telemetry.MetricResult
 import software.aws.toolkits.telemetry.Telemetry
 import java.util.concurrent.CompletableFuture
@@ -245,6 +249,10 @@ class BrowserConnector(
                     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 @@ class BrowserConnector(
                     }
                 } ?: (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 @@ class BrowserConnector(
             CHAT_TAB_REMOVE -> {
                 handleChatNotification<TabEventRequest, TabEventParams>(node) { server, params ->
                     chatCommunicationManager.removePartialChatMessage(params.tabId)
+                    cancelInflightRequests(params.tabId)
                     server.tabRemove(params)
                 }
             }
@@ -414,6 +427,30 @@ class BrowserConnector(
                     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 = Gson().toJson(
+                    // https://github.com/aws/language-servers/blob/1c0d88806087125b6fc561f610cc15e98127c6bf/server/aws-lsp-codewhisperer/src/language-server/agenticChat/agenticChatController.ts#L403
+                    ChatUiMessageParams(
+                        title = AwsCoreBundle.message("amazonqChat.stopChatResponse"),
+                        body = ""
+                    )
+                )
+
+                val uiMessage = convertToJsonToSendToChat(
+                    command = SEND_CHAT_COMMAND_PROMPT,
+                    tabId = stopResponseRequest.params.tabId,
+                    params = paramsJson.toString(),
+                    isPartialResult = false
+                )
+                browser.postChat(uiMessage)
+            }
         }
     }
 
@@ -433,6 +470,14 @@ class BrowserConnector(
                 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/ChatUiMessageParams.kt

@@ -0,0 +1,14 @@
+// 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
+import java.util.UUID
+
+data class ChatUiMessageParams(
+    val title: String,
+    val additionalMessages: List<String> = emptyList(),
+    val messageId: String = UUID.randomUUID().toString(),
+    val buttons: List<String> = emptyList(),
+    val codeReference: List<String> = emptyList(),
+    val body: String = "",
+)

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,
+)

plugins/core/resources/resources/software/aws/toolkits/resources/MessagesBundle.properties

@@ -38,6 +38,7 @@ action.aws.toolkit.toolwindow.newConnection.text=Add Another Connection...
 action.dynamic.open.text=Open Resource...
 action.q.openchat.text=Open Chat Panel
 amazonqChat.project_context.index_in_progress=By the way, I'm still indexing this project for full context from your workspace. I may have a better response in a few minutes when it's complete if you'd like to try again then.
+amazonqChat.stopChatResponse=You stopped your current work, please provide additional examples or ask another question.
 amazonqDoc.answer.codeResult=You can accept the changes to your files, or describe any additional changes you'd like me to make.
 amazonqDoc.answer.readmeCreated=I've created a README for your code.
 amazonqDoc.answer.readmeUpdated=I've updated your README.