fix(amazonq): Stop button disappears after being clicked

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

@@ -41,6 +41,7 @@
 import software.aws.toolkits.jetbrains.services.amazonq.lsp.JsonRpcRequest
 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.ChatAsyncResultManager
 import software.aws.toolkits.jetbrains.services.amazonq.lsp.flareChat.ChatCommunicationManager
 import software.aws.toolkits.jetbrains.services.amazonq.lsp.flareChat.FlareUiMessage
 import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.AUTH_FOLLOW_UP_CLICKED
@@ -106,6 +107,7 @@
 ) {
     var uiReady = CompletableDeferred<Boolean>()
     private val chatCommunicationManager = ChatCommunicationManager.getInstance(project)
+    private val chatAsyncResultManager = ChatAsyncResultManager.getInstance(project)
 
     suspend fun connect(
         browser: Browser,
@@ -227,6 +229,7 @@
 
                 val tabId = requestFromUi.params.tabId
                 val partialResultToken = chatCommunicationManager.addPartialChatMessage(tabId)
+                chatCommunicationManager.registerPartialResultToken(partialResultToken)
 
                 var encryptionManager: JwtEncryptionManager? = null
                 val result = AmazonQLspService.executeIfRunning(project) { server ->
@@ -247,6 +250,7 @@
                 val tabId = requestFromUi.params.tabId
                 val quickActionParams = node.params ?: error("empty payload")
                 val partialResultToken = chatCommunicationManager.addPartialChatMessage(tabId)
+                chatCommunicationManager.registerPartialResultToken(partialResultToken)
                 var encryptionManager: JwtEncryptionManager? = null
                 val result = AmazonQLspService.executeIfRunning(project) { server ->
                     encryptionManager = this.encryptionManager
@@ -476,15 +480,32 @@
                 )
                 browser.postChat(messageToChat)
                 chatCommunicationManager.removeInflightRequestForTab(tabId)
-            } catch (_: CancellationException) {
+            } catch (e: CancellationException) {
                 LOG.warn { "Cancelled chat generation" }
+                try {
+                    chatAsyncResultManager.createRequestId(partialResultToken)
+                    chatAsyncResultManager.getResult(partialResultToken)
+                    handleCancellation(tabId, browser)
+                } catch (ex: Exception) {
+                    LOG.warn(ex) { "An error occurred while processing cancellation" }
+                } finally {
+                    chatAsyncResultManager.removeRequestId(partialResultToken)
+                    chatCommunicationManager.removePartialResultLock(partialResultToken)
+                    chatCommunicationManager.removeFinalResultProcessed(partialResultToken)
+                }
             } catch (e: Exception) {
                 LOG.warn(e) { "Failed to send chat message" }
                 browser.postChat(chatCommunicationManager.getErrorUiMessage(tabId, e, partialResultToken))
             }
         }
     }
 
