diff --git a/packages/core/src/amazonq/webview/ui/apps/cwChatConnector.ts b/packages/core/src/amazonq/webview/ui/apps/cwChatConnector.ts index 113ca8dd23c..111350b8e6a 100644 --- a/packages/core/src/amazonq/webview/ui/apps/cwChatConnector.ts +++ b/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 } diff --git a/packages/core/src/codewhispererChat/clients/chat/v0/chat.ts b/packages/core/src/codewhispererChat/clients/chat/v0/chat.ts index ae3d6e8ee26..d7616b214b2 100644 --- a/packages/core/src/codewhispererChat/clients/chat/v0/chat.ts +++ b/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 = new Map() contexts: Map = 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 } diff --git a/packages/core/src/codewhispererChat/controllers/chat/controller.ts b/packages/core/src/codewhispererChat/controllers/chat/controller.ts index 49ecdd9ff2d..a00ec8df99f 100644 --- a/packages/core/src/codewhispererChat/controllers/chat/controller.ts +++ b/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, diff --git a/packages/core/src/codewhispererChat/controllers/chat/messenger/messenger.ts b/packages/core/src/codewhispererChat/controllers/chat/messenger/messenger.ts index af9afeb9eb5..dcc20d7e6c4 100644 --- a/packages/core/src/codewhispererChat/controllers/chat/messenger/messenger.ts +++ b/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, }) diff --git a/packages/core/src/codewhispererChat/tools/fsWrite.ts b/packages/core/src/codewhispererChat/tools/fsWrite.ts index 847d60bc331..54f00514289 100644 --- a/packages/core/src/codewhispererChat/tools/fsWrite.ts +++ b/packages/core/src/codewhispererChat/tools/fsWrite.ts @@ -83,31 +83,39 @@ export class FsWrite { } public async getDiffChanges(): Promise { - 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 { switch (this.params.command) { case 'create': diff --git a/packages/core/src/codewhispererChat/tools/toolUtils.ts b/packages/core/src/codewhispererChat/tools/toolUtils.ts index 93d51244ebf..1549f89b0dc 100644 --- a/packages/core/src/codewhispererChat/tools/toolUtils.ts +++ b/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: