fix: ux polish for list directory tool messages

Author: laileni-aws

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

@@ -691,34 +691,37 @@ 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)
+        const operationType = toolUse.name === 'fsRead' ? 'read' : 'listDir'
+        let messageIdToUpdate = toolUse.toolUseId!
+
+        if (operationType === 'read' || operationType === 'listDir') {
+            const currentId = chatResultStream.getMessageIdToUpdateForTool(operationType)
+
+            if (currentId) {
+                messageIdToUpdate = currentId
+            } else {
+                chatResultStream.setMessageIdToUpdateForTool(operationType, 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, operationType, [
+                ...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 +745,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,12 +41,13 @@ export class AgenticChatResultStream {
     getResult(only?: string): ChatResult {
         return this.#joinResults(this.#state.chatResultBlocks, only)
     }
-    getMessageIdToUpdate(): string | undefined {
-        return this.#state.messageIdToUpdate
+
+    setMessageIdToUpdateForTool(toolName: OperationType, messageId: string) {
+        this.#state.messageIdToUpdateForTool.set(toolName, messageId)
     }
 
-    setMessageIdToUpdate(messageId: string) {
-        this.#state.messageIdToUpdate = messageId
+    getMessageIdToUpdateForTool(toolName: OperationType): string | undefined {
+        return this.#state.messageIdToUpdateForTool.get(toolName)
     }
 
     /**