+    private fun handleCancellation(tabId: String, browser: Browser) {
+        // Send a message to hide the stop button without showing an error
+        val cancelMessage = chatCommunicationManager.getCancellationUiMessage(tabId)
+        browser.postChat(cancelMessage)
+    }
+
     private fun cancelInflightRequests(tabId: String) {
         chatCommunicationManager.getInflightRequestForTab(tabId)?.let { request ->
             request.cancel(true)

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

@@ -0,0 +1,74 @@
+// 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.flareChat
+
+import com.intellij.openapi.components.Service
+import com.intellij.openapi.components.service
+import com.intellij.openapi.project.Project
+import java.util.concurrent.CompletableFuture
+import java.util.concurrent.ConcurrentHashMap
+import java.util.concurrent.TimeUnit
+import java.util.concurrent.TimeoutException
+
+/**
+ * Manages asynchronous results for chat operations, particularly handling the coordination
+ * between partial results and final results during cancellation.
+ */
+@Service(Service.Level.PROJECT)
+class ChatAsyncResultManager {
+    private val results = ConcurrentHashMap<String, CompletableFuture<Any>>()
+    private val completedResults = ConcurrentHashMap<String, Any>()
+    private val timeout = 30L
+    private val timeUnit = TimeUnit.SECONDS
+
+    fun createRequestId(requestId: String) {
+        if (!completedResults.containsKey(requestId)) {
+            results[requestId] = CompletableFuture()
+        }
+    }
+
+    fun removeRequestId(requestId: String) {
+        val future = results.remove(requestId)
+        if (future != null && !future.isDone) {
+            future.cancel(true)
+        }
+        completedResults.remove(requestId)
+    }
+
+    fun setResult(requestId: String, result: Any) {
+        val future = results[requestId]
+        if (future != null) {
+            future.complete(result)
+            results.remove(requestId)
+        }
+        completedResults[requestId] = result
+    }
+
+    fun getResult(requestId: String): Any? =
+        getResult(requestId, timeout, timeUnit)
+
+    private fun getResult(requestId: String, timeout: Long, unit: TimeUnit): Any? {
+        val completedResult = completedResults[requestId]
+        if (completedResult != null) {
+            return completedResult
+        }
+
+        val future = results[requestId] ?: throw IllegalArgumentException("Request ID not found: $requestId")
+
+        try {
+            val result = future.get(timeout, unit)
+            completedResults[requestId] = result
+            results.remove(requestId)
+            return result
+        } catch (e: TimeoutException) {
+            future.cancel(true)
+            results.remove(requestId)
+            throw TimeoutException("Operation timed out for requestId: $requestId")
+        }
+    }
+
+    companion object {
+        fun getInstance(project: Project) = project.service<ChatAsyncResultManager>()
+    }
+}

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

@@ -43,6 +43,8 @@
     private val inflightRequestByTabId = ConcurrentHashMap<String, CompletableFuture<String>>()
     private val pendingSerializedChatRequests = ConcurrentHashMap<String, CompletableFuture<GetSerializedChatResult>>()
     private val pendingTabRequests = ConcurrentHashMap<String, CompletableFuture<LSPAny>>()
+    private val partialResultLocks = ConcurrentHashMap<String, Any>()
+    private val finalResultProcessed = ConcurrentHashMap<String, Boolean>()
 
     fun setUiReady() {
         uiReady.complete(true)
@@ -97,6 +99,20 @@
     fun removeTabOpenRequest(requestId: String) =
         pendingTabRequests.remove(requestId)
 
+    fun removePartialResultLock(token: String) {
+        partialResultLocks.remove(token)
+    }
+
+    fun removeFinalResultProcessed(token: String) {
+        finalResultProcessed.remove(token)
+    }
+
+    fun registerPartialResultToken(partialResultToken: String) {
+        val lock = Any()
+        partialResultLocks[partialResultToken] = lock
+        finalResultProcessed[partialResultToken] = false
+    }
+
     fun handlePartialResultProgressNotification(project: Project, params: ProgressParams) {
         val token = ProgressNotificationUtils.getToken(params)
         val tabId = getPartialChatMessage(token)
@@ -112,13 +128,49 @@
         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,
-                params = partialChatResult,
-                isPartialResult = true
-            )
-            AsyncChatUiListener.notifyPartialMessageUpdate(uiMessage)
+
+            // Special case: check for stop message before proceeding
+            val partialResultMap = tryOrNull {
+                Gson().fromJson(partialChatResult, Map::class.java)
+            }
+
+            if (partialResultMap != null) {
+                @Suppress("UNCHECKED_CAST")
+                val additionalMessages = partialResultMap["additionalMessages"] as? List<Map<String, Any>>
+                if (additionalMessages != null) {
+                    for (message in additionalMessages) {
+                        val messageId = message["messageId"] as? String
+                        if (messageId != null && messageId.startsWith("stopped")) {
+                            // Process stop messages immediately
+                            val uiMessage = convertToJsonToSendToChat(
+                                command = SEND_CHAT_COMMAND_PROMPT,
+                                tabId = tabId,
+                                params = partialChatResult,
+                                isPartialResult = true
+                            )
+                            AsyncChatUiListener.notifyPartialMessageUpdate(uiMessage)
+                            finalResultProcessed[token] = true
+                            ChatAsyncResultManager.getInstance(project).setResult(token, partialResultMap)
+                            return
+                        }
+                    }
+                }
+            }
+
+            // Normal processing for non-stop messages
+            val lock = partialResultLocks[token] ?: return
+            synchronized(lock) {
+                if (finalResultProcessed[token] == true || partialResultLocks[token] == null) {
+                    return@synchronized
+                }
+                val uiMessage = convertToJsonToSendToChat(
+                    command = SEND_CHAT_COMMAND_PROMPT,
+                    tabId = tabId,
+                    params = partialChatResult,
+                    isPartialResult = true
+                )
+                AsyncChatUiListener.notifyPartialMessageUpdate(uiMessage)
+            }
         }
     }
 
@@ -148,6 +200,21 @@
         return uiMessage
     }
 
+    fun getCancellationUiMessage(tabId: String): String {
+        // Create a minimal error params with empty error message to hide the stop button
+        // without showing an actual error message to the user
+        val errorParams = Gson().toJson(ErrorParams(tabId, null, "", "")).toString()
+
+        return """
+            {
+            "command":"$CHAT_ERROR_PARAMS",
+            "tabId": "$tabId",
+            "params": $errorParams,
+            "isPartialResult": false
+            }
+        """.trimIndent()
+    }
+
     fun handleAuthFollowUpClicked(project: Project, params: AuthFollowUpClickedParams) {
         val incomingType = params.authFollowupType
         val connectionManager = ToolkitConnectionManager.getInstance(project)