fix: Reduce perceived latency of fsWrite. Show fsWrite errors in the UX

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

@@ -3,7 +3,7 @@
  * Will be deleted or merged.
  */
 
 import * as path from 'path'
 import {
     ChatTriggerType,
     CodeWhispererStreamingServiceException,
@@ -91,7 +91,7 @@
     ChatResultWithMetadata as AgenticChatResultWithMetadata,
 } from './agenticChatEventParser'
 import { ChatSessionService } from '../chat/chatSessionService'
-import { AgenticChatResultStream, ResultStreamWriter } from './agenticChatResultStream'
+import { AgenticChatResultStream, progressPrefix, ResultStreamWriter } from './agenticChatResultStream'
 import { executeToolMessage, toolErrorMessage, toolResultMessage } from './textFormatting'
 import {
     AdditionalContentEntryAddition,
@@ -834,18 +834,15 @@
                 if (!availableToolNames.includes(toolUse.name)) {
                     throw new Error(`Tool ${toolUse.name} is not available in the current mode`)
                 }
+
+                // remove progress UI
+                await chatResultStream.removeResultBlockAndUpdateUI(progressPrefix + toolUse.toolUseId)
+
                 // fsRead and listDirectory write to an existing card and could show nothing in the current position
                 if (!['fsWrite', 'fsRead', 'listDirectory'].includes(toolUse.name)) {
                     await this.#showUndoAllIfRequired(chatResultStream, session)
                 }
-                const { explanation } = toolUse.input as unknown as ExplanatoryParams
-                if (explanation) {
-                    await chatResultStream.writeResultBlock({
-                        type: 'directive',
-                        messageId: toolUse.toolUseId + '_explanation',
-                        body: explanation,
-                    })
-                }
+
                 switch (toolUse.name) {
                     case 'fsRead':
                     case 'listDirectory':
@@ -1056,6 +1053,36 @@
                         continue
                     }
                 }
+                // display fs write failure status in the UX of that file card
+                if (toolUse.name === 'fsWrite' && toolUse.toolUseId) {
+                    const existingCard = chatResultStream.getMessageBlockId(toolUse.toolUseId)
+                    const fsParam = toolUse.input as unknown as FsWriteParams
+                    const fileName = path.basename(fsParam.path)
+                    const errorResult = {
+                        type: 'tool',
+                        messageId: toolUse.toolUseId,
+                        header: {
+                            fileList: {
+                                filePaths: [fileName],
+                                details: {
+                                    [fileName]: {
+                                        description: fsParam.path,
+                                    },
+                                },
+                            },
+                            status: {
+                                status: 'error',
+                                icon: 'error',
+                                text: 'Error',
+                            },
+                        },
+                    } as ChatResult
+                    if (existingCard) {
+                        await chatResultStream.overwriteResultBlock(errorResult, existingCard)
+                    } else {
+                        await chatResultStream.writeResultBlock(errorResult)
+                    }
+                }
                 const errMsg = err instanceof Error ? err.message : 'unknown error'
                 this.#log(`Error running tool ${toolUse.name}:`, errMsg)
                 results.push({
@@ -2203,6 +2230,64 @@
         }
     }
 
+    async #showToolUseIntermediateResult(
+        data: AgenticChatResultWithMetadata,
+        chatResultStream: AgenticChatResultStream,
+        streamWriter: ResultStreamWriter
+    ) {
+        // extract the key value from incomplete JSON response stream
+        function extractKey(incompleteJson: string, key: string): string | undefined {
+            const pattern = new RegExp(`"${key}":\\s*"([^"]*)"`, 'g')
+            const match = pattern.exec(incompleteJson)
+            return match?.[1]
+        }
+        const toolUses = Object.values(data.toolUses)
+        for (const toolUse of toolUses) {
+            if (toolUse.name === 'fsWrite' && typeof toolUse.input === 'string') {
+                const filepath = extractKey(toolUse.input, 'path')
+                const msgId = progressPrefix + toolUse.toolUseId
+                // render fs write UI as soon as fs write starts
+                if (filepath && !chatResultStream.hasMessage(msgId)) {
+                    const fileName = path.basename(filepath)
+                    await streamWriter.close()
+                    await chatResultStream.writeResultBlock({
+                        type: 'tool',
+                        messageId: msgId,
+                        header: {
+                            fileList: {
+                                filePaths: [fileName],
+                                details: {
+                                    [fileName]: {
+                                        description: filepath,
+                                    },
+                                },
+                            },
+                            status: {
+                                status: 'info',
+                                icon: 'progress',
+                                text: '',
+                            },
+                        },
+                    })
+                }
+            }
+
+            // render the tool use explanatory as soon as this is received for all tool uses.
+            if (typeof toolUse.input === 'string') {
+                const explanation = extractKey(toolUse.input, 'explanation')
+                const messageId = progressPrefix + toolUse.toolUseId + '_explanation'
+                if (explanation && !chatResultStream.hasMessage(messageId)) {
+                    await streamWriter.close()
+                    await chatResultStream.writeResultBlock({
+                        type: 'directive',
+                        messageId: messageId,
+                        body: explanation,
+                    })
+                }
+            }
+        }
+    }
+
     async #processGenerateAssistantResponseResponse(
         response: GenerateAssistantResponseCommandOutput,
         metric: Metric<AddMessageEvent>,
@@ -2251,6 +2336,7 @@
                     toolUseStartTimes,
                     toolUseLoadingTimeouts
                 )
+                await this.#showToolUseIntermediateResult(result.data, chatResultStream, streamWriter)
             }
         }
         await streamWriter.close()

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

@@ -1,5 +1,5 @@
 import { ChatResult, FileDetails, ChatMessage } from '@aws/language-server-runtimes/protocol'
 import { randomUUID } from 'crypto'
 
 export interface ResultStreamWriter {
     /**
@@ -44,6 +44,8 @@
 
 type OperationType = 'read' | 'write' | 'listDir'
 
+export const progressPrefix = 'progress_'
+
 interface FileOperation {
     type: OperationType
     filePaths: FileDetailsWithPath[]
@@ -201,6 +203,29 @@
         this.#state.chatResultBlocks = this.#state.chatResultBlocks.filter(block => block.messageId !== messageId)
     }
 
+    async removeResultBlockAndUpdateUI(messageId: string) {
+        if (this.hasMessage(messageId)) {
+            const blockId = this.getMessageBlockId(messageId)
+            if (blockId !== undefined) {
+                await this.overwriteResultBlock({ body: '', messageId: messageId }, blockId)
+            }
+            await this.removeResultBlock(messageId)
+        }
+    }
+
+    hasMessage(messageId: string): boolean {
+        return this.#state.chatResultBlocks.some(block => block.messageId === messageId)
+    }
+
+    getMessageBlockId(messageId: string): number | undefined {
+        for (const [i, block] of this.#state.chatResultBlocks.entries()) {
+            if (block.messageId === messageId) {
+                return i
+            }
+        }
+        return undefined
+    }
+
     getResultStreamWriter(): ResultStreamWriter {
         // Note: if write calls are not awaited, stream can be out-of-order.
         if (this.#state.isLocked) {