Handle error stages in Chat

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

@@ -178,19 +178,31 @@ class BrowserConnector(
                 var encryptionManager: JwtEncryptionManager? = null
                 val result = AmazonQLspService.executeIfRunning(project) { server ->
                     encryptionManager = this.encryptionManager
+
                     encryptionManager?.encrypt(chatParams)?.let { EncryptedChatParams(it, partialResultToken) }?.let { server.sendChatPrompt(it) }
+
                 } ?: (CompletableFuture.failedFuture(IllegalStateException("LSP Server not running")))
 
                 result.whenComplete {
                         value, error ->
-                    chatCommunicationManager.removePartialChatMessage(partialResultToken)
-                    val messageToChat = ChatCommunicationManager.convertToJsonToSendToChat(
-                        node.command,
-                        requestFromUi.params.tabId,
-                        encryptionManager?.decrypt(value).orEmpty(),
-                        isPartialResult = false
-                    )
-                    browser.postChat(messageToChat)
+                    if(error != null) {
+                        throw error
+                    }
+                    try {
+                        chatCommunicationManager.removePartialChatMessage(partialResultToken)
+                        val messageToChat = ChatCommunicationManager.convertToJsonToSendToChat(
+                            node.command,
+                            requestFromUi.params.tabId,
+                            encryptionManager?.decrypt(value) ?: "",
+                            isPartialResult = false
+                        )
+                        throw Exception("Trying an error message")
+                        browser.postChat(messageToChat)
+                    } catch (e: Exception) {
+                        chatCommunicationManager.sendErrorToUi(requestFromUi.params.tabId, e, partialResultToken)
+                    }
+
+
                 }
             }
         }

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

@@ -3,16 +3,24 @@
 
 package software.aws.toolkits.jetbrains.services.amazonq.lsp.flareChat
 
+import com.fasterxml.jackson.module.kotlin.jacksonObjectMapper
+import com.google.gson.Gson
+import com.google.gson.JsonElement
+import com.google.gson.JsonParser
+import com.intellij.docker.agent.util.toJson
 import com.intellij.openapi.components.Service
 import com.intellij.openapi.components.service
 import com.intellij.openapi.project.Project
 import org.eclipse.lsp4j.ProgressParams
 import software.aws.toolkits.jetbrains.services.amazonq.lsp.AmazonQLspService
 import software.aws.toolkits.jetbrains.services.amazonq.lsp.flareChat.ProgressNotificationUtils.getObject
+import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.CHAT_ERROR_PARAMS
+import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.ErrorParams
 import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.SEND_CHAT_COMMAND_PROMPT
 import java.util.UUID
 import java.util.concurrent.ConcurrentHashMap
 
+
 @Service(Service.Level.PROJECT)
 class ChatCommunicationManager {
     private val chatPartialResultMap = ConcurrentHashMap<String, String>()
@@ -34,36 +42,79 @@
         if (tabId == null || tabId.isEmpty()) {
             return
         }
-        if (params.value.isLeft || params.value.right == null) {
-            error(
-                "Error handling partial result notification: expected value of type Object"
-            )
+        try {
+            if (params.value.isLeft || params.value.right == null) {
+                error(
+                    "Error handling partial result notification: expected value of type Object"
+                )
+            }
+
+            val encryptedPartialChatResult = getObject(params, String::class.java)
+            if (encryptedPartialChatResult != null) {
+                val partialChatResult = AmazonQLspService.getInstance(project).encryptionManager.decrypt(encryptedPartialChatResult)
+                sendErrorToUi(tabId, IllegalStateException("Try out an error"), token)
+  //              sendMessageToChatUi(SEND_CHAT_COMMAND_PROMPT, tabId, partialChatResult, isPartialResult = true)
+            }
+        } catch (e: Exception) {
+            sendErrorToUi(tabId, e, token)
         }
 
-        val encryptedPartialChatResult = getObject(params, String::class.java)
-        if (encryptedPartialChatResult != null) {
-            val partialChatResult = AmazonQLspService.getInstance(project).encryptionManager.decrypt(encryptedPartialChatResult)
+    }
+
+    private fun sendMessageToChatUi(command: String, tabId: String, partialChatResult: String, isPartialResult: Boolean) {
+        val uiMessage = convertToJsonToSendToChat(
+            command = command,
+            tabId = tabId,
+            params = partialChatResult,
+            isPartialResult = isPartialResult
+        )
+        AsyncChatUiListener.notifyPartialMessageUpdate(uiMessage)
+    }
 
-            val uiMessage = convertToJsonToSendToChat(
-                command = SEND_CHAT_COMMAND_PROMPT,
-                tabId = tabId,
-                params = partialChatResult,
-                isPartialResult = true
-            )
-            AsyncChatUiListener.notifyPartialMessageUpdate(uiMessage)
-        }
+
+    fun sendErrorToUi(tabId: String, exception: Exception, token: String) {
+        removePartialChatMessage(token)
+        val errorTitle = "An error occurred while processing your request."
+        val errorMessage = "Details: ${exception.message}"
+        val errorParams = Gson().toJsonTree(ErrorParams(tabId, null, errorMessage, errorTitle))
+        sendErrorMessageToChatUi(CHAT_ERROR_PARAMS, tabId, errorParams, false)
+    }
+
+    private fun sendErrorMessageToChatUi(command: String, tabId: String, partialChatResult: JsonElement, isPartialResult: Boolean) {
+        val uiMessage =  """
+                {
+                "command":"$command",
+                "tabId": "$tabId",
+                "params": $partialChatResult,
+                "isPartialResult": $isPartialResult
+                }
+            """.trimIndent()
+        AsyncChatUiListener.notifyPartialMessageUpdate(uiMessage)
     }
     companion object {
         fun getInstance(project: Project) = project.service<ChatCommunicationManager>()
 
-        fun convertToJsonToSendToChat(command: String, tabId: String, params: String, isPartialResult: Boolean): String =
-            """
+        fun convertToJsonToSendToChat(command: String, tabId: String, params: String, isPartialResult: Boolean): String {
+            if(command == CHAT_ERROR_PARAMS) {
+                val param = JsonParser.parseString(params)
+                return """
+                {
+                "command":"$command",
+                "tabId": "$tabId",
+                "params": $param,
+                "isPartialResult": $isPartialResult
+                }
+            """.trimIndent()
+            }
+            return """
                 {
                 "command":"$command",
                 "tabId": "$tabId",
                 "params": $params,
                 "isPartialResult": $isPartialResult
                 }
             """.trimIndent()
+        }
+
     }
 }

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

@@ -0,0 +1,11 @@
+// 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 ErrorParams(
+    val tabID: String,
+    val triggerType: String?,
+    val message: String,
+    val title: String
+)

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

@@ -4,3 +4,4 @@
 package software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat
 
 const val SEND_CHAT_COMMAND_PROMPT = "aws/chat/sendChatPrompt"
+const val CHAT_ERROR_PARAMS = "errorMessage"