fix: further improvements for thinking/loading

Author: zixlin7

chat-client/src/client/mynahUi.ts

@@ -433,7 +433,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.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 ?? '', {
@@ -509,6 +513,13 @@ export class AgenticChatController implements ChatHandlers {
                 break
             }
 
+            // 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
+            }
+
             // Process tool uses and update the request input for the next iteration
             const toolResults = await this.#processToolUses(pendingToolUses, chatResultStream, session, tabId, token)
             currentRequestInput = this.#updateRequestInputWithToolResults(currentRequestInput, toolResults)
@@ -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