feat: send chat update on undo

Author: ctlai95

chat-client/src/client/utils.ts

@@ -11,7 +11,12 @@ export function toMynahButtons(buttons: Button[] | undefined): ChatItemButton[]
 }
 
 export function toMynahHeader(header: ChatMessage['header']): ChatItemContent['header'] {
-    return { ...header, icon: toMynahIcon(header?.icon), buttons: toMynahButtons(header?.buttons) }
+    return {
+        ...header,
+        icon: toMynahIcon(header?.icon),
+        buttons: toMynahButtons(header?.buttons),
+        status: { ...header?.status, icon: toMynahIcon(header?.status?.icon) },
+    }
 }
 
 export function toDetailsWithoutIcon(

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

@@ -171,6 +171,26 @@ export class AgenticChatController implements ChatHandlers {
             const toolUseId = params.messageId
             try {
                 await this.#undoFileChange(toolUseId)
+                const cachedToolUse = this.#triggerContext.getToolUseLookup().get(toolUseId)
+                if (cachedToolUse) {
+                    this.#features.chat.sendChatUpdate({
+                        tabId: params.tabId,
+                        data: {
+                            messages: [
+                                {
+                                    ...cachedToolUse.chatResult,
+                                    header: {
+                                        ...cachedToolUse.chatResult?.header,
+                                        buttons: cachedToolUse.chatResult?.header?.buttons?.filter(
+                                            button => button.id !== 'undo-changes'
+                                        ),
+                                        status: { status: 'error', icon: 'cancel', text: 'Change discarded' },
+                                    },
+                                },
+                            ],
+                        },
+                    })
+                }
             } catch (err: any) {
                 return { success: false, failureReason: err.message }
             }
@@ -554,6 +574,11 @@ export class AgenticChatController implements ChatHandlers {
                         break
                     case 'fsWrite':
                         const chatResult = await this.#getFsWriteChatResult(toolUse)
+                        const toolUseLookup = this.#triggerContext.getToolUseLookup()
+                        const cachedToolUse = toolUseLookup.get(toolUse.toolUseId)
+                        if (cachedToolUse) {
+                            toolUseLookup.set(toolUse.toolUseId, { ...cachedToolUse, chatResult })
+                        }
                         await chatResultStream.writeResultBlock(chatResult)
                         break
                     case 'executeBash':

server/aws-lsp-codewhisperer/src/language-server/agenticChat/context/agenticChatTriggerContext.ts

@@ -20,6 +20,7 @@ import {
     FileList,
     TextDocument,
     OPEN_WORKSPACE_INDEX_SETTINGS_BUTTON_ID,
+    ChatResult,
 } from '@aws/language-server-runtimes/server-interface'
 import { Features } from '../../types'
 import { DocumentContext, DocumentContextExtractor } from '../../chat/contexts/documentContext'
@@ -58,7 +59,7 @@ export class AgenticChatTriggerContext {
     #lsp: Features['lsp']
     #logging: Features['logging']
     #documentContextExtractor: DocumentContextExtractor
-    #toolUseLookup: Map<string, ToolUse & { oldContent?: string }>
+    #toolUseLookup: Map<string, ToolUse & { oldContent?: string; chatResult?: ChatResult }>
 
     constructor({ workspace, lsp, logging }: Pick<Features, 'workspace' | 'lsp' | 'logging'> & Partial<Features>) {
         this.#workspace = workspace