feat(chat): support multi file write with reject and open diff

Author: ctlai95

packages/core/package.nls.json

@@ -454,6 +454,8 @@
     "AWS.amazonq.doc.pillText.makeChanges": "Make changes",
     "AWS.amazonq.inline.invokeChat": "Inline chat",
     "AWS.amazonq.opensettings:": "Open settings",
+    "AWS.amazonq.executeBash.run": "Run",
+    "AWS.amazonq.executeBash.reject": "Reject",
     "AWS.toolkit.lambda.walkthrough.quickpickTitle": "Application Builder Walkthrough",
     "AWS.toolkit.lambda.walkthrough.title": "Get started building your application",
     "AWS.toolkit.lambda.walkthrough.description": "Your quick guide to build an application visually, iterate locally, and deploy to the cloud!",

packages/core/src/amazonq/webview/ui/apps/cwChatConnector.ts

@@ -185,7 +185,10 @@ export class Connector extends BaseConnector {
         this.chatItems.get(tabId)?.set(messageId, { ...item })
     }
 
-    private getCurrentChatItem(tabId: string, messageId: string): ChatItem | undefined {
+    private getCurrentChatItem(tabId: string, messageId: string | undefined): ChatItem | undefined {
+        if (!messageId) {
+            return
+        }
         return this.chatItems.get(tabId)?.get(messageId)
     }
 
@@ -293,7 +296,7 @@ export class Connector extends BaseConnector {
 
     onCustomFormAction(
         tabId: string,
-        messageId: string,
+        messageId: string | undefined,
         action: {
             id: string
             text?: string | undefined
@@ -304,6 +307,10 @@ export class Connector extends BaseConnector {
             return
         }
 
+        if (messageId?.startsWith('tooluse_')) {
+            action.formItemValues = { ...action.formItemValues, toolUseId: messageId }
+        }
+
         this.sendMessageToExtension({
             command: 'form-action-click',
             action: action,

packages/core/src/codewhispererChat/clients/chat/v0/chat.ts

@@ -15,6 +15,7 @@ import { createCodeWhispererChatStreamingClient } from '../../../../shared/clien
 import { createQDeveloperStreamingClient } from '../../../../shared/clients/qDeveloperChatClient'
 import { UserWrittenCodeTracker } from '../../../../codewhisperer/tracker/userWrittenCodeTracker'
 import { PromptMessage } from '../../../controllers/chat/model'
+import { FsWriteBackup } from '../../../../codewhispererChat/tools/fsWrite'
 
 export type ToolUseWithError = {
     toolUse: ToolUse
@@ -33,6 +34,7 @@ export class ChatSession {
     private _showDiffOnFileWrite: boolean = false
     private _context: PromptMessage['context']
     private _pairProgrammingModeOn: boolean = true
+    private _fsWriteBackups: Map<string, FsWriteBackup> = new Map()
     /**
      * True if messages from local history have been sent to session.
      */
@@ -70,6 +72,14 @@ export class ChatSession {
         this._context = context
     }
 
+    public get fsWriteBackups(): Map<string, FsWriteBackup> {
+        return this._fsWriteBackups
+    }
+
+    public setFsWriteBackup(toolUseId: string, backup: FsWriteBackup) {
+        this._fsWriteBackups.set(toolUseId, backup)
+    }
+
     public tokenSource!: vscode.CancellationTokenSource
 
     constructor() {

packages/core/src/codewhispererChat/controllers/chat/controller.ts

@@ -97,7 +97,7 @@ import { maxToolOutputCharacterLength, OutputKind } from '../../tools/toolShared
 import { ToolUtils, Tool, ToolType } from '../../tools/toolUtils'
 import { ChatStream } from '../../tools/chatStream'
 import { ChatHistoryStorage } from '../../storages/chatHistoryStorage'
-import { FsWrite, FsWriteParams } from '../../tools/fsWrite'
+import { FsWriteParams } from '../../tools/fsWrite'
 import { tempDirPath } from '../../../shared/filesystemUtilities'
 import { Database } from '../../../shared/db/chatDb/chatDb'
 import { TabBarController } from './tabBarController'
@@ -722,6 +722,10 @@ export class ChatController {
                             const chatStream = new ChatStream(this.messenger, tabID, triggerID, toolUse, {
                                 requiresAcceptance: false,
                             })
+                            if (tool.type === ToolType.FsWrite && toolUse.toolUseId) {
+                                const backup = await tool.tool.getBackup()
+                                session.setFsWriteBackup(toolUse.toolUseId, backup)
+                            }
                             const output = await ToolUtils.invoke(tool, chatStream)
                             if (output.output.content.length > maxToolOutputCharacterLength) {
                                 throw Error(
@@ -814,13 +818,15 @@ export class ChatController {
             case 'submit-create-prompt':
                 await this.handleCreatePrompt(message)
                 break
-            case 'accept-code-diff':
             case 'run-shell-command':
             case 'generic-tool-execution':
-                await this.closeDiffView()
                 await this.processToolUseMessage(message)
                 break
+            case 'accept-code-diff':
+                await this.closeDiffView()
+                break
             case 'reject-code-diff':
+                await this.restoreBackup(message)
                 await this.closeDiffView()
                 break
             case 'reject-shell-command':
@@ -831,6 +837,22 @@ export class ChatController {
         }
     }
 
+    private async restoreBackup(message: CustomFormActionMessage) {
+        const tabID = message.tabID
+        const toolUseId = message.action.formItemValues?.toolUseId
+        if (!tabID || !toolUseId) {
+            return
+        }
+
+        const session = this.sessionStorage.getSession(tabID)
+        const { content, filePath, isNew } = session.fsWriteBackups.get(toolUseId) ?? {}
+        if (filePath && isNew) {
+            await fs.delete(filePath)
+        } else if (filePath && content !== undefined) {
+            await fs.writeFile(filePath, content)
+        }
+    }
+
     private async processContextSelected(message: ContextSelectedMessage) {
         if (message.tabID && message.contextItem.id === createSavedPromptCommandId) {
             this.handlePromptCreate(message.tabID)
@@ -852,8 +874,15 @@ export class ChatController {
         const session = this.sessionStorage.getSession(message.tabID)
         // Check if user clicked on filePath in the contextList or in the fileListTree and perform the functionality accordingly.
         if (session.showDiffOnFileWrite) {
+            const toolUseId = message.messageId
+            const { filePath, content } = session.fsWriteBackups.get(toolUseId) ?? {}
+            if (!filePath || content === undefined) {
+                return
+            }
+
             try {
                 // Create a temporary file path to show the diff view
+                // TODO: Use amazonQDiffScheme for temp file
                 const pathToArchiveDir = path.join(tempDirPath, 'q-chat')
                 const archivePathExists = await fs.existsDir(pathToArchiveDir)
                 if (archivePathExists) {
@@ -862,39 +891,12 @@ export class ChatController {
                 await fs.mkdir(pathToArchiveDir)
                 const resultArtifactsDir = path.join(pathToArchiveDir, 'resultArtifacts')
                 await fs.mkdir(resultArtifactsDir)
-                const tempFilePath = path.join(
-                    resultArtifactsDir,
-                    `temp-${path.basename((session.toolUseWithError?.toolUse.input as unknown as FsWriteParams).path)}`
-                )
 
-                // If we have existing filePath copy file content from existing file to temporary file.
-                const filePath = (session.toolUseWithError?.toolUse.input as any).path ?? message.filePath
-                const fileExists = await fs.existsFile(filePath)
-                if (fileExists) {
-                    const fileContent = await fs.readFileText(filePath)
-                    await fs.writeFile(tempFilePath, fileContent)
-                }
+                const tempFilePath = path.join(resultArtifactsDir, `temp-${path.basename(filePath)}`)
+                await fs.writeFile(tempFilePath, content)
 
-                // Create a deep clone of the toolUse object and pass this toolUse to FsWrite tool execution to get the modified temporary file.
-                const clonedToolUse = structuredClone(session.toolUseWithError?.toolUse)
-                if (!clonedToolUse) {
-                    return
-                }
-                const input = clonedToolUse.input as unknown as FsWriteParams
-                input.path = tempFilePath
-
-                const fsWrite = new FsWrite(input)
-                await fsWrite.invoke()
-
-                // Check if fileExists=false, If yes, return instead of showing broken diff experience.
-                if (!tempFilePath) {
-                    void vscode.window.showInformationMessage(
-                        'Generated code changes have been reviewed and processed.'
-                    )
-                    return
-                }
-                const leftUri = fileExists ? vscode.Uri.file(filePath) : vscode.Uri.from({ scheme: 'untitled' })
-                const rightUri = vscode.Uri.file(tempFilePath ?? filePath)
+                const leftUri = vscode.Uri.file(tempFilePath)
+                const rightUri = vscode.Uri.file(filePath)
                 const fileName = path.basename(filePath)
                 await vscode.commands.executeCommand(
                     'vscode.diff',

packages/core/src/codewhispererChat/controllers/chat/messenger/messenger.ts

@@ -20,6 +20,7 @@ import {
     UpdateDetailedListMessage,
     CloseDetailedListMessage,
     SelectTabMessage,
+    ChatItemHeader,
 } from '../../../view/connector/connector'
 import { EditorContextCommandType } from '../../../commands/registerCommands'
 import { ChatResponseStream as qdevChatResponseStream } from '@amzn/amazon-q-developer-streaming-client'
@@ -65,6 +66,8 @@ import { noWriteTools, tools } from '../../../constants'
 import { Change } from 'diff'
 import { FsWriteParams } from '../../../tools/fsWrite'
 import { AsyncEventProgressMessage } from '../../../../amazonq/commons/connector/connectorMessages'
+import { localize } from '../../../../shared/utilities/vsCodeUtils'
+import { getDiffLinesFromChanges } from '../../../../shared/utilities/diffUtils'
 
 export type StaticTextResponseType =
     | 'quick-action-help'
@@ -496,49 +499,44 @@ export class Messenger {
         changeList?: Change[]
     ) {
         const buttons: ChatItemButton[] = []
-        let fileList: ChatItemContent['fileList'] = undefined
-        let shellCommandHeader = undefined
+        let header: ChatItemHeader | undefined = undefined
+        let fullWidth: boolean | undefined = undefined
+        let padding: boolean | undefined = undefined
+        let codeBlockActions: ChatItemContent['codeBlockActions'] = undefined
         if (toolUse?.name === ToolType.ExecuteBash && message.startsWith('```shell')) {
             if (validation.requiresAcceptance) {
-                buttons.push({
-                    id: 'run-shell-command',
-                    text: 'Run',
-                    status: 'main',
-                    icon: 'play' as MynahIconsType,
-                })
-                buttons.push({
-                    id: 'reject-shell-command',
-                    text: 'Reject',
-                    status: 'clear',
-                    icon: 'cancel' as MynahIconsType,
-                })
-            }
-
-            shellCommandHeader = {
-                icon: 'code-block' as MynahIconsType,
-                body: 'shell',
-                buttons: buttons,
+                const buttons: ChatItemButton[] = [
+                    {
+                        id: 'run-shell-command',
+                        text: localize('AWS.amazonq.executeBash.run', 'Run'),
+                        status: 'main',
+                        icon: 'play' as MynahIconsType,
+                    },
+                    {
+                        id: 'reject-shell-command',
+                        text: localize('AWS.amazonq.executeBash.reject', 'Reject'),
+                        status: 'clear',
+                        icon: 'cancel' as MynahIconsType,
+                    },
+                ]
+                header = {
+                    icon: 'code-block' as MynahIconsType,
+                    body: 'shell',
+                    buttons,
+                }
             }
-
             if (validation.warning) {
                 message = validation.warning + message
             }
+            fullWidth = true
+            padding = false
+            // eslint-disable-next-line unicorn/no-null
+            codeBlockActions = { 'insert-to-cursor': null, copy: null }
         } else if (toolUse?.name === ToolType.FsWrite) {
             const input = toolUse.input as unknown as FsWriteParams
             const fileName = path.basename(input.path)
-            const changes = changeList?.reduce(
-                (acc, { count = 0, added, removed }) => {
-                    if (added) {
-                        acc.added += count
-                    } else if (removed) {
-                        acc.deleted += count
-                    }
-                    return acc
-                },
-                { added: 0, deleted: 0 }
-            )
-            // FileList
-            fileList = {
+            const changes = getDiffLinesFromChanges(changeList)
+            const fileList: ChatItemContent['fileList'] = {
                 fileTreeTitle: '',
                 hideFileCount: true,
                 filePaths: [fileName],
@@ -550,17 +548,27 @@ export class Messenger {
                     },
                 },
             }
-            // Buttons
-            buttons.push({
-                id: 'reject-code-diff',
-                status: 'clear',
-                icon: 'cancel' as MynahIconsType,
-            })
-            buttons.push({
-                id: 'accept-code-diff',
-                status: 'clear',
-                icon: 'ok' as MynahIconsType,
-            })
+            const buttons: ChatItemButton[] = [
+                {
+                    id: 'reject-code-diff',
+                    status: 'clear',
+                    icon: 'cancel' as MynahIconsType,
+                },
+                {
+                    id: 'accept-code-diff',
+                    status: 'clear',
+                    icon: 'ok' as MynahIconsType,
+                },
+            ]
+            header = {
+                icon: 'code-block' as MynahIconsType,
+                buttons,
+                fileList,
+            }
+            fullWidth = true
+            padding = false
+            // eslint-disable-next-line unicorn/no-null
+            codeBlockActions = { 'insert-to-cursor': null, copy: null }
         }
 
         this.dispatcher.sendChatMessage(
@@ -577,23 +585,11 @@ export class Messenger {
                     codeBlockLanguage: undefined,
                     contextList: undefined,
                     canBeVoted: false,
-                    buttons:
-                        toolUse?.name === ToolType.FsWrite || toolUse?.name === ToolType.ExecuteBash
-                            ? undefined
-                            : buttons,
-                    fullWidth: toolUse?.name === ToolType.FsWrite || toolUse?.name === ToolType.ExecuteBash,
-                    padding: !(toolUse?.name === ToolType.FsWrite || toolUse?.name === ToolType.ExecuteBash),
-                    header:
-                        toolUse?.name === ToolType.FsWrite
-                            ? { icon: 'code-block' as MynahIconsType, buttons: buttons, fileList: fileList }
-                            : toolUse?.name === ToolType.ExecuteBash
-                              ? shellCommandHeader
-                              : undefined,
-                    codeBlockActions:
-                        // eslint-disable-next-line unicorn/no-null, prettier/prettier
-                        toolUse?.name === ToolType.FsWrite || toolUse?.name === ToolType.ExecuteBash
-                            ? { 'insert-to-cursor': undefined, copy: undefined }
-                            : undefined,
+                    buttons,
+                    fullWidth,
+                    padding,
+                    header,
+                    codeBlockActions,
                 },
                 tabID
             )