feat(amazonq): read tool ui revamp (#2113)

ctlai95 · web-flow · commit 301f2c59b17e · 2025-08-12T15:14:04.000-07:00
* feat(amazonq): read tool ui revamp * feat(amazonq): read tool message revamp (#2049) * feat(amazonq): read tool message revamp * fix tests * feat: file search ui (#2078) * feat: file search ui * fix tests * fix integration tests * remove unnecessary type check * fix: use quotes instead of backticks * fix header update issue * fix integration test
diff --git a/chat-client/src/client/mynahUi.ts b/chat-client/src/client/mynahUi.ts
@@ -1353,10 +1353,15 @@ export const createMynahUi = (
                     fileTreeTitle: '',
                     hideFileCount: true,
                     details: toDetailsWithoutIcon(header.fileList.details),
+                    renderAsPills:
+                        !header.fileList.details ||
+                        (Object.values(header.fileList.details).every(detail => !detail.changes) &&
+                            (!header.buttons || !header.buttons.some(button => button.id === 'undo-changes')) &&
+                            !header.status?.icon),
                 }
             }
             if (!isPartialResult) {
-                if (processedHeader) {
+                if (processedHeader && !message.header?.status) {
                     processedHeader.status = undefined
                 }
             }
@@ -1369,7 +1374,8 @@ export const createMynahUi = (
                 processedHeader.buttons !== null &&
                 processedHeader.buttons.length > 0) ||
                 processedHeader.status !== undefined ||
-                processedHeader.icon !== undefined)
+                processedHeader.icon !== undefined ||
+                processedHeader.fileList !== undefined)
 
         const padding =
             message.type === 'tool' ? (fileList ? true : message.messageId?.endsWith('_permission')) : undefined
@@ -1380,8 +1386,10 @@ export const createMynahUi = (
         // Adding this conditional check to show the stop message in the center.
         const contentHorizontalAlignment: ChatItem['contentHorizontalAlignment'] = undefined
 
-        // If message.header?.status?.text is Stopped or Rejected or Ignored or Completed etc.. card should be in disabled state.
-        const shouldMute = message.header?.status?.text !== undefined && message.header?.status?.text !== 'Completed'
+        // If message.header?.status?.text is Stopped or Rejected or Ignored etc.. card should be in disabled state.
+        const shouldMute =
+            message.header?.status?.text !== undefined &&
+            ['Stopped', 'Rejected', 'Ignored', 'Failed', 'Error'].includes(message.header.status.text)
 
         return {
             body: message.body,
diff --git a/integration-tests/q-agentic-chat-server/src/tests/agenticChatInteg.test.ts b/integration-tests/q-agentic-chat-server/src/tests/agenticChatInteg.test.ts
@@ -169,11 +169,11 @@ describe('Q Agentic Chat Server Integration Tests', async () => {
 
         expect(decryptedResult.additionalMessages).to.be.an('array')
         const fsReadMessage = decryptedResult.additionalMessages?.find(
-            msg => msg.type === 'tool' && msg.fileList?.rootFolderTitle === '1 file read'
+            msg => msg.type === 'tool' && msg.header?.body === '1 file read'
         )
         expect(fsReadMessage).to.exist
         const expectedPath = path.join(rootPath, 'test.py')
-        const actualPaths = fsReadMessage?.fileList?.filePaths?.map(normalizePath) || []
+        const actualPaths = fsReadMessage?.header?.fileList?.filePaths?.map(normalizePath) || []
         expect(actualPaths).to.include.members([normalizePath(expectedPath)])
         expect(fsReadMessage?.messageId?.startsWith('tooluse_')).to.be.true
     })
@@ -191,10 +191,10 @@ describe('Q Agentic Chat Server Integration Tests', async () => {
 
         expect(decryptedResult.additionalMessages).to.be.an('array')
         const listDirectoryMessage = decryptedResult.additionalMessages?.find(
-            msg => msg.type === 'tool' && msg.fileList?.rootFolderTitle === '1 directory listed'
+            msg => msg.type === 'tool' && msg.header?.body === '1 directory listed'
         )
         expect(listDirectoryMessage).to.exist
-        const actualPaths = listDirectoryMessage?.fileList?.filePaths?.map(normalizePath) || []
+        const actualPaths = listDirectoryMessage?.header?.fileList?.filePaths?.map(normalizePath) || []
         expect(actualPaths).to.include.members([normalizePath(rootPath)])
         expect(listDirectoryMessage?.messageId?.startsWith('tooluse_')).to.be.true
     })
@@ -371,11 +371,12 @@ describe('Q Agentic Chat Server Integration Tests', async () => {
 
         expect(decryptedResult.additionalMessages).to.be.an('array')
         const fileSearchMessage = decryptedResult.additionalMessages?.find(
-            msg => msg.type === 'tool' && msg.fileList?.rootFolderTitle === '1 directory searched'
+            msg => msg.type === 'tool' && msg.header?.body === 'Searched for "test" in '
         )
         expect(fileSearchMessage).to.exist
         expect(fileSearchMessage?.messageId?.startsWith('tooluse_')).to.be.true
-        const actualPaths = fileSearchMessage?.fileList?.filePaths?.map(normalizePath) || []
+        expect(fileSearchMessage?.header?.status?.text).to.equal('3 results found')
+        const actualPaths = fileSearchMessage?.header?.fileList?.filePaths?.map(normalizePath) || []
         expect(actualPaths).to.include.members([normalizePath(rootPath)])
     })
 })
diff --git a/server/aws-lsp-codewhisperer/src/language-server/agenticChat/agenticChatController.test.ts b/server/aws-lsp-codewhisperer/src/language-server/agenticChat/agenticChatController.test.ts
@@ -451,7 +451,7 @@ describe('AgenticChatController', () => {
 
             assert.deepStrictEqual(chatResult, {
                 additionalMessages: [],
-                body: '\n\nHello World!',
+                body: '\nHello World!',
                 messageId: 'mock-message-id',
                 buttons: [],
                 codeReference: [],
@@ -1150,7 +1150,7 @@ describe('AgenticChatController', () => {
             sinon.assert.callCount(testFeatures.lsp.sendProgress, mockChatResponseList.length + 1) // response length + 1 loading messages
             assert.deepStrictEqual(chatResult, {
                 additionalMessages: [],
-                body: '\n\nHello World!',
+                body: '\nHello World!',
                 messageId: 'mock-message-id',
                 codeReference: [],
                 buttons: [],
@@ -1169,7 +1169,7 @@ describe('AgenticChatController', () => {
             sinon.assert.callCount(testFeatures.lsp.sendProgress, mockChatResponseList.length + 1) // response length + 1 loading message
             assert.deepStrictEqual(chatResult, {
                 additionalMessages: [],
-                body: '\n\nHello World!',
+                body: '\nHello World!',
                 messageId: 'mock-message-id',
                 buttons: [],
                 codeReference: [],
diff --git a/server/aws-lsp-codewhisperer/src/language-server/agenticChat/agenticChatController.ts b/server/aws-lsp-codewhisperer/src/language-server/agenticChat/agenticChatController.ts
@@ -169,7 +169,7 @@ import { ExecuteBash, ExecuteBashParams } from './tools/executeBash'
 import { ExplanatoryParams, InvokeOutput, ToolApprovalException } from './tools/toolShared'
 import { validatePathBasic, validatePathExists, validatePaths as validatePathsSync } from './utils/pathValidation'
 import { GrepSearch, SanitizedRipgrepOutput } from './tools/grepSearch'
-import { FileSearch, FileSearchParams } from './tools/fileSearch'
+import { FileSearch, FileSearchParams, isFileSearchParams } from './tools/fileSearch'
 import { FsReplace, FsReplaceParams } from './tools/fsReplace'
 import { loggingUtils, timeoutUtils } from '@aws/lsp-core'
 import { diffLines } from 'diff'
@@ -1695,8 +1695,7 @@ export class AgenticChatController implements ChatHandlers {
                 // 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 (![FS_WRITE, FS_REPLACE, FS_READ, LIST_DIRECTORY].includes(toolUse.name)) {
+                if (![FS_WRITE, FS_REPLACE].includes(toolUse.name)) {
                     await this.#showUndoAllIfRequired(chatResultStream, session)
                 }
                 // fsWrite can take a long time, so we render fsWrite  Explanatory upon partial streaming responses.
@@ -1911,10 +1910,19 @@ export class AgenticChatController implements ChatHandlers {
                 switch (toolUse.name) {
                     case FS_READ:
                     case LIST_DIRECTORY:
+                        const readToolResult = await this.#processReadTool(toolUse, chatResultStream)
+                        if (readToolResult) {
+                            await chatResultStream.writeResultBlock(readToolResult)
+                        }
+                        break
                     case FILE_SEARCH:
-                        const initialListDirResult = this.#processReadOrListOrSearch(toolUse, chatResultStream)
-                        if (initialListDirResult) {
-                            await chatResultStream.writeResultBlock(initialListDirResult)
+                        if (isFileSearchParams(toolUse.input)) {
+                            await this.#processFileSearchTool(
+                                toolUse.input,
+                                toolUse.toolUseId,
+                                result,
+                                chatResultStream
+                            )
                         }
                         break
                     // no need to write tool result for listDir,fsRead,fileSearch into chat stream
@@ -2315,7 +2323,6 @@ export class AgenticChatController implements ChatHandlers {
         }
 
         const toolMsgId = toolUse.toolUseId!
-        const chatMsgId = chatResultStream.getResult().messageId
         let headerEmitted = false
 
         const initialHeader: ChatMessage['header'] = {
@@ -2353,13 +2360,6 @@ export class AgenticChatController implements ChatHandlers {
                     header: completedHeader,
                 })
 
-                await chatResultStream.writeResultBlock({
-                    type: 'answer',
-                    messageId: chatMsgId,
-                    body: '',
-                    header: undefined,
-                })
-
                 this.#stoppedToolUses.add(toolMsgId)
             },
         })
@@ -2877,70 +2877,135 @@ export class AgenticChatController implements ChatHandlers {
         }
     }
 
-    #processReadOrListOrSearch(toolUse: ToolUse, chatResultStream: AgenticChatResultStream): ChatMessage | undefined {
-        let messageIdToUpdate = toolUse.toolUseId!
-        const currentId = chatResultStream.getMessageIdToUpdateForTool(toolUse.name!)
+    async #processFileSearchTool(
+        toolInput: FileSearchParams,
+        toolUseId: string,
+        result: InvokeOutput,
+        chatResultStream: AgenticChatResultStream
+    ): Promise<void> {
+        if (typeof result.output.content !== 'string') return
 
-        if (currentId) {
-            messageIdToUpdate = currentId
-        } else {
-            chatResultStream.setMessageIdToUpdateForTool(toolUse.name!, messageIdToUpdate)
+        const { queryName, path: inputPath } = toolInput
+        const resultCount = result.output.content
+            .split('\n')
+            .filter(line => line.trim().startsWith('[F]') || line.trim().startsWith('[D]')).length
+
+        const chatMessage: ChatMessage = {
+            type: 'tool',
+            messageId: toolUseId,
+            header: {
+                body: `Searched for "${queryName}" in `,
+                icon: 'search',
+                status: {
+                    text: `${resultCount} result${resultCount !== 1 ? 's' : ''} found`,
+                },
+                fileList: {
+                    filePaths: [inputPath],
+                    details: {
+                        [inputPath]: {
+                            description: inputPath,
+                            visibleName: path.basename(inputPath),
+                            clickable: false,
+                        },
+                    },
+                },
+            },
         }
-        let currentPaths = []
+        await chatResultStream.writeResultBlock(chatMessage)
+    }
+
+    async #processReadTool(
+        toolUse: ToolUse,
+        chatResultStream: AgenticChatResultStream
+    ): Promise<ChatMessage | undefined> {
+        let currentPaths: string[] = []
         if (toolUse.name === FS_READ) {
-            currentPaths = (toolUse.input as unknown as FsReadParams)?.paths
+            currentPaths = (toolUse.input as unknown as FsReadParams)?.paths || []
+        } else if (toolUse.name === LIST_DIRECTORY) {
+            const singlePath = (toolUse.input as unknown as ListDirectoryParams)?.path
+            if (singlePath) {
+                currentPaths = [singlePath]
+            }
+        } else if (toolUse.name === FILE_SEARCH) {
+            const queryName = (toolUse.input as unknown as FileSearchParams)?.queryName
+            if (queryName) {
+                currentPaths = [queryName]
+            }
         } else {
-            currentPaths.push((toolUse.input as unknown as ListDirectoryParams | FileSearchParams)?.path)
+            return
         }
 
-        if (!currentPaths) return
+        if (currentPaths.length === 0) return
 
-        for (const currentPath of currentPaths) {
-            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 }],
-                }
-                chatResultStream.addMessageOperation(messageIdToUpdate, toolUse.name!, [
-                    ...existingPaths,
-                    currentFileDetail,
-                ])
+        // Check if the last message is the same tool type
+        const lastMessage = chatResultStream.getLastMessage()
+        const isSameToolType =
+            lastMessage?.type === 'tool' && lastMessage.header?.icon === this.#toolToIcon(toolUse.name)
+
+        let allPaths = currentPaths
+
+        if (isSameToolType && lastMessage.messageId) {
+            // Combine with existing paths and overwrite the last message
+            const existingPaths = lastMessage.header?.fileList?.filePaths || []
+            allPaths = [...existingPaths, ...currentPaths]
+
+            const blockId = chatResultStream.getMessageBlockId(lastMessage.messageId)
+            if (blockId !== undefined) {
+                // Create the updated message with combined paths
+                const updatedMessage = this.#createFileListToolMessage(toolUse, allPaths, lastMessage.messageId)
+                // Overwrite the existing block
+                await chatResultStream.overwriteResultBlock(updatedMessage, blockId)
+                return undefined // Don't return a message since we already wrote it
             }
         }
+
+        // Create new message with current paths
+        return this.#createFileListToolMessage(toolUse, allPaths, toolUse.toolUseId!)
+    }
+
+    #createFileListToolMessage(toolUse: ToolUse, filePaths: string[], messageId: string): ChatMessage {
+        const itemCount = filePaths.length
         let title: string
-        const itemCount = chatResultStream.getMessageOperation(messageIdToUpdate)?.filePaths.length
-        const filePathsPushed = chatResultStream.getMessageOperation(messageIdToUpdate)?.filePaths ?? []
-        if (!itemCount) {
+        if (itemCount === 0) {
             title = 'Gathering context'
         } else {
             title =
                 toolUse.name === FS_READ
                     ? `${itemCount} file${itemCount > 1 ? 's' : ''} read`
-                    : toolUse.name === FILE_SEARCH
-                      ? `${itemCount} ${itemCount === 1 ? 'directory' : 'directories'} searched`
-                      : `${itemCount} ${itemCount === 1 ? 'directory' : 'directories'} listed`
+                    : toolUse.name === LIST_DIRECTORY
+                      ? `${itemCount} ${itemCount === 1 ? 'directory' : 'directories'} listed`
+                      : ''
         }
         const details: Record<string, FileDetails> = {}
-        for (const item of filePathsPushed) {
-            details[item.relativeFilePath] = {
-                lineRanges: item.lineRanges,
-                description: item.relativeFilePath,
+        for (const filePath of filePaths) {
+            details[filePath] = {
+                description: filePath,
+                visibleName: path.basename(filePath),
+                clickable: toolUse.name === FS_READ,
             }
         }
-
-        const fileList: FileList = {
-            rootFolderTitle: title,
-            filePaths: filePathsPushed.map(item => item.relativeFilePath),
-            details,
-        }
         return {
             type: 'tool',
-            fileList,
-            messageId: messageIdToUpdate,
-            body: '',
+            header: {
+                body: title,
+                icon: this.#toolToIcon(toolUse.name),
+                fileList: {
+                    filePaths,
+                    details,
+                },
+            },
+            messageId,
+        }
+    }
+
+    #toolToIcon(toolName: string | undefined): string | undefined {
+        switch (toolName) {
+            case FS_READ:
+                return 'eye'
+            case LIST_DIRECTORY:
+                return 'check-list'
+            default:
+                return undefined
         }
     }
 
@@ -2956,14 +3021,7 @@ export class AgenticChatController implements ChatHandlers {
             return undefined
         }
 
-        let messageIdToUpdate = toolUse.toolUseId!
-        const currentId = chatResultStream.getMessageIdToUpdateForTool(toolUse.name!)
-
-        if (currentId) {
-            messageIdToUpdate = currentId
-        } else {
-            chatResultStream.setMessageIdToUpdateForTool(toolUse.name!, messageIdToUpdate)
-        }
+        const messageIdToUpdate = toolUse.toolUseId!
 
         // Extract search results from the tool output
         const output = result.output.content as SanitizedRipgrepOutput
diff --git a/server/aws-lsp-codewhisperer/src/language-server/agenticChat/agenticChatResultStream.ts b/server/aws-lsp-codewhisperer/src/language-server/agenticChat/agenticChatResultStream.ts
diff --git a/server/aws-lsp-codewhisperer/src/language-server/agenticChat/tools/fileSearch.ts b/server/aws-lsp-codewhisperer/src/language-server/agenticChat/tools/fileSearch.ts
