fix(amazonq): Stop button disappears after being clicked (#5765)

Fixed a race condition between two streams receiving cancellation information. Previously, if the CancellationException from the sendChatPrompt stream was processed before the "You stopped your work" message from the partialResult stream, the UI would not properly display the stop message.

The fix prioritizes stop messages by checking for them before acquiring any locks, ensuring they're processed immediately when detected. This allows the cancellation to complete properly regardless of which stream receives information first.

Key changes:

* Added special case handling for stop messages before lock acquisition
* Immediately mark final result as processed when stop message is detected
* Set result in ChatAsyncResultManager to coordinate between streams

This ensures a consistent user experience when cancelling operations, with the stop message always being displayed properly.

Author: samgst-amazon

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.JsonRpcNotification
 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 @@ class BrowserConnector(
 ) {
     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 @@ class BrowserConnector(
 
                 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 @@ class BrowserConnector(
                 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 @@ class BrowserConnector(
                 )
                 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 @@ class ChatCommunicationManager(private val cs: CoroutineScope) {
     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 @@ class ChatCommunicationManager(private val cs: CoroutineScope) {
     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 @@ class ChatCommunicationManager(private val cs: CoroutineScope) {
         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 @@ class ChatCommunicationManager(private val cs: CoroutineScope) {
         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)