From be636c877b57d1eb4a231b635d2f840ea3649e32 Mon Sep 17 00:00:00 2001 From: Tai Lai Date: Mon, 7 Apr 2025 16:57:31 -0700 Subject: [PATCH] feat(chat): support multi file write with reject and open diff --- packages/core/package.nls.json | 2 + .../webview/ui/apps/cwChatConnector.ts | 11 +- .../codewhispererChat/clients/chat/v0/chat.ts | 10 ++ .../controllers/chat/controller.ts | 70 +++++----- .../controllers/chat/messenger/messenger.ts | 120 +++++++++--------- .../src/codewhispererChat/tools/fsWrite.ts | 38 +++--- .../src/codewhispererChat/tools/toolUtils.ts | 2 +- .../core/src/shared/utilities/diffUtils.ts | 25 ++++ .../tools/toolShared.test.ts | 4 +- 9 files changed, 166 insertions(+), 116 deletions(-) diff --git a/packages/core/package.nls.json b/packages/core/package.nls.json index 22d2ce3fc3b..50f8f3ffe28 100644 --- a/packages/core/package.nls.json +++ b/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!", diff --git a/packages/core/src/amazonq/webview/ui/apps/cwChatConnector.ts b/packages/core/src/amazonq/webview/ui/apps/cwChatConnector.ts index 67469ac391a..3486b1841a1 100644 --- a/packages/core/src/amazonq/webview/ui/apps/cwChatConnector.ts +++ b/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, diff --git a/packages/core/src/codewhispererChat/clients/chat/v0/chat.ts b/packages/core/src/codewhispererChat/clients/chat/v0/chat.ts index 87ff36dee6b..738923ecddc 100644 --- a/packages/core/src/codewhispererChat/clients/chat/v0/chat.ts +++ b/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 = 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 { + return this._fsWriteBackups + } + + public setFsWriteBackup(toolUseId: string, backup: FsWriteBackup) { + this._fsWriteBackups.set(toolUseId, backup) + } + public tokenSource!: vscode.CancellationTokenSource constructor() { diff --git a/packages/core/src/codewhispererChat/controllers/chat/controller.ts b/packages/core/src/codewhispererChat/controllers/chat/controller.ts index 32ab4968aed..e3ecd20f3d3 100644 --- a/packages/core/src/codewhispererChat/controllers/chat/controller.ts +++ b/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', diff --git a/packages/core/src/codewhispererChat/controllers/chat/messenger/messenger.ts b/packages/core/src/codewhispererChat/controllers/chat/messenger/messenger.ts index 32a82c3d249..ac9c85a3d94 100644 --- a/packages/core/src/codewhispererChat/controllers/chat/messenger/messenger.ts +++ b/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 ) diff --git a/packages/core/src/codewhispererChat/tools/fsWrite.ts b/packages/core/src/codewhispererChat/tools/fsWrite.ts index 847d60bc331..dff68d4f425 100644 --- a/packages/core/src/codewhispererChat/tools/fsWrite.ts +++ b/packages/core/src/codewhispererChat/tools/fsWrite.ts @@ -9,8 +9,6 @@ import vscode from 'vscode' import { fs } from '../../shared/fs/fs' import { Writable } from 'stream' import { Change, diffLines } from 'diff' -import { getDiffMarkdown } from '../../shared/utilities/diffUtils' -import { getLanguageForFilePath } from '../../shared/utilities/textDocumentUtilities' interface BaseParams { path: string @@ -41,6 +39,12 @@ export interface AppendParams extends BaseParams { export type FsWriteParams = CreateParams | StrReplaceParams | InsertParams | AppendParams +export interface FsWriteBackup { + filePath: string + content: string + isNew: boolean +} + export class FsWrite { private readonly logger = getLogger('fsWrite') @@ -73,24 +77,14 @@ export class FsWrite { } public async queueDescription(updates: Writable): Promise { - const sanitizedPath = sanitizePath(this.params.path) - const changes = await this.getDiffChanges() - const language = await getLanguageForFilePath(sanitizedPath) - - const diff = getDiffMarkdown(changes, language) - updates.write(diff) + // Write an empty string because FsWrite should only show a chat message with header + updates.write(' ') updates.end() } public async getDiffChanges(): Promise { - const sanitizedPath = sanitizePath(this.params.path) let newContent - let oldContent - try { - oldContent = await fs.readFileText(sanitizedPath) - } catch (err) { - oldContent = '' - } + const { filePath: sanitizedPath, content: oldContent } = await this.getBackup() switch (this.params.command) { case 'create': newContent = this.getCreateCommandText(this.params) @@ -108,6 +102,20 @@ export class FsWrite { return diffLines(oldContent, newContent) } + public async getBackup(): Promise { + 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: diff --git a/packages/core/src/shared/utilities/diffUtils.ts b/packages/core/src/shared/utilities/diffUtils.ts index 96b0a766133..65957bbe1f2 100644 --- a/packages/core/src/shared/utilities/diffUtils.ts +++ b/packages/core/src/shared/utilities/diffUtils.ts @@ -150,6 +150,31 @@ export function getDiffCharsAndLines( } } +/** + * Calculates the number of added and deleted lines from a set of changes. + * + * @param {Change[] | undefined} changes - An array of diff Change objects containing added, removed, or unchanged content + * @returns {Object} An object containing the number of added and deleted lines + * @example + * // Calculate the number of added and deleted lines from diff changes + * const diffChanges = diffLines(originalCode, newCode); + * const lineDiffs = getDiffLinesFromChanges(diffChanges); + * // Result will be an object with added and deleted line counts + */ +export function getDiffLinesFromChanges(changes: Change[] | undefined) { + return changes?.reduce( + (acc, { count = 0, added, removed }) => { + if (added) { + acc.added += count + } else if (removed) { + acc.deleted += count + } + return acc + }, + { added: 0, deleted: 0 } + ) +} + /** * Converts diff changes into a markdown-formatted string with syntax highlighting. * This function takes an array of diff changes and formats them as a markdown code block diff --git a/packages/core/src/test/codewhispererChat/tools/toolShared.test.ts b/packages/core/src/test/codewhispererChat/tools/toolShared.test.ts index b6dec845e3e..e6209561f93 100644 --- a/packages/core/src/test/codewhispererChat/tools/toolShared.test.ts +++ b/packages/core/src/test/codewhispererChat/tools/toolShared.test.ts @@ -67,9 +67,9 @@ describe('ToolUtils', function () { assert.strictEqual(ToolUtils.requiresAcceptance(tool).requiresAcceptance, false) }) - it('returns true for FsWrite tool', function () { + it('returns false for FsWrite tool', function () { const tool: Tool = { type: ToolType.FsWrite, tool: mockFsWrite as unknown as FsWrite } - assert.strictEqual(ToolUtils.requiresAcceptance(tool).requiresAcceptance, true) + assert.strictEqual(ToolUtils.requiresAcceptance(tool).requiresAcceptance, false) }) it('delegates to the tool for ExecuteBash', function () {