fix: ux polishing for read tool messages.

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,
     GenerateAssistantResponseCommandInput,
@@ -703,22 +703,34 @@
     }
 
     #processReadOrList(toolUse: ToolUse, chatResultStream: AgenticChatResultStream): ChatMessage | undefined {
-        // return initial message about fsRead or listDir
-        const toolUseId = toolUse.toolUseId!
-        const currentPath = (toolUse.input as unknown as FsReadParams | ListDirectoryParams).path
+        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 currentPath = (toolUse.input as unknown as FsReadParams | ListDirectoryParams)?.path
         if (!currentPath) return
-        const currentFileList = chatResultStream.getContextFileList(toolUseId)
-        if (!currentFileList.some(path => path.relativeFilePath === currentPath)) {
+        const existingPaths = chatResultStream.getMessageOperation(messageId)?.filePaths || []
+        // Check if path already exists in the list
+        const isPathAlreadyProcessed = existingPaths.some(path => path.relativeFilePath === currentPath)
+        if (!isPathAlreadyProcessed) {
             const currentFileDetail = {
-                relativeFilePath: (toolUse.input as any)?.path,
+                relativeFilePath: currentPath,
                 lineRanges: [{ first: -1, second: -1 }],
             }
-            chatResultStream.addContextFileList(toolUseId, currentFileDetail)
-            currentFileList.push(currentFileDetail)
+            const operationType = toolUse.name === 'fsRead' ? 'read' : 'listDir'
+            if (operationType === 'read') {
+                chatResultStream.addMessageOperation(messageId, operationType, [...existingPaths, currentFileDetail])
+            }
         }
-
         let title: string
-        const itemCount = currentFileList.length
+        const itemCount = chatResultStream.getMessageOperation(messageId)?.filePaths.length
+        const filePathsPushed = chatResultStream.getMessageOperation(messageId)?.filePaths ?? []
         if (!itemCount) {
             title = 'Gathering context'
         } else {
@@ -728,22 +740,21 @@
                     : `${itemCount} ${itemCount === 1 ? 'directory' : 'directories'} listed`
         }
         const fileDetails: Record<string, FileDetails> = {}
-        for (const item of currentFileList) {
+        for (const item of filePathsPushed) {
             fileDetails[item.relativeFilePath] = {
                 lineRanges: item.lineRanges,
             }
         }
 
         const contextList: FileList = {
             rootFolderTitle: title,
-            filePaths: currentFileList.map(item => item.relativeFilePath),
+            filePaths: filePathsPushed.map(item => item.relativeFilePath),
             details: fileDetails,
         }
-
         return {
             type: 'tool',
             contextList,
-            messageId: toolUseId,
+            messageId,
             body: '',
         }
     }

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'
 
 interface ResultStreamWriter {
     write(chunk: ChatResult, final?: boolean): Promise<void>
@@ -15,14 +15,22 @@
 interface FileDetailsWithPath extends FileDetails {
     relativeFilePath: string
 }
+
+type OperationType = 'read' | 'write' | 'listDir'
+
+interface FileOperation {
+    type: OperationType
+    filePaths: FileDetailsWithPath[]
+}
 export class AgenticChatResultStream {
     static readonly resultDelimiter = '\n\n'
     #state = {
         chatResultBlocks: [] as ChatMessage[],
         isLocked: false,
-        contextFileList: {} as Record<string, FileDetailsWithPath[]>,
         uuid: randomUUID(),
         messageId: undefined as string | undefined,
+        messageIdToUpdate: undefined as string | undefined,
+        messageOperations: new Map<string, FileOperation>(),
     }
     readonly #sendProgress: (newChatResult: ChatResult | string) => Promise<void>
 
@@ -33,16 +41,31 @@
     getResult(only?: string): ChatResult {
         return this.#joinResults(this.#state.chatResultBlocks, only)
     }
+    getMessageIdToUpdate(): string | undefined {
+        return this.#state.messageIdToUpdate
+    }
 
-    getContextFileList(toolUseId: string): FileDetailsWithPath[] {
-        return this.#state.contextFileList[toolUseId] ?? []
+    setMessageIdToUpdate(messageId: string) {
+        this.#state.messageIdToUpdate = messageId
     }
 
-    addContextFileList(toolUseId: string, fileDetails: FileDetailsWithPath) {
-        if (!this.#state.contextFileList[toolUseId]) {
-            this.#state.contextFileList[toolUseId] = []
-        }
-        this.#state.contextFileList[toolUseId].push(fileDetails)
+    /**
+     * 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 filePaths Array of FileDetailsWithPath involved in the operation
+     */
+    addMessageOperation(messageId: string, type: OperationType, filePaths: FileDetailsWithPath[]) {
+        this.#state.messageOperations.set(messageId, { type, filePaths })
+    }
+
+    /**
+     * Gets the file operation details for a specific message
+     * @param messageId The ID of the message
+     * @returns The file operation details or undefined if not found
+     */
+    getMessageOperation(messageId: string): FileOperation | undefined {
+        return this.#state.messageOperations.get(messageId)
     }
 
     #joinResults(chatResults: ChatMessage[], only?: string): ChatResult {
@@ -74,17 +97,18 @@
                                 am.messageId === c.messageId
                                     ? am.body + AgenticChatResultStream.resultDelimiter + c.body
                                     : am.body,
-                            ...((c.contextList || acc.contextList) && {
-                                contextList: {
-                                    filePaths: [
-                                        ...(acc.contextList?.filePaths ?? []),
-                                        ...(c.contextList?.filePaths ?? []),
-                                    ],
-                                    rootFolderTitle: c.contextList?.rootFolderTitle
-                                        ? c.contextList.rootFolderTitle
-                                        : (acc.contextList?.rootFolderTitle ?? ''),
-                                },
-                            }),
+                            ...(am.messageId === c.messageId &&
+                                (c.contextList || acc.contextList) && {
+                                    contextList: {
+                                        filePaths: [
+                                            ...(acc.contextList?.filePaths ?? []),
+                                            ...(c.contextList?.filePaths ?? []),
+                                        ],
+                                        rootFolderTitle: c.contextList?.rootFolderTitle
+                                            ? c.contextList.rootFolderTitle
+                                            : (acc.contextList?.rootFolderTitle ?? ''),
+                                    },
+                                }),
                             header: c.header ? { ...c.header } : { ...am.header },
                         })),
                     }