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

Author: ctlai95

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

@@ -314,7 +314,7 @@ export class Connector extends BaseConnector {
 
         if (
             !this.onChatAnswerUpdated ||
-            !['accept-code-diff', 'reject-code-diff', 'confirm-tool-use'].includes(action.id)
+            !(['accept-code-diff', 'confirm-tool-use'].includes(action.id) || action.id.startsWith('reject-code-diff'))
         ) {
             return
         }
@@ -341,17 +341,6 @@ export class Connector extends BaseConnector {
                     answer.body = ' '
                 }
                 break
-            case 'reject-code-diff':
-                if (answer.header) {
-                    answer.header.status = {
-                        icon: 'cancel' as MynahIconsType,
-                        text: 'Rejected',
-                        status: 'error',
-                    }
-                    answer.header.buttons = []
-                    answer.body = ' '
-                }
-                break
             case 'confirm-tool-use':
                 answer.buttons = [
                     {
@@ -367,6 +356,17 @@ export class Connector extends BaseConnector {
             default:
                 break
         }
+        if (action.id.startsWith('reject-code-diff')) {
+            if (answer.header) {
+                answer.header.status = {
+                    icon: 'cancel' as MynahIconsType,
+                    text: 'Rejected',
+                    status: 'error',
+                }
+                answer.header.buttons = []
+                answer.body = ' '
+            }
+        }
 
         if (currentChatItem && answer.messageId) {
             const updatedItem = { ...currentChatItem, ...answer }

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

@@ -28,6 +28,7 @@ export class ChatSession {
     private _showDiffOnFileWrite: boolean = false
     private _context: PromptMessage['context']
     private _pairProgrammingModeOn: boolean = true
+    private _fsWriteBackups: Map<string, { filePath: string; content: string; isNew: boolean }> = new Map()
 
     contexts: Map<string, { first: number; second: number }[]> = new Map()
     // TODO: doesn't handle the edge case when two files share the same relativePath string but from different root
@@ -53,6 +54,20 @@ export class ChatSession {
         this._toolUse = toolUse
     }
 
+    public get fsWriteBackups() {
+        return this._fsWriteBackups
+    }
+
+    public setFsWriteBackups(
+        toolUseId: string | undefined,
+        content: { filePath: string; content: string; isNew: boolean }
+    ) {
+        if (!toolUseId) {
+            return
+        }
+        this._fsWriteBackups.set(toolUseId, content)
+    }
+
     public get context(): PromptMessage['context'] {
         return this._context
     }

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

@@ -92,7 +92,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'
 
 export interface ChatControllerMessagePublishers {
@@ -739,6 +739,10 @@ export class ChatController {
                         const chatStream = new ChatStream(this.messenger, tabID, triggerID, toolUse, {
                             requiresAcceptance: false,
                         })
+                        if (tool.type === ToolType.FsWrite) {
+                            const backup = await tool.tool.getOldContent()
+                            session.setFsWriteBackups(toolUse.toolUseId, backup)
+                        }
                         const output = await ToolUtils.invoke(tool, chatStream)
                         if (output.output.content.length > maxToolOutputCharacterLength) {
                             throw Error(
@@ -818,22 +822,29 @@ export class ChatController {
             case 'submit-create-prompt':
                 await this.handleCreatePrompt(message)
                 break
-            case 'accept-code-diff':
-                await this.closeDiffView()
-                break
             case 'confirm-tool-use':
             case 'generic-tool-execution':
                 await this.processToolUseMessage(message)
                 break
-            case 'reject-code-diff':
-                await this.closeDiffView()
-                break
             case 'tool-unavailable':
                 await this.processUnavailableToolUseMessage(message)
                 break
             default:
                 getLogger().warn(`Unhandled action: ${message.action.id}`)
         }
+
+        if (message.action.id.startsWith('reject-code-diff')) {
+            // revert the changes
+            const toolUseId = message.action.id.split('/')[1]
+            const backups = this.sessionStorage.getSession(message.tabID!).fsWriteBackups
+            const { filePath, content, isNew } = backups.get(toolUseId) ?? {}
+            if (filePath && isNew) {
+                await fs.delete(filePath)
+            } else if (filePath && content !== undefined) {
+                await fs.writeFile(filePath, content)
+            }
+            await this.closeDiffView()
+        }
     }
 
     private async processContextSelected(message: ContextSelectedMessage) {
@@ -855,8 +866,10 @@ export class ChatController {
 
     private async processFileClickMessage(message: FileClick) {
         const session = this.sessionStorage.getSession(message.tabID)
+        const toolUseId = message.messageId
+        const backup = session.fsWriteBackups.get(toolUseId)
         // Check if user clicked on filePath in the contextList or in the fileListTree and perform the functionality accordingly.
-        if (session.showDiffOnFileWrite) {
+        if (session.showDiffOnFileWrite && backup?.filePath) {
             try {
                 // Create a temporary file path to show the diff view
                 const pathToArchiveDir = path.join(tempDirPath, 'q-chat')
@@ -867,40 +880,13 @@ 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.toolUse?.input as unknown as FsWriteParams).path)}`
-                )
+                const tempFilePath = path.join(resultArtifactsDir, `temp-${path.basename(backup.filePath)}`)
 
-                // If we have existing filePath copy file content from existing file to temporary file.
-                const filePath = (session.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)
-                }
+                await fs.writeFile(tempFilePath, backup.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.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 fileName = path.basename(filePath)
+                const leftUri = vscode.Uri.file(tempFilePath)
+                const rightUri = vscode.Uri.file(backup.filePath)
+                const fileName = path.basename(backup.filePath)
                 await vscode.commands.executeCommand(
                     'vscode.diff',
                     leftUri,

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

@@ -489,7 +489,7 @@ export class Messenger {
             }
             // Buttons
             buttons.push({
-                id: 'reject-code-diff',
+                id: `reject-code-diff/${toolUse.toolUseId}`,
                 status: 'clear',
                 icon: 'cancel' as MynahIconsType,
             })

packages/core/src/codewhispererChat/tools/fsWrite.ts

@@ -83,31 +83,39 @@ export class FsWrite {
     }
 
     public async getDiffChanges(): Promise<Change[]> {
-        const sanitizedPath = sanitizePath(this.params.path)
+        const { filePath, content: oldContent } = await this.getOldContent()
         let newContent
-        let oldContent
-        try {
-            oldContent = await fs.readFileText(sanitizedPath)
-        } catch (err) {
-            oldContent = ''
-        }
         switch (this.params.command) {
             case 'create':
                 newContent = this.getCreateCommandText(this.params)
                 break
             case 'strReplace':
-                newContent = await this.getStrReplaceContent(this.params, sanitizedPath)
+                newContent = await this.getStrReplaceContent(this.params, filePath)
                 break
             case 'insert':
-                newContent = await this.getInsertContent(this.params, sanitizedPath)
+                newContent = await this.getInsertContent(this.params, filePath)
                 break
             case 'append':
-                newContent = await this.getAppendContent(this.params, sanitizedPath)
+                newContent = await this.getAppendContent(this.params, filePath)
                 break
         }
         return diffLines(oldContent, newContent)
     }
 
+    public async getOldContent(): Promise<{ filePath: string; content: string; isNew: boolean }> {
+        const sanitizedPath = sanitizePath(this.params.path)
+        let oldContent
+        let isNew
+        try {
+            oldContent = await fs.readFileText(sanitizedPath)
+            isNew = false
+        } catch (err) {
+            oldContent = ''
+            isNew = true
+        }
+        return { filePath: sanitizedPath, content: oldContent, isNew }
+    }
+
     public async validate(): Promise<void> {
         switch (this.params.command) {
             case 'create':

packages/core/src/codewhispererChat/tools/toolUtils.ts

@@ -42,7 +42,7 @@ export class ToolUtils {
             case ToolType.FsRead:
                 return { requiresAcceptance: false }
             case ToolType.FsWrite:
-                return { requiresAcceptance: true }
+                return { requiresAcceptance: false }
             case ToolType.ExecuteBash:
                 return tool.tool.requiresAcceptance()
             case ToolType.ListDirectory: