Skip to content

Commit b38ae21

Browse files
joyceerhljoyceerhl
authored
feat: enable retrying and deleting requests in an editing session (microsoft#231143)
* feat: enable retrying and deleting requests in an editing session * fix: allow the user to delete or retry the first request by storing snapshots as soon as a request is sent * fix: don't advertise checkpointing for now
1 parent 61b0fe7 commit b38ae21

File tree

6 files changed

+214
-51
lines changed

6 files changed

+214
-51
lines changed

‎src/vs/workbench/contrib/chat/browser/actions/chatTitleActions.ts

+41-5
Original file line numberDiff line numberDiff line change
@@ -8,9 +8,10 @@ import { KeyCode, KeyMod } from '../../../../../base/common/keyCodes.js';
88
import { marked } from '../../../../../base/common/marked/marked.js';
99
import { ServicesAccessor } from '../../../../../editor/browser/editorExtensions.js';
1010
import { IBulkEditService } from '../../../../../editor/browser/services/bulkEditService.js';
11-
import { localize2 } from '../../../../../nls.js';
11+
import { localize, localize2 } from '../../../../../nls.js';
1212
import { Action2, MenuId, registerAction2 } from '../../../../../platform/actions/common/actions.js';
1313
import { ContextKeyExpr } from '../../../../../platform/contextkey/common/contextkey.js';
14+
import { IDialogService } from '../../../../../platform/dialogs/common/dialogs.js';
1415
import { KeybindingWeight } from '../../../../../platform/keybinding/common/keybindingsRegistry.js';
1516
import { IEditorService } from '../../../../services/editor/common/editorService.js';
1617
import { ResourceNotebookCellEdit } from '../../../bulkEdit/browser/bulkCellEdits.js';
@@ -20,9 +21,10 @@ import { CellEditType, CellKind, NOTEBOOK_EDITOR_ID } from '../../../notebook/co
2021
import { NOTEBOOK_IS_ACTIVE_EDITOR } from '../../../notebook/common/notebookContextKeys.js';
2122
import { ChatAgentLocation } from '../../common/chatAgents.js';
2223
import { CONTEXT_CHAT_LOCATION, CONTEXT_CHAT_RESPONSE_SUPPORT_ISSUE_REPORTING, CONTEXT_IN_CHAT_INPUT, CONTEXT_IN_CHAT_SESSION, CONTEXT_ITEM_ID, CONTEXT_LAST_ITEM_ID, CONTEXT_REQUEST, CONTEXT_RESPONSE, CONTEXT_RESPONSE_ERROR, CONTEXT_RESPONSE_FILTERED, CONTEXT_RESPONSE_VOTE } from '../../common/chatContextKeys.js';
24+
import { IChatEditingService } from '../../common/chatEditingService.js';
2325
import { ChatAgentVoteDirection, ChatAgentVoteDownReason, IChatService } from '../../common/chatService.js';
2426
import { isRequestVM, isResponseVM } from '../../common/chatViewModel.js';
25-
import { IChatWidgetService } from '../chat.js';
27+
import { ChatTreeItem, IChatWidgetService } from '../chat.js';
2628
import { CHAT_CATEGORY } from './chatActions.js';
2729

2830
export const MarkUnhelpfulActionId = 'workbench.action.chat.markUnhelpful';
@@ -188,14 +190,38 @@ export function registerChatTitleActions() {
188190
});
189191
}
190192

