Reverting the PR 5765

Author: LokeshDogga13

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

@@ -42,7 +42,6 @@ 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
@@ -122,7 +121,6 @@ class BrowserConnector(
 ) {
     val uiReady = CompletableDeferred<Boolean>()
     private val chatCommunicationManager = ChatCommunicationManager.getInstance(project)
-    private val chatAsyncResultManager = ChatAsyncResultManager.getInstance(project)
 
     suspend fun connect(
         browser: Browser,
@@ -239,7 +237,6 @@ class BrowserConnector(
 
                 val tabId = requestFromUi.params.tabId
                 val partialResultToken = chatCommunicationManager.addPartialChatMessage(tabId)
-                chatCommunicationManager.registerPartialResultToken(partialResultToken)
 
                 var encryptionManager: JwtEncryptionManager? = null
                 val result = AmazonQLspService.executeAsyncIfRunning(project) { server ->
@@ -260,7 +257,6 @@ 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.executeAsyncIfRunning(project) { server ->
                     encryptionManager = this.encryptionManager
@@ -587,30 +583,13 @@ class BrowserConnector(
                 chatCommunicationManager.removeInflightRequestForTab(tabId)
             } 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 updateQuickActionsInBrowser(browser: Browser) {
         val isFeatureDevAvailable = isFeatureDevAvailable(project)
         val isCodeTransformAvailable = isCodeTransformAvailable(project)

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

@@ -43,8 +43,6 @@ class ChatCommunicationManager(private val project: Project, private val cs: Cor
     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>()
     private val openTabs = mutableSetOf<String>()
 
     fun setUiReady() {
@@ -118,20 +116,6 @@ class ChatCommunicationManager(private val project: Project, private val cs: Cor
     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)
@@ -147,49 +131,13 @@ class ChatCommunicationManager(private val project: Project, private val cs: Cor
         val encryptedPartialChatResult = getObject(params, String::class.java)
         if (encryptedPartialChatResult != null) {
             val partialChatResult = AmazonQLspService.getInstance(project).encryptionManager.decrypt(encryptedPartialChatResult)
-
-            // 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(project, 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(project, uiMessage)
-            }
+            val uiMessage = convertToJsonToSendToChat(
+                command = SEND_CHAT_COMMAND_PROMPT,
+                tabId = tabId,
+                params = partialChatResult,
+                isPartialResult = true
+            )
+            AsyncChatUiListener.notifyPartialMessageUpdate(project, uiMessage)
         }
     }
 
@@ -219,21 +167,6 @@ class ChatCommunicationManager(private val project: Project, private val cs: Cor
         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)