fix: further improvements for thinking/loading (#1125)

## Problem
- thinking don't show up in second round of messages because the ID used wasn't unique
- sometimes it takes a while from us rendering text response to processing toolUse
- listDir/read/search shows result before running the tool

## Solution
- switch to use `uuid` instead in iterations
- add additional loading message between text response to toolUse processing
- move list/read/search result to after tool run

Author: zixlin7

chat-client/src/client/mynahUi.ts

@@ -456,7 +456,7 @@ export const createMynahUi = (
             const chatItems = store.chatItems || []
             const updatedItems = chatItems.map(item => ({
                 ...item,
-                type: item.type === ChatItemType.ANSWER_STREAM ? ChatItemType.ANSWER : item.type,
+                type: item.type === ChatItemType.ANSWER_STREAM && !item.body ? ChatItemType.ANSWER : item.type,
             }))
             mynahUi.updateStore(tabId, { loadingChat: false, cancelButtonWhenLoading: true, chatItems: updatedItems })
             messager.onStopChatResponse(tabId)

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

@@ -894,7 +894,7 @@ describe('AgenticChatController', () => {
 
             const chatResult = await chatResultPromise
 
-            sinon.assert.callCount(testFeatures.lsp.sendProgress, mockChatResponseList.length + 1) // response length + loading message
+            sinon.assert.callCount(testFeatures.lsp.sendProgress, mockChatResponseList.length + 2) // response length + 2 loading messages
             assert.deepStrictEqual(chatResult, {
                 additionalMessages: [],
                 body: '\n\nHello World!',
@@ -911,7 +911,7 @@ describe('AgenticChatController', () => {
 
             const chatResult = await chatResultPromise
 
-            sinon.assert.callCount(testFeatures.lsp.sendProgress, mockChatResponseList.length + 1) // response length + loading message
+            sinon.assert.callCount(testFeatures.lsp.sendProgress, mockChatResponseList.length + 2) // response length + 2 loading message
             assert.deepStrictEqual(chatResult, {
                 additionalMessages: [],
                 body: '\n\nHello World!',

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

@@ -427,7 +427,7 @@ export class AgenticChatController implements ChatHandlers {
             const conversationId = conversationIdentifier ?? ''
 
             // show loading message while we process request
-            loadingMessageId = `loading-${conversationId}-${iterationCount}`
+            loadingMessageId = `loading-${uuid()}-${iterationCount}`
             await chatResultStream.writeResultBlock({ messageId: loadingMessageId, type: 'answer' })
 
             if (!currentMessage || !conversationId) {
@@ -481,6 +481,10 @@ export class AgenticChatController implements ChatHandlers {
                 documentReference
             )
 
+            // show loading message after we render text response and before processing toolUse
+            loadingMessageId = `loading-${uuid()}-${iterationCount}`
+            await chatResultStream.writeResultBlock({ messageId: loadingMessageId, type: 'answer' })
+
             //  Add the current assistantResponse message to the history DB
             if (result.data?.chatResult.body !== undefined) {
                 this.#chatHistoryDb.addMessage(tabId, 'cwc', conversationIdentifier ?? '', {
@@ -503,6 +507,13 @@ export class AgenticChatController implements ChatHandlers {
             // Check if we have any tool uses that need to be processed
             const pendingToolUses = this.#getPendingToolUses(result.data?.toolUses || {})
 
+            // remove the temp loading message when we are going to process toolUse
+            if (loadingMessageId) {
+                await chatResultStream.removeResultBlock(loadingMessageId)
+                this.#features.chat.sendChatUpdate({ tabId, state: { inProgress: false } })
+                loadingMessageId = undefined
+            }
+
             if (pendingToolUses.length === 0) {
                 // No more tool uses, we're done
                 finalResult = result
@@ -629,12 +640,7 @@ export class AgenticChatController implements ChatHandlers {
                 await chatResultStream.writeResultBlock({ messageId: loadingMessageId, type: 'answer' })
                 this.#features.chat.sendChatUpdate({ tabId, state: { inProgress: true } })
 
-                if (['fsRead', 'listDirectory', 'fileSearch'].includes(toolUse.name)) {
-                    const initialListDirResult = this.#processReadOrListOrSearch(toolUse, chatResultStream)
-                    if (initialListDirResult) {
-                        await chatResultStream.writeResultBlock(initialListDirResult)
-                    }
-                } else if (toolUse.name === 'fsWrite') {
+                if (toolUse.name === 'fsWrite') {
                     const input = toolUse.input as unknown as FsWriteParams
                     const document = await this.#triggerContext.getTextDocument(input.path)
                     this.#triggerContext
@@ -673,6 +679,11 @@ export class AgenticChatController implements ChatHandlers {
                     case 'fsRead':
                     case 'listDirectory':
                     case 'fileSearch':
+                        const initialListDirResult = this.#processReadOrListOrSearch(toolUse, chatResultStream)
+                        if (initialListDirResult) {
+                            await chatResultStream.writeResultBlock(initialListDirResult)
+                        }
+                        break
                     // no need to write tool result for listDir,fsRead,fileSearch into chat stream
                     case 'executeBash':
                         // no need to write tool result for listDir and fsRead into chat stream