Skip to content

feat(chat): undo all file changes #7093

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
wants to merge 1 commit into from
Closed

feat(chat): undo all file changes #7093

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions packages/core/package.nls.json
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.",
Expand Down
38 changes: 38 additions & 0 deletions packages/core/src/codewhispererChat/clients/chat/v0/chat.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -46,6 +46,8 @@ export class ChatSession {
private _fsWriteBackups: Map<string, FsWriteBackup> = new Map()
private _agenticLoopInProgress: boolean = false
private _messageOperations: Map<string, FileOperation> = new Map()
private _fsWriteGroupsForUndoAll: Map<string, Set<string>> = new Map()
private _currentFsWriteIdForUndoAll: string | undefined

/**
* True if messages from local history have been sent to session.
Expand Down Expand Up @@ -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<string, Set<string>> {
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<string>())
}
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
}
}
37 changes: 37 additions & 0 deletions packages/core/src/codewhispererChat/controllers/chat/controller.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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)
Expand Down
91 changes: 64 additions & 27 deletions packages/core/src/codewhispererChat/controllers/chat/messenger/messenger.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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'

Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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) {
Expand All @@ -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.
}
Expand All @@ -470,6 +451,8 @@ export class Messenger {
return
}

this.showUndoAllIfRequired(session, tabID, triggerID, true)

this.dispatcher.sendChatMessage(
new ChatMessage(
{
Expand Down Expand Up @@ -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)
)
}
}
5 changes: 4 additions & 1 deletion packages/core/src/codewhispererChat/view/connector/connector.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -276,6 +276,7 @@ export class CustomFormActionMessage extends UiMessage {
formItemValues?: Record<string, string> | undefined
}
readonly triggerId: string
readonly messageId?: string

constructor(
tabID: string,
Expand All @@ -284,11 +285,13 @@ export class CustomFormActionMessage extends UiMessage {
text?: string | undefined
formItemValues?: Record<string, string> | undefined
},
triggerId: string
triggerId: string,
messageId?: string
) {
super(tabID)
this.action = action
this.triggerId = triggerId
this.messageId = messageId
}
}

Expand Down