fix: Render timeout error, JSON parse error, cancellation to the in progress fs.write UI (#1382)

* fix: Render timeout error, JSON parse error, cancellation to the in progress fs.write UI

Author: leigaol

server/aws-lsp-codewhisperer/src/language-server/agenticChat/agenticChatController.ts

@@ -410,10 +410,11 @@ export class AgenticChatController implements ChatHandlers {
                 // depends on it
                 session.conversationId = uuid()
             }
-
+            const chatResultStream = this.#getChatResultStream(params.partialResultToken)
             token.onCancellationRequested(async () => {
                 this.#log('cancellation requested')
                 await this.#showUndoAllIfRequired(chatResultStream, session)
+                await chatResultStream.updateOngoingProgressResult('Canceled')
                 await this.#getChatResultStream(params.partialResultToken).writeResultBlock({
                     type: 'directive',
                     messageId: 'stopped' + uuid(),
@@ -431,8 +432,6 @@ export class AgenticChatController implements ChatHandlers {
             })
             session.setConversationType('AgenticChat')
 
-            const chatResultStream = this.#getChatResultStream(params.partialResultToken)
-
             const additionalContext = await this.#additionalContextProvider.getAdditionalContext(
                 triggerContext,
                 params.context
@@ -655,6 +654,8 @@ export class AgenticChatController implements ChatHandlers {
                 }
                 currentRequestInput = this.#updateRequestInputWithToolResults(currentRequestInput, [], content)
                 shouldDisplayMessage = false
+                // set the in progress tool use UI status to Error
+                await chatResultStream.updateOngoingProgressResult('Error')
                 continue
             }
 
@@ -756,6 +757,8 @@ export class AgenticChatController implements ChatHandlers {
                         'Your toolUse input is incomplete, try again. If the error happens consistently, break this task down into multiple tool uses with smaller input. Do not apologize.'
                     shouldDisplayMessage = false
                 }
+                // set the in progress tool use UI status to Error
+                await chatResultStream.updateOngoingProgressResult('Error')
             }
             if (result.success && this.#toolUseLatencies.length > 0) {
                 // Clear latencies for the next LLM call

server/aws-lsp-codewhisperer/src/language-server/agenticChat/agenticChatResultStream.ts

@@ -203,6 +203,10 @@ export class AgenticChatResultStream {
         this.#state.chatResultBlocks = this.#state.chatResultBlocks.filter(block => block.messageId !== messageId)
     }
 
+    /**
+     * Removes a specific messageId and renders the result on UI
+     * @param messageId
+     */
     async removeResultBlockAndUpdateUI(messageId: string) {
         if (this.hasMessage(messageId)) {
             const blockId = this.getMessageBlockId(messageId)
@@ -213,6 +217,22 @@ export class AgenticChatResultStream {
         }
     }
 
+    async updateOngoingProgressResult(errorMessage: string) {
+        for (const block of this.#state.chatResultBlocks) {
+            if (block.messageId?.startsWith(progressPrefix) && block.header?.status?.icon === 'progress') {
+                await this.removeResultBlockAndUpdateUI(block.messageId)
+                block.header.status = {
+                    status: 'error',
+                    icon: 'error',
+                    text: errorMessage,
+                }
+                block.messageId = block.messageId.substring(progressPrefix.length)
+                await this.writeResultBlock(block)
+                break
+            }
+        }
+    }
+
     hasMessage(messageId: string): boolean {
         return this.#state.chatResultBlocks.some(block => block.messageId === messageId)
     }