diff --git a/packages/core/package.nls.json b/packages/core/package.nls.json index ee9b63c220d..d899974c0f4 100644 --- a/packages/core/package.nls.json +++ b/packages/core/package.nls.json @@ -457,6 +457,7 @@ "AWS.amazonq.opensettings:": "Open settings", "AWS.amazonq.executeBash.run": "Run", "AWS.amazonq.executeBash.reject": "Reject", + "AWS.amazonq.fsWrite.undoAll": "Undo all changes", "AWS.amazonq.chat.directive.pairProgrammingModeOn": "You are using **pair programming**: Q can now list files, preview code diffs and allow you to run shell commands.", "AWS.amazonq.chat.directive.pairProgrammingModeOff": "You turned off **pair programming**. Q will not include code diffs or run commands in the chat.", "AWS.amazonq.chat.directive.permission.readAndList": "I need permission to read files and list directories outside the workspace.", diff --git a/packages/core/src/codewhispererChat/clients/chat/v0/chat.ts b/packages/core/src/codewhispererChat/clients/chat/v0/chat.ts index aac87312c3f..83555f44bd7 100644 --- a/packages/core/src/codewhispererChat/clients/chat/v0/chat.ts +++ b/packages/core/src/codewhispererChat/clients/chat/v0/chat.ts @@ -46,6 +46,8 @@ export class ChatSession { private _fsWriteBackups: Map = new Map() private _agenticLoopInProgress: boolean = false private _messageOperations: Map = new Map() + private _fsWriteGroupsForUndoAll: Map> = new Map() + private _currentFsWriteIdForUndoAll: string | undefined /** * True if messages from local history have been sent to session. @@ -221,4 +223,40 @@ export class ChatSession { public getOperationTypeByMessageId(messageId: string): OperationType | undefined { return this._messageOperations.get(messageId)?.type } + + /** + * Gets the fsWrite groups for undo all operations + * @returns Map where key is the first fsWriteId in a group and value is a Set of all fsWriteIds in that group + */ + public get fsWriteGroupsForUndoAll(): Map> { + return this._fsWriteGroupsForUndoAll + } + + /** + * Adds a single fsWriteId to a group for undo all operations + * @param groupId The first fsWriteId in the group (used as key) + * @param fsWriteId A single fsWriteId to add to the group + */ + public addToFsWriteGroupForUndoAll(groupId: string, fsWriteId: string): void { + if (!this._fsWriteGroupsForUndoAll.has(groupId)) { + this._fsWriteGroupsForUndoAll.set(groupId, new Set()) + } + this._fsWriteGroupsForUndoAll.get(groupId)?.add(fsWriteId) + } + + /** + * Gets the current fsWriteId for undo all operations + * @returns The first fsWriteId in the current undo all group, or undefined if not set + */ + public get currentFsWriteIdForUndoAll(): string | undefined { + return this._currentFsWriteIdForUndoAll + } + + /** + * Sets the current fsWriteId for undo all operations + * @param fsWriteId The first fsWriteId in the current undo all group + */ + public setCurrentFsWriteIdForUndoAll(fsWriteId: string | undefined): void { + this._currentFsWriteIdForUndoAll = fsWriteId + } } diff --git a/packages/core/src/codewhispererChat/controllers/chat/controller.ts b/packages/core/src/codewhispererChat/controllers/chat/controller.ts index a9711c3ab02..d261945f802 100644 --- a/packages/core/src/codewhispererChat/controllers/chat/controller.ts +++ b/packages/core/src/codewhispererChat/controllers/chat/controller.ts @@ -778,6 +778,18 @@ export class ChatController { if (tool.type === ToolType.FsWrite && toolUse.toolUseId) { const backup = await tool.tool.getBackup() session.setFsWriteBackup(toolUse.toolUseId, backup) + + if (session.currentFsWriteIdForUndoAll === undefined) { + session.setCurrentFsWriteIdForUndoAll(toolUse.toolUseId) + session.addToFsWriteGroupForUndoAll(toolUse.toolUseId, toolUse.toolUseId) + } else { + session.addToFsWriteGroupForUndoAll( + session.currentFsWriteIdForUndoAll, + toolUse.toolUseId + ) + } + } else { + session.setCurrentFsWriteIdForUndoAll(undefined) } // Check again if cancelled before invoking the tool @@ -934,11 +946,36 @@ export class ChatController { ConversationTracker.getInstance().markTriggerCompleted(message.triggerId) } break + case 'undo-all': + await this.undoAllFileChanges(message) + break default: getLogger().warn(`Unhandled action: ${message.action.id}`) } } + private async undoAllFileChanges(message: CustomFormActionMessage) { + const tabID = message.tabID + // UndoAll button chat messageId is expected to have the /undoAll suffix + const toolUseIdWithSuffix = message.action.formItemValues?.toolUseId + const toolUseId = toolUseIdWithSuffix?.split('/').shift() + if (!tabID || !toolUseId) { + return + } + + const session = this.sessionStorage.getSession(tabID) + const fsWriteIdSet = session.fsWriteGroupsForUndoAll.get(toolUseId) + if (!fsWriteIdSet) { + return + } + + for (const fsWriteId of [...fsWriteIdSet].reverse()) { + this.messenger.sendCustomFormActionMessage(tabID, 'reject-code-diff', message.triggerId, fsWriteId) + } + + session.fsWriteGroupsForUndoAll.delete(toolUseId) + } + private async sendCommandRejectMessage(tabID: string) { const session = this.sessionStorage.getSession(tabID) session.setAgenticLoopInProgress(false) diff --git a/packages/core/src/codewhispererChat/controllers/chat/messenger/messenger.ts b/packages/core/src/codewhispererChat/controllers/chat/messenger/messenger.ts index 34e2ddc093c..4e627d8aa7e 100644 --- a/packages/core/src/codewhispererChat/controllers/chat/messenger/messenger.ts +++ b/packages/core/src/codewhispererChat/controllers/chat/messenger/messenger.ts @@ -71,6 +71,7 @@ import { ConversationTracker } from '../../../storages/conversationTracker' import { waitTimeout, Timeout } from '../../../../shared/utilities/timeoutUtils' import { FsReadParams } from '../../../tools/fsRead' import { ListDirectoryParams } from '../../../tools/listDirectory' +import { i18n } from '../../../../shared/i18n-helper' export type StaticTextResponseType = 'quick-action-help' | 'onboarding-help' | 'transform' | 'help' @@ -327,6 +328,10 @@ export class Messenger { } const tool = ToolUtils.tryFromToolUse(toolUse) if ('type' in tool) { + if (tool.type !== ToolType.FsWrite) { + this.showUndoAllIfRequired(session, tabID, triggerID) + } + let explanation: string | undefined = undefined let changeList: Change[] | undefined = undefined let messageIdToUpdate: string | undefined = undefined @@ -397,28 +402,12 @@ export class Messenger { if (this.isTriggerCancelled(triggerID)) { return } - this.dispatcher.sendCustomFormActionMessage( - new CustomFormActionMessage( - tabID, - { - id: 'run-shell-command', - }, - triggerID - ) - ) + this.sendCustomFormActionMessage(tabID, 'run-shell-command', triggerID) } else { if (this.isTriggerCancelled(triggerID)) { return } - this.dispatcher.sendCustomFormActionMessage( - new CustomFormActionMessage( - tabID, - { - id: 'generic-tool-execution', - }, - triggerID - ) - ) + this.sendCustomFormActionMessage(tabID, 'generic-tool-execution', triggerID) } } else { if (tool.type === ToolType.ExecuteBash) { @@ -444,15 +433,7 @@ export class Messenger { ) session.setToolUseWithError({ toolUse, error }) // trigger processToolUseMessage to handle the error - this.dispatcher.sendCustomFormActionMessage( - new CustomFormActionMessage( - tabID, - { - id: 'generic-tool-execution', - }, - triggerID - ) - ) + this.sendCustomFormActionMessage(tabID, 'generic-tool-execution', triggerID) } // TODO: Add a spinner component for fsWrite, previous implementation is causing lag in mynah UX. } @@ -470,6 +451,8 @@ export class Messenger { return } + this.showUndoAllIfRequired(session, tabID, triggerID, true) + this.dispatcher.sendChatMessage( new ChatMessage( { @@ -1140,4 +1123,58 @@ export class Messenger { const conversationTracker = ConversationTracker.getInstance() return conversationTracker.isTriggerCancelled(triggerId) } + + private showUndoAllIfRequired( + session: ChatSession, + tabID: string, + triggerID: string, + shouldSendInitialStream = false + ) { + if (session.currentFsWriteIdForUndoAll === undefined) { + return + } + const fsWriteGroup = session.fsWriteGroupsForUndoAll.get(session.currentFsWriteIdForUndoAll) + if (!fsWriteGroup || fsWriteGroup.size <= 1) { + return + } + + this.dispatcher.sendChatMessage( + new ChatMessage( + { + message: '', + messageType: 'answer', + followUps: undefined, + followUpsHeader: undefined, + relatedSuggestions: undefined, + triggerID, + // Add a suffix to avoid collision with the actual tool messageId + messageID: `${session.currentFsWriteIdForUndoAll}/undoAll`, + userIntent: undefined, + codeBlockLanguage: undefined, + contextList: undefined, + buttons: [ + { + id: 'undo-all', + text: i18n('AWS.amazonq.fsWrite.undoAll'), + icon: 'revert', + position: 'outside', + status: 'clear', + keepCardAfterClick: false, + }, + ], + }, + tabID + ) + ) + session.setCurrentFsWriteIdForUndoAll(undefined) + if (shouldSendInitialStream) { + this.sendInitalStream(tabID, triggerID) + } + } + + public sendCustomFormActionMessage(tabID: string, actionID: string, triggerID: string, messageID?: string) { + this.dispatcher.sendCustomFormActionMessage( + new CustomFormActionMessage(tabID, { id: actionID }, triggerID, messageID) + ) + } } diff --git a/packages/core/src/codewhispererChat/view/connector/connector.ts b/packages/core/src/codewhispererChat/view/connector/connector.ts index a37c6279585..b400e56ad3a 100644 --- a/packages/core/src/codewhispererChat/view/connector/connector.ts +++ b/packages/core/src/codewhispererChat/view/connector/connector.ts @@ -276,6 +276,7 @@ export class CustomFormActionMessage extends UiMessage { formItemValues?: Record | undefined } readonly triggerId: string + readonly messageId?: string constructor( tabID: string, @@ -284,11 +285,13 @@ export class CustomFormActionMessage extends UiMessage { text?: string | undefined formItemValues?: Record | undefined }, - triggerId: string + triggerId: string, + messageId?: string ) { super(tabID) this.action = action this.triggerId = triggerId + this.messageId = messageId } }