191-
run(accessor: ServicesAccessor, ...args: any[]) {
193+
async run(accessor: ServicesAccessor, ...args: any[]) {
192194
const item = args[0];
193195
if (!isResponseVM(item)) {
194196
return;
195197
}
196198

197199
const chatService = accessor.get(IChatService);
198-
const request = chatService.getSession(item.sessionId)?.getRequests().find(candidate => candidate.id === item.requestId);
200+
const chatEditingService = accessor.get(IChatEditingService);
201+
const chatModel = chatService.getSession(item.sessionId);
202+
const chatRequests = chatModel?.getRequests();
203+
if (!chatRequests) {
204+
return;
205+
}
206+
const itemIndex = chatRequests?.findIndex(request => request.id === item.requestId);
207+
if (chatModel?.initialLocation === ChatAgentLocation.EditingSession) {
208+
const dialogService = accessor.get(IDialogService);
209+
const confirmation = await dialogService.confirm({
210+
title: localize('chat.removeLast.confirmation.title', "Do you want to retry your last edit?"),
211+
message: localize('chat.remove.confirmation.message', "This will also undo any edits made to your working set from this request."),
212+
primaryButton: localize('chat.remove.confirmation.primaryButton', "Yes"),
213+
type: 'info'
214+
});
215+
if (!confirmation) {
216+
return;
217+
}
218+
// Reset the snapshot
219+
const snapshotRequest = chatRequests[itemIndex];
220+
if (snapshotRequest) {
221+
await chatEditingService.restoreSnapshot(snapshotRequest.id);
222+
}
223+
}
224+
const request = chatModel?.getRequests().find(candidate => candidate.id === item.requestId);
199225
chatService.resendRequest(request!);
200226
}
201227
});
@@ -300,13 +326,23 @@ export function registerChatTitleActions() {
300326
}
301327

302328
run(accessor: ServicesAccessor, ...args: any[]) {
303-
let item = args[0];
329+
let item: ChatTreeItem | undefined = args[0];
304330
if (!isRequestVM(item)) {
305331
const chatWidgetService = accessor.get(IChatWidgetService);
306332
const widget = chatWidgetService.lastFocusedWidget;
307333
item = widget?.getFocus();
308334
}
309335

336+
if (!item) {
337+
return;
338+
}
339+
340+
const chatService = accessor.get(IChatService);
341+
const chatModel = chatService.getSession(item.sessionId);
342+
if (chatModel?.initialLocation !== ChatAgentLocation.EditingSession) {
343+
return;
344+
}
345+
310346
const requestId = isRequestVM(item) ? item.id :
311347
isResponseVM(item) ? item.requestId : undefined;
312348

‎src/vs/workbench/contrib/chat/browser/chatContentParts/chatMarkdownContentPart.ts

+19-8
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,7 @@ import { IEditorService } from '../../../../services/editor/common/editorService
3030
import { IMarkdownVulnerability } from '../../common/annotations.js';
3131
import { IChatEditingService } from '../../common/chatEditingService.js';
3232
import { IChatProgressRenderableResponseContent } from '../../common/chatModel.js';
33+
import { IChatService } from '../../common/chatService.js';
3334
import { isRequestVM, isResponseVM } from '../../common/chatViewModel.js';
3435
import { CodeBlockModelCollection } from '../../common/codeBlockModelCollection.js';
3536
import { IChatCodeBlockInfo, IChatListItemRendererOptions } from '../chat.js';
@@ -136,7 +137,7 @@ export class ChatMarkdownContentPart extends Disposable implements IChatContentP
136137
return ref.object.element;
137138
} else {
138139
const requestId = isRequestVM(element) ? element.id : element.requestId;
139-
const ref = this.renderCodeBlockPill(requestId, codeBlockInfo.codemapperUri, isCodeBlockComplete);
140+
const ref = this.renderCodeBlockPill(element.sessionId, requestId, codeBlockInfo.codemapperUri, isCodeBlockComplete);
140141
if (isResponseVM(codeBlockInfo.element)) {
141142
// TODO@joyceerhl: remove this code when we change the codeblockUri API to make the URI available synchronously
142143
this.codeBlockModelCollection.update(codeBlockInfo.element.sessionId, codeBlockInfo.element, codeBlockInfo.codeBlockIndex, { text, languageId: codeBlockInfo.languageId }).then((e) => {
@@ -177,8 +178,8 @@ export class ChatMarkdownContentPart extends Disposable implements IChatContentP
177178
this.domNode = result.element;
178179
}
179180

180-
private renderCodeBlockPill(requestId: string, codemapperUri: URI | undefined, isCodeBlockComplete?: boolean): IDisposableReference<CollapsedCodeBlock> {
181-
const codeBlock = this.instantiationService.createInstance(CollapsedCodeBlock, requestId);
181+
private renderCodeBlockPill(sessionId: string, requestId: string, codemapperUri: URI | undefined, isCodeBlockComplete?: boolean): IDisposableReference<CollapsedCodeBlock> {
182+
const codeBlock = this.instantiationService.createInstance(CollapsedCodeBlock, sessionId, requestId);
182183
if (codemapperUri) {
183184
codeBlock.render(codemapperUri, !isCodeBlockComplete);
184185
}
@@ -275,10 +276,12 @@ class CollapsedCodeBlock extends Disposable {
275276
private isStreaming: boolean | undefined;
276277

277278
constructor(
279+
private readonly sessionId: string,
278280
private readonly requestId: string,
279281
@ILabelService private readonly labelService: ILabelService,
280282
@IEditorService private readonly editorService: IEditorService,
281283
@IModelService private readonly modelService: IModelService,
284+
@IChatService private readonly chatService: IChatService,
282285
@ILanguageService private readonly languageService: ILanguageService,
283286
@IChatEditingService private readonly chatEditingService: IChatEditingService,
284287
) {
@@ -287,11 +290,19 @@ class CollapsedCodeBlock extends Disposable {
287290
this.element.classList.add('show-file-icons');
288291
this._register(dom.addDisposableListener(this.element, 'click', async () => {
289292
if (this.uri) {
290-
const snapshot = this.chatEditingService.getSnapshotUri(this.requestId, this.uri);
291-
if (snapshot) {
292-
const editor = await this.editorService.openEditor({ resource: snapshot, label: localize('chatEditing.snapshot', '{0} (Working Set History)', basename(this.uri)), options: { transient: true, activation: EditorActivation.ACTIVATE } });
293-
if (isCodeEditor(editor)) {
294-
editor.updateOptions({ readOnly: true });
293+
const chatModel = this.chatService.getSession(this.sessionId);
294+
const requests = chatModel?.getRequests();
295+
if (!requests) {
296+
return;
297+
}
298+
const snapshotRequestId = requests?.find((v, i) => i > 0 && requests[i - 1]?.id === this.requestId)?.id;
299+
if (snapshotRequestId) {
300+
const snapshot = this.chatEditingService.getSnapshotUri(snapshotRequestId, this.uri);
301+
if (snapshot) {
302+
const editor = await this.editorService.openEditor({ resource: snapshot, label: localize('chatEditing.snapshot', '{0} (Working Set History)', basename(this.uri)), options: { transient: true, activation: EditorActivation.ACTIVATE } });
303+
if (isCodeEditor(editor)) {
304+
editor.updateOptions({ readOnly: true });
305+
}
295306
}
296307
} else {
297308
this.editorService.openEditor({ resource: this.uri });

‎src/vs/workbench/contrib/chat/browser/chatEditingActions.ts

+108-10
Original file line numberDiff line numberDiff line change
@@ -4,22 +4,26 @@
44
*--------------------------------------------------------------------------------------------*/
55

66
import { Codicon } from '../../../../base/common/codicons.js';
7+
import { KeyCode, KeyMod } from '../../../../base/common/keyCodes.js';
78
import { ResourceSet } from '../../../../base/common/map.js';
89
import { URI } from '../../../../base/common/uri.js';
910
import { ServicesAccessor } from '../../../../editor/browser/editorExtensions.js';
1011
import { localize, localize2 } from '../../../../nls.js';
1112
import { Action2, MenuId, registerAction2 } from '../../../../platform/actions/common/actions.js';
1213
import { ContextKeyExpr } from '../../../../platform/contextkey/common/contextkey.js';
14+
import { IDialogService } from '../../../../platform/dialogs/common/dialogs.js';
1315
import { EditorActivation } from '../../../../platform/editor/common/editor.js';
16+
import { KeybindingWeight } from '../../../../platform/keybinding/common/keybindingsRegistry.js';
1417
import { IListService } from '../../../../platform/list/browser/listService.js';
1518
import { GroupsOrder, IEditorGroupsService } from '../../../services/editor/common/editorGroupsService.js';
1619
import { IEditorService } from '../../../services/editor/common/editorService.js';
1720
import { ChatAgentLocation } from '../common/chatAgents.js';
18-
import { CONTEXT_CHAT_LOCATION, CONTEXT_CHAT_REQUEST_IN_PROGRESS, CONTEXT_ITEM_ID, CONTEXT_LAST_ITEM_ID, CONTEXT_RESPONSE } from '../common/chatContextKeys.js';
21+
import { CONTEXT_CHAT_LOCATION, CONTEXT_CHAT_REQUEST_IN_PROGRESS, CONTEXT_IN_CHAT_INPUT, CONTEXT_IN_CHAT_SESSION, CONTEXT_REQUEST, CONTEXT_RESPONSE } from '../common/chatContextKeys.js';
1922
import { applyingChatEditsContextKey, CHAT_EDITING_MULTI_DIFF_SOURCE_RESOLVER_SCHEME, chatEditingResourceContextKey, chatEditingWidgetFileStateContextKey, decidedChatEditingResourceContextKey, IChatEditingService, IChatEditingSession, inChatEditingSessionContextKey, isChatRequestCheckpointed, WorkingSetEntryState } from '../common/chatEditingService.js';
20-
import { isResponseVM } from '../common/chatViewModel.js';
23+
import { IChatService } from '../common/chatService.js';
24+
import { isRequestVM, isResponseVM } from '../common/chatViewModel.js';
2125
import { CHAT_CATEGORY } from './actions/chatActions.js';
22-
import { IChatWidget, IChatWidgetService } from './chat.js';
26+
import { ChatTreeItem, IChatWidget, IChatWidgetService } from './chat.js';
2327

2428
abstract class WorkingSetAction extends Action2 {
2529
run(accessor: ServicesAccessor, ...args: any[]) {
@@ -295,11 +299,8 @@ registerAction2(class RestoreWorkingSetAction extends Action2 {
295299
id: MenuId.ChatMessageFooter,
296300
group: 'navigation',
297301
order: 1000,
298-
when: ContextKeyExpr.and(
299-
CONTEXT_CHAT_LOCATION.isEqualTo(ChatAgentLocation.EditingSession),
300-
CONTEXT_RESPONSE,
301-
ContextKeyExpr.notIn(CONTEXT_ITEM_ID.key, CONTEXT_LAST_ITEM_ID.key)
302-
)
302+
when: ContextKeyExpr.false()
303+
// when: ContextKeyExpr.and(CONTEXT_CHAT_LOCATION.isEqualTo(ChatAgentLocation.EditingSession), CONTEXT_RESPONSE, ContextKeyExpr.notIn(CONTEXT_ITEM_ID.key, CONTEXT_LAST_ITEM_ID.key)
303304
}
304305
});
305306
}
@@ -312,13 +313,110 @@ registerAction2(class RestoreWorkingSetAction extends Action2 {
312313
}
313314

314315
const { session, requestId } = item.model;
315-
if (requestId === session.checkpoint?.id) {
316+
const shouldUnsetCheckpoint = requestId === session.checkpoint?.id;
317+
if (shouldUnsetCheckpoint) {
316318
// Unset the existing checkpoint
317319
session.setCheckpoint(undefined);
318320
} else {
319321
session.setCheckpoint(requestId);
320322
}
321323

322-
chatEditingService.restoreSnapshot(requestId);
324+
// The next request is associated with the working set snapshot representing
325+
// the 'good state' from this checkpointed request
326+
const chatService = accessor.get(IChatService);
327+
const chatModel = chatService.getSession(item.sessionId);
328+
const chatRequests = chatModel?.getRequests();
329+
const snapshot = chatRequests?.find((v, i) => i > 0 && chatRequests[i - 1]?.id === requestId);
330+
if (!shouldUnsetCheckpoint && snapshot !== undefined) {
331+
chatEditingService.restoreSnapshot(snapshot.id);
332+
} else if (shouldUnsetCheckpoint) {
333+
chatEditingService.restoreSnapshot(undefined);
334+
}
335+
}
336+
});
337+
338+
registerAction2(class RemoveAction extends Action2 {
339+
constructor() {
340+
super({
341+
id: 'workbench.action.chat.undoEdits',
342+
title: localize2('chat.undoEdits.label', "Undo Edits"),
343+
f1: false,
344+
category: CHAT_CATEGORY,
345+
icon: Codicon.trashcan,
346+
keybinding: {
347+
primary: KeyCode.Delete,
348+
mac: {
349+
primary: KeyMod.CtrlCmd | KeyCode.Backspace,
350+
},
351+
when: ContextKeyExpr.and(CONTEXT_CHAT_LOCATION.isEqualTo(ChatAgentLocation.EditingSession), CONTEXT_IN_CHAT_SESSION, CONTEXT_IN_CHAT_INPUT.negate()),
352+
weight: KeybindingWeight.WorkbenchContrib,
353+
},
354+
menu: [
355+
{
356+
id: MenuId.ChatMessageFooter,
357+
group: 'navigation',
358+
order: 4,
359+
when: ContextKeyExpr.and(CONTEXT_CHAT_LOCATION.isEqualTo(ChatAgentLocation.EditingSession), CONTEXT_RESPONSE)
360+
},
361+
{
362+
id: MenuId.ChatMessageTitle,
363+
group: 'navigation',
364+
order: 2,
365+
when: ContextKeyExpr.and(CONTEXT_CHAT_LOCATION.isEqualTo(ChatAgentLocation.EditingSession), CONTEXT_REQUEST)
366+
}
367+
]
368+
});
369+
}
370+
371+
async run(accessor: ServicesAccessor, ...args: any[]) {
372+
let item: ChatTreeItem | undefined = args[0];
373+
if (!isResponseVM(item)) {
374+
const chatWidgetService = accessor.get(IChatWidgetService);
375+
const widget = chatWidgetService.lastFocusedWidget;
376+
item = widget?.getFocus();
377+
}
378+
379+
if (!item) {
380+
return;
381+
}
382+
383+
const chatService = accessor.get(IChatService);
384+
const chatModel = chatService.getSession(item.sessionId);
385+
if (chatModel?.initialLocation !== ChatAgentLocation.EditingSession) {
386+
return;
387+
}
388+
389+
const requestId = isRequestVM(item) ? item.id :
390+
isResponseVM(item) ? item.requestId : undefined;
391+
392+
if (requestId) {
393+
const dialogService = accessor.get(IDialogService);
394+
const chatEditingService = accessor.get(IChatEditingService);
395+
const chatRequests = chatModel.getRequests();
396+
const itemIndex = chatRequests.findIndex(request => request.id === requestId);
397+
const editsToUndo = chatRequests.length - itemIndex;
398+
399+
const confirmation = await dialogService.confirm({
400+
title: editsToUndo === 1
401+
? localize('chat.removeLast.confirmation.title', "Do you want to undo your last edit?")
402+
: localize('chat.remove.confirmation.title', "Do you want to undo {0} edits?", editsToUndo),
403+
message: editsToUndo === 1
404+
? localize('chat.removeLast.confirmation.message', "This will remove your last request and undo the edits it made to your working set.")
405+
: localize('chat.remove.confirmation.message', "This will remove all subsequent requests and undo the edits they made to your working set."),
406+
primaryButton: localize('chat.remove.confirmation.primaryButton', "Yes"),
407+
type: 'info'
408+
});
409+
410+
if (confirmation.confirmed) {
411+
// Restore the snapshot to what it was before the request(s) that we deleted
412+
const snapshotRequestId = chatRequests[itemIndex].id;
413+
await chatEditingService.restoreSnapshot(snapshotRequestId);
414+
415+
// Remove the request and all that come after it
416+
for (const request of chatRequests.slice(itemIndex)) {
417+
await chatService.removeRequest(item.sessionId, request.id);
418+
}
419+
}
420+
}
323421
}
324422
});

0 commit comments

Comments
 (0)