Skip to content

Commit cb3e6a5

Browse files
ctlai95ctlai95
committed
feat(chat): undo all file changes
1 parent 3b71170 commit cb3e6a5

File tree

5 files changed

+144
-28
lines changed

5 files changed

+144
-28
lines changed

‎packages/core/package.nls.json‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -457,6 +457,7 @@
457457
"AWS.amazonq.opensettings:": "Open settings",
458458
"AWS.amazonq.executeBash.run": "Run",
459459
"AWS.amazonq.executeBash.reject": "Reject",
460+
"AWS.amazonq.fsWrite.undoAll": "Undo all changes",
460461
"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.",
461462
"AWS.amazonq.chat.directive.pairProgrammingModeOff": "You turned off **pair programming**. Q will not include code diffs or run commands in the chat.",
462463
"AWS.amazonq.chat.directive.permission.readAndList": "I need permission to read files and list directories outside the workspace.",

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

Lines changed: 38 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -46,6 +46,8 @@ export class ChatSession {
4646
private _fsWriteBackups: Map<string, FsWriteBackup> = new Map()
4747
private _agenticLoopInProgress: boolean = false
4848
private _messageOperations: Map<string, FileOperation> = new Map()
49+
private _fsWriteGroupsForUndoAll: Map<string, Set<string>> = new Map()
50+
private _currentFsWriteIdForUndoAll: string | undefined
4951

5052
/**
5153
* True if messages from local history have been sent to session.
@@ -221,4 +223,40 @@ export class ChatSession {
221223
public getOperationTypeByMessageId(messageId: string): OperationType | undefined {
222224
return this._messageOperations.get(messageId)?.type
223225
}
226+
227+
/**
228+
* Gets the fsWrite groups for undo all operations
229+
* @returns Map where key is the first fsWriteId in a group and value is a Set of all fsWriteIds in that group
230+
*/
231+
public get fsWriteGroupsForUndoAll(): Map<string, Set<string>> {
232+
return this._fsWriteGroupsForUndoAll
233+
}
234+
235+
/**
236+
* Adds a single fsWriteId to a group for undo all operations
237+
* @param groupId The first fsWriteId in the group (used as key)
238+
* @param fsWriteId A single fsWriteId to add to the group
239+
*/
240+
public addToFsWriteGroupForUndoAll(groupId: string, fsWriteId: string): void {
241+
if (!this._fsWriteGroupsForUndoAll.has(groupId)) {
242+
this._fsWriteGroupsForUndoAll.set(groupId, new Set<string>())
243+
}
244+
this._fsWriteGroupsForUndoAll.get(groupId)?.add(fsWriteId)
245+
}
246+
247+
/**
248+
* Gets the current fsWriteId for undo all operations
249+
* @returns The first fsWriteId in the current undo all group, or undefined if not set
250+
*/
251+
public get currentFsWriteIdForUndoAll(): string | undefined {
252+
return this._currentFsWriteIdForUndoAll
253+
}
254+
255+
/**
256+
* Sets the current fsWriteId for undo all operations
257+
* @param fsWriteId The first fsWriteId in the current undo all group
258+
*/
259+
public setCurrentFsWriteIdForUndoAll(fsWriteId: string | undefined): void {
260+
this._currentFsWriteIdForUndoAll = fsWriteId
261+
}
224262
}

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

Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -778,6 +778,18 @@ export class ChatController {
778778
if (tool.type === ToolType.FsWrite && toolUse.toolUseId) {
779779
const backup = await tool.tool.getBackup()
780780
session.setFsWriteBackup(toolUse.toolUseId, backup)
781+
782+
if (session.currentFsWriteIdForUndoAll === undefined) {
783+
session.setCurrentFsWriteIdForUndoAll(toolUse.toolUseId)
784+
session.addToFsWriteGroupForUndoAll(toolUse.toolUseId, toolUse.toolUseId)
785+
} else {
786+
session.addToFsWriteGroupForUndoAll(
787+
session.currentFsWriteIdForUndoAll,
788+
toolUse.toolUseId
789+
)
790+
}
791+
} else {
792+
session.setCurrentFsWriteIdForUndoAll(undefined)
781793
}
782794

783795
// Check again if cancelled before invoking the tool
@@ -934,11 +946,36 @@ export class ChatController {
934946
ConversationTracker.getInstance().markTriggerCompleted(message.triggerId)
935947
}
936948
break
949+
case 'undo-all':
950+
await this.undoAllFileChanges(message)
951+
break
937952
default:
938953
getLogger().warn(`Unhandled action: ${message.action.id}`)
939954
}
940955
}
941956

957+
private async undoAllFileChanges(message: CustomFormActionMessage) {
958+
const tabID = message.tabID
959+
// UndoAll button chat messageId is expected to have the /undoAll suffix
960+
const toolUseIdWithSuffix = message.action.formItemValues?.toolUseId
961+
const toolUseId = toolUseIdWithSuffix?.split('/').shift()
962+
if (!tabID || !toolUseId) {
963+
return
964+
}
965+
966+
const session = this.sessionStorage.getSession(tabID)
967+
const fsWriteIdSet = session.fsWriteGroupsForUndoAll.get(toolUseId)
968+
if (!fsWriteIdSet) {
969+
return
970+
}
971+
972+
for (const fsWriteId of [...fsWriteIdSet].reverse()) {
973+
this.messenger.sendCustomFormActionMessage(tabID, 'reject-code-diff', message.triggerId, fsWriteId)
974+
}
975+
976+
session.fsWriteGroupsForUndoAll.delete(toolUseId)
977+
}
978+
942979
private async sendCommandRejectMessage(tabID: string) {
943980
const session = this.sessionStorage.getSession(tabID)
944981
session.setAgenticLoopInProgress(false)

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

Lines changed: 64 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -71,6 +71,7 @@ import { ConversationTracker } from '../../../storages/conversationTracker'
7171
import { waitTimeout, Timeout } from '../../../../shared/utilities/timeoutUtils'
7272
import { FsReadParams } from '../../../tools/fsRead'
7373
import { ListDirectoryParams } from '../../../tools/listDirectory'
74+
import { i18n } from '../../../../shared/i18n-helper'
7475

7576
export type StaticTextResponseType = 'quick-action-help' | 'onboarding-help' | 'transform' | 'help'
7677

@@ -327,6 +328,10 @@ export class Messenger {
327328
}
328329
const tool = ToolUtils.tryFromToolUse(toolUse)
329330
if ('type' in tool) {
331+
if (tool.type !== ToolType.FsWrite) {
332+
this.showUndoAllIfRequired(session, tabID, triggerID)
333+
}
334+
330335
let explanation: string | undefined = undefined
331336
let changeList: Change[] | undefined = undefined
332337
let messageIdToUpdate: string | undefined = undefined
@@ -397,28 +402,12 @@ export class Messenger {
397402
if (this.isTriggerCancelled(triggerID)) {
398403
return
399404
}
400-
this.dispatcher.sendCustomFormActionMessage(
401-
new CustomFormActionMessage(
402-
tabID,
403-
{
404-
id: 'run-shell-command',
405-
},
406-
triggerID
407-
)
408-
)
405+
this.sendCustomFormActionMessage(tabID, 'run-shell-command', triggerID)
409406
} else {
410407
if (this.isTriggerCancelled(triggerID)) {
411408
return
412409
}
413-
this.dispatcher.sendCustomFormActionMessage(
414-
new CustomFormActionMessage(
415-
tabID,
416-
{
417-
id: 'generic-tool-execution',
418-
},
419-
triggerID
420-
)
421-
)
410+
this.sendCustomFormActionMessage(tabID, 'generic-tool-execution', triggerID)
422411
}
423412
} else {
424413
if (tool.type === ToolType.ExecuteBash) {
@@ -444,15 +433,7 @@ export class Messenger {
444433
)
445434
session.setToolUseWithError({ toolUse, error })
446435
// trigger processToolUseMessage to handle the error
447-
this.dispatcher.sendCustomFormActionMessage(
448-
new CustomFormActionMessage(
449-
tabID,
450-
{
451-
id: 'generic-tool-execution',
452-
},
453-
triggerID
454-
)
455-
)
436+
this.sendCustomFormActionMessage(tabID, 'generic-tool-execution', triggerID)
456437
}
457438
// TODO: Add a spinner component for fsWrite, previous implementation is causing lag in mynah UX.
458439
}
@@ -470,6 +451,8 @@ export class Messenger {
470451
return
471452
}
472453

454+
this.showUndoAllIfRequired(session, tabID, triggerID, true)
455+
473456
this.dispatcher.sendChatMessage(
474457
new ChatMessage(
475458
{
@@ -1140,4 +1123,58 @@ export class Messenger {
11401123
const conversationTracker = ConversationTracker.getInstance()
11411124
return conversationTracker.isTriggerCancelled(triggerId)
11421125
}
1126+
1127+
private showUndoAllIfRequired(
1128+
session: ChatSession,
1129+
tabID: string,
1130+
triggerID: string,
1131+
shouldSendInitialStream = false
1132+
) {
1133+
if (session.currentFsWriteIdForUndoAll === undefined) {
1134+
return
1135+
}
1136+
const fsWriteGroup = session.fsWriteGroupsForUndoAll.get(session.currentFsWriteIdForUndoAll)
1137+
if (!fsWriteGroup || fsWriteGroup.size <= 1) {
1138+
return
1139+
}
1140+
1141+
this.dispatcher.sendChatMessage(
1142+
new ChatMessage(
1143+
{
1144+
message: '',
1145+
messageType: 'answer',
1146+
followUps: undefined,
1147+
followUpsHeader: undefined,
1148+
relatedSuggestions: undefined,
1149+
triggerID,
1150+
// Add a suffix to avoid collision with the actual tool messageId
1151+
messageID: `${session.currentFsWriteIdForUndoAll}/undoAll`,
1152+
userIntent: undefined,
1153+
codeBlockLanguage: undefined,
1154+
contextList: undefined,
1155+
buttons: [
1156+
{
1157+
id: 'undo-all',
1158+
text: i18n('AWS.amazonq.fsWrite.undoAll'),
1159+
icon: 'revert',
1160+
position: 'outside',
1161+
status: 'clear',
1162+
keepCardAfterClick: false,
1163+
},
1164+
],
1165+
},
1166+
tabID
1167+
)
1168+
)
1169+
session.setCurrentFsWriteIdForUndoAll(undefined)
1170+
if (shouldSendInitialStream) {
1171+
this.sendInitalStream(tabID, triggerID)
1172+
}
1173+
}
1174+
1175+
public sendCustomFormActionMessage(tabID: string, actionID: string, triggerID: string, messageID?: string) {
1176+
this.dispatcher.sendCustomFormActionMessage(
1177+
new CustomFormActionMessage(tabID, { id: actionID }, triggerID, messageID)
1178+
)
1179+
}
11431180
}

‎packages/core/src/codewhispererChat/view/connector/connector.ts‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -276,6 +276,7 @@ export class CustomFormActionMessage extends UiMessage {
276276
formItemValues?: Record<string, string> | undefined
277277
}
278278
readonly triggerId: string
279+
readonly messageId?: string
279280

280281
constructor(
281282
tabID: string,
@@ -284,11 +285,13 @@ export class CustomFormActionMessage extends UiMessage {
284285
text?: string | undefined
285286
formItemValues?: Record<string, string> | undefined
286287
},
287-
triggerId: string
288+
triggerId: string,
289+
messageId?: string
288290
) {
289291
super(tabID)
290292
this.action = action
291293
this.triggerId = triggerId
294+
this.messageId = messageId
292295
}
293296
}
294297

0 commit comments

Comments
 (0)