fix: ux polish for list directory tool messages

Author: laileni-aws

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

@@ -691,34 +691,33 @@ export class AgenticChatController implements ChatHandlers {
     }
 
     #processReadOrList(toolUse: ToolUse, chatResultStream: AgenticChatResultStream): ChatMessage | undefined {
-        if (toolUse.name !== 'fsRead') {
-            //TODO: Implement list directory UX in next PR.
-            return {}
-        }
-        let messageId = toolUse.toolUseId || ''
-        if (chatResultStream.getMessageIdToUpdate()) {
-            messageId = chatResultStream.getMessageIdToUpdate()!
-        } else if (messageId) {
-            chatResultStream.setMessageIdToUpdate(messageId)
+        let messageIdToUpdate = toolUse.toolUseId!
+        const currentId = chatResultStream.getMessageIdToUpdateForTool(toolUse.name!)
+
+        if (currentId) {
+            messageIdToUpdate = currentId
+        } else {
+            chatResultStream.setMessageIdToUpdateForTool(toolUse.name!, messageIdToUpdate)
         }
+
         const currentPath = (toolUse.input as unknown as FsReadParams | ListDirectoryParams)?.path
         if (!currentPath) return
-        const existingPaths = chatResultStream.getMessageOperation(messageId)?.filePaths || []
+        const existingPaths = chatResultStream.getMessageOperation(messageIdToUpdate)?.filePaths || []
         // Check if path already exists in the list
         const isPathAlreadyProcessed = existingPaths.some(path => path.relativeFilePath === currentPath)
         if (!isPathAlreadyProcessed) {
             const currentFileDetail = {
                 relativeFilePath: currentPath,
                 lineRanges: [{ first: -1, second: -1 }],
             }
-            const operationType = toolUse.name === 'fsRead' ? 'read' : 'listDir'
-            if (operationType === 'read') {
-                chatResultStream.addMessageOperation(messageId, operationType, [...existingPaths, currentFileDetail])
-            }
+            chatResultStream.addMessageOperation(messageIdToUpdate, toolUse.name!, [
+                ...existingPaths,
+                currentFileDetail,
+            ])
         }
         let title: string
-        const itemCount = chatResultStream.getMessageOperation(messageId)?.filePaths.length
-        const filePathsPushed = chatResultStream.getMessageOperation(messageId)?.filePaths ?? []
+        const itemCount = chatResultStream.getMessageOperation(messageIdToUpdate)?.filePaths.length
+        const filePathsPushed = chatResultStream.getMessageOperation(messageIdToUpdate)?.filePaths ?? []
         if (!itemCount) {
             title = 'Gathering context'
         } else {
@@ -742,7 +741,7 @@ export class AgenticChatController implements ChatHandlers {
         return {
             type: 'tool',
             contextList,
-            messageId,
+            messageId: messageIdToUpdate,
             body: '',
         }
     }

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

@@ -29,7 +29,7 @@ export class AgenticChatResultStream {
         isLocked: false,
         uuid: randomUUID(),
         messageId: undefined as string | undefined,
-        messageIdToUpdate: undefined as string | undefined,
+        messageIdToUpdateForTool: new Map<OperationType, string>(),
         messageOperations: new Map<string, FileOperation>(),
     }
     readonly #sendProgress: (newChatResult: ChatResult | string) => Promise<void>
@@ -41,22 +41,23 @@ export class AgenticChatResultStream {
     getResult(only?: string): ChatResult {
         return this.#joinResults(this.#state.chatResultBlocks, only)
     }
-    getMessageIdToUpdate(): string | undefined {
-        return this.#state.messageIdToUpdate
+
+    setMessageIdToUpdateForTool(toolName: string, messageId: string) {
+        this.#state.messageIdToUpdateForTool.set(toolName as OperationType, messageId)
     }
 
-    setMessageIdToUpdate(messageId: string) {
-        this.#state.messageIdToUpdate = messageId
+    getMessageIdToUpdateForTool(toolName: string): string | undefined {
+        return this.#state.messageIdToUpdateForTool.get(toolName as OperationType)
     }
 
     /**
      * Adds a file operation for a specific message
      * @param messageId The ID of the message
-     * @param type The type of operation ('read' or 'listDir' or 'write')
+     * @param type The type of operation ('fsRead' or 'listDirectory' or 'fsWrite')
      * @param filePaths Array of FileDetailsWithPath involved in the operation
      */
-    addMessageOperation(messageId: string, type: OperationType, filePaths: FileDetailsWithPath[]) {
-        this.#state.messageOperations.set(messageId, { type, filePaths })
+    addMessageOperation(messageId: string, type: string, filePaths: FileDetailsWithPath[]) {
+        this.#state.messageOperations.set(messageId, { type: type as OperationType, filePaths })
     }
 
     /**