Skip to content

Commit

Permalink
feat: enable retrying and deleting requests in an editing session (#2…
Browse files Browse the repository at this point in the history 
…31143)

* 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
  • Loading branch information
@joyceerhl
joyceerhl authored Oct 11, 2024
1 parent 61b0fe7 commit b38ae21
Show file tree
Hide file tree
Showing 6 changed files with 214 additions and 51 deletions.
46 changes: 41 additions & 5 deletions src/vs/workbench/contrib/chat/browser/actions/chatTitleActions.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -8,9 +8,10 @@ import { KeyCode, KeyMod } from '../../../../../base/common/keyCodes.js';
import { marked } from '../../../../../base/common/marked/marked.js';
import { ServicesAccessor } from '../../../../../editor/browser/editorExtensions.js';
import { IBulkEditService } from '../../../../../editor/browser/services/bulkEditService.js';
import { localize2 } from '../../../../../nls.js';
import { localize, localize2 } from '../../../../../nls.js';
import { Action2, MenuId, registerAction2 } from '../../../../../platform/actions/common/actions.js';
import { ContextKeyExpr } from '../../../../../platform/contextkey/common/contextkey.js';
import { IDialogService } from '../../../../../platform/dialogs/common/dialogs.js';
import { KeybindingWeight } from '../../../../../platform/keybinding/common/keybindingsRegistry.js';
import { IEditorService } from '../../../../services/editor/common/editorService.js';
import { ResourceNotebookCellEdit } from '../../../bulkEdit/browser/bulkCellEdits.js';
Expand All @@ -20,9 +21,10 @@ import { CellEditType, CellKind, NOTEBOOK_EDITOR_ID } from '../../../notebook/co
import { NOTEBOOK_IS_ACTIVE_EDITOR } from '../../../notebook/common/notebookContextKeys.js';
import { ChatAgentLocation } from '../../common/chatAgents.js';
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';
import { IChatEditingService } from '../../common/chatEditingService.js';
import { ChatAgentVoteDirection, ChatAgentVoteDownReason, IChatService } from '../../common/chatService.js';
import { isRequestVM, isResponseVM } from '../../common/chatViewModel.js';
import { IChatWidgetService } from '../chat.js';
import { ChatTreeItem, IChatWidgetService } from '../chat.js';
import { CHAT_CATEGORY } from './chatActions.js';

export const MarkUnhelpfulActionId = 'workbench.action.chat.markUnhelpful';
Expand Down Expand Up @@ -188,14 +190,38 @@ export function registerChatTitleActions() {
});
}

run(accessor: ServicesAccessor, ...args: any[]) {
async run(accessor: ServicesAccessor, ...args: any[]) {
const item = args[0];
if (!isResponseVM(item)) {
return;
}

const chatService = accessor.get(IChatService);
const request = chatService.getSession(item.sessionId)?.getRequests().find(candidate => candidate.id === item.requestId);
const chatEditingService = accessor.get(IChatEditingService);
const chatModel = chatService.getSession(item.sessionId);
const chatRequests = chatModel?.getRequests();
if (!chatRequests) {
return;
}
const itemIndex = chatRequests?.findIndex(request => request.id === item.requestId);
if (chatModel?.initialLocation === ChatAgentLocation.EditingSession) {
const dialogService = accessor.get(IDialogService);
const confirmation = await dialogService.confirm({
title: localize('chat.removeLast.confirmation.title', "Do you want to retry your last edit?"),
message: localize('chat.remove.confirmation.message', "This will also undo any edits made to your working set from this request."),
primaryButton: localize('chat.remove.confirmation.primaryButton', "Yes"),
type: 'info'
});
if (!confirmation) {
return;
}
// Reset the snapshot
const snapshotRequest = chatRequests[itemIndex];
if (snapshotRequest) {
await chatEditingService.restoreSnapshot(snapshotRequest.id);
}
}
const request = chatModel?.getRequests().find(candidate => candidate.id === item.requestId);
chatService.resendRequest(request!);
}
});
Expand Down Expand Up @@ -300,13 +326,23 @@ export function registerChatTitleActions() {
}

run(accessor: ServicesAccessor, ...args: any[]) {
let item = args[0];
let item: ChatTreeItem | undefined = args[0];
if (!isRequestVM(item)) {
const chatWidgetService = accessor.get(IChatWidgetService);
const widget = chatWidgetService.lastFocusedWidget;
item = widget?.getFocus();
}

if (!item) {
return;
}

const chatService = accessor.get(IChatService);
const chatModel = chatService.getSession(item.sessionId);
if (chatModel?.initialLocation !== ChatAgentLocation.EditingSession) {
return;
}

const requestId = isRequestVM(item) ? item.id :
isResponseVM(item) ? item.requestId : undefined;

Expand Down
27 changes: 19 additions & 8 deletions src/vs/workbench/contrib/chat/browser/chatContentParts/chatMarkdownContentPart.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,7 @@ import { IEditorService } from '../../../../services/editor/common/editorService
import { IMarkdownVulnerability } from '../../common/annotations.js';
import { IChatEditingService } from '../../common/chatEditingService.js';
import { IChatProgressRenderableResponseContent } from '../../common/chatModel.js';
import { IChatService } from '../../common/chatService.js';
import { isRequestVM, isResponseVM } from '../../common/chatViewModel.js';
import { CodeBlockModelCollection } from '../../common/codeBlockModelCollection.js';
import { IChatCodeBlockInfo, IChatListItemRendererOptions } from '../chat.js';
Expand Down Expand Up @@ -136,7 +137,7 @@ export class ChatMarkdownContentPart extends Disposable implements IChatContentP
return ref.object.element;
} else {
const requestId = isRequestVM(element) ? element.id : element.requestId;
const ref = this.renderCodeBlockPill(requestId, codeBlockInfo.codemapperUri, isCodeBlockComplete);
const ref = this.renderCodeBlockPill(element.sessionId, requestId, codeBlockInfo.codemapperUri, isCodeBlockComplete);
if (isResponseVM(codeBlockInfo.element)) {
// TODO@joyceerhl: remove this code when we change the codeblockUri API to make the URI available synchronously
this.codeBlockModelCollection.update(codeBlockInfo.element.sessionId, codeBlockInfo.element, codeBlockInfo.codeBlockIndex, { text, languageId: codeBlockInfo.languageId }).then((e) => {
Expand Down Expand Up @@ -177,8 +178,8 @@ export class ChatMarkdownContentPart extends Disposable implements IChatContentP
this.domNode = result.element;
}

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

constructor(
private readonly sessionId: string,
private readonly requestId: string,
@ILabelService private readonly labelService: ILabelService,
@IEditorService private readonly editorService: IEditorService,
@IModelService private readonly modelService: IModelService,
@IChatService private readonly chatService: IChatService,
@ILanguageService private readonly languageService: ILanguageService,
@IChatEditingService private readonly chatEditingService: IChatEditingService,
) {
Expand All @@ -287,11 +290,19 @@ class CollapsedCodeBlock extends Disposable {
this.element.classList.add('show-file-icons');
this._register(dom.addDisposableListener(this.element, 'click', async () => {
if (this.uri) {
const snapshot = this.chatEditingService.getSnapshotUri(this.requestId, this.uri);
if (snapshot) {
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 } });
if (isCodeEditor(editor)) {
editor.updateOptions({ readOnly: true });
const chatModel = this.chatService.getSession(this.sessionId);
const requests = chatModel?.getRequests();
if (!requests) {
return;
}
const snapshotRequestId = requests?.find((v, i) => i > 0 && requests[i - 1]?.id === this.requestId)?.id;
if (snapshotRequestId) {
const snapshot = this.chatEditingService.getSnapshotUri(snapshotRequestId, this.uri);
if (snapshot) {
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 } });
if (isCodeEditor(editor)) {
editor.updateOptions({ readOnly: true });
}
}
} else {
this.editorService.openEditor({ resource: this.uri });
Expand Down
118 changes: 108 additions & 10 deletions src/vs/workbench/contrib/chat/browser/chatEditingActions.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,22 +4,26 @@
*--------------------------------------------------------------------------------------------*/

import { Codicon } from '../../../../base/common/codicons.js';
import { KeyCode, KeyMod } from '../../../../base/common/keyCodes.js';
import { ResourceSet } from '../../../../base/common/map.js';
import { URI } from '../../../../base/common/uri.js';
import { ServicesAccessor } from '../../../../editor/browser/editorExtensions.js';
import { localize, localize2 } from '../../../../nls.js';
import { Action2, MenuId, registerAction2 } from '../../../../platform/actions/common/actions.js';
import { ContextKeyExpr } from '../../../../platform/contextkey/common/contextkey.js';
import { IDialogService } from '../../../../platform/dialogs/common/dialogs.js';
import { EditorActivation } from '../../../../platform/editor/common/editor.js';
import { KeybindingWeight } from '../../../../platform/keybinding/common/keybindingsRegistry.js';
import { IListService } from '../../../../platform/list/browser/listService.js';
import { GroupsOrder, IEditorGroupsService } from '../../../services/editor/common/editorGroupsService.js';
import { IEditorService } from '../../../services/editor/common/editorService.js';
import { ChatAgentLocation } from '../common/chatAgents.js';
import { CONTEXT_CHAT_LOCATION, CONTEXT_CHAT_REQUEST_IN_PROGRESS, CONTEXT_ITEM_ID, CONTEXT_LAST_ITEM_ID, CONTEXT_RESPONSE } from '../common/chatContextKeys.js';
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';
import { applyingChatEditsContextKey, CHAT_EDITING_MULTI_DIFF_SOURCE_RESOLVER_SCHEME, chatEditingResourceContextKey, chatEditingWidgetFileStateContextKey, decidedChatEditingResourceContextKey, IChatEditingService, IChatEditingSession, inChatEditingSessionContextKey, isChatRequestCheckpointed, WorkingSetEntryState } from '../common/chatEditingService.js';
import { isResponseVM } from '../common/chatViewModel.js';
import { IChatService } from '../common/chatService.js';
import { isRequestVM, isResponseVM } from '../common/chatViewModel.js';
import { CHAT_CATEGORY } from './actions/chatActions.js';
import { IChatWidget, IChatWidgetService } from './chat.js';
import { ChatTreeItem, IChatWidget, IChatWidgetService } from './chat.js';

abstract class WorkingSetAction extends Action2 {
run(accessor: ServicesAccessor, ...args: any[]) {
Expand Down Expand Up @@ -295,11 +299,8 @@ registerAction2(class RestoreWorkingSetAction extends Action2 {
id: MenuId.ChatMessageFooter,
group: 'navigation',
order: 1000,
when: ContextKeyExpr.and(
CONTEXT_CHAT_LOCATION.isEqualTo(ChatAgentLocation.EditingSession),
CONTEXT_RESPONSE,
ContextKeyExpr.notIn(CONTEXT_ITEM_ID.key, CONTEXT_LAST_ITEM_ID.key)
)
when: ContextKeyExpr.false()
// when: ContextKeyExpr.and(CONTEXT_CHAT_LOCATION.isEqualTo(ChatAgentLocation.EditingSession), CONTEXT_RESPONSE, ContextKeyExpr.notIn(CONTEXT_ITEM_ID.key, CONTEXT_LAST_ITEM_ID.key)
}
});
}
Expand All @@ -312,13 +313,110 @@ registerAction2(class RestoreWorkingSetAction extends Action2 {
}

const { session, requestId } = item.model;
if (requestId === session.checkpoint?.id) {
const shouldUnsetCheckpoint = requestId === session.checkpoint?.id;
if (shouldUnsetCheckpoint) {
// Unset the existing checkpoint
session.setCheckpoint(undefined);
} else {
session.setCheckpoint(requestId);
}

chatEditingService.restoreSnapshot(requestId);
// The next request is associated with the working set snapshot representing
// the 'good state' from this checkpointed request
const chatService = accessor.get(IChatService);
const chatModel = chatService.getSession(item.sessionId);
const chatRequests = chatModel?.getRequests();
const snapshot = chatRequests?.find((v, i) => i > 0 && chatRequests[i - 1]?.id === requestId);
if (!shouldUnsetCheckpoint && snapshot !== undefined) {
chatEditingService.restoreSnapshot(snapshot.id);
} else if (shouldUnsetCheckpoint) {
chatEditingService.restoreSnapshot(undefined);
}
}
});

registerAction2(class RemoveAction extends Action2 {
constructor() {
super({
id: 'workbench.action.chat.undoEdits',
title: localize2('chat.undoEdits.label', "Undo Edits"),
f1: false,
category: CHAT_CATEGORY,
icon: Codicon.trashcan,
keybinding: {
primary: KeyCode.Delete,
mac: {
primary: KeyMod.CtrlCmd | KeyCode.Backspace,
},
when: ContextKeyExpr.and(CONTEXT_CHAT_LOCATION.isEqualTo(ChatAgentLocation.EditingSession), CONTEXT_IN_CHAT_SESSION, CONTEXT_IN_CHAT_INPUT.negate()),
weight: KeybindingWeight.WorkbenchContrib,
},
menu: [
{
id: MenuId.ChatMessageFooter,
group: 'navigation',
order: 4,
when: ContextKeyExpr.and(CONTEXT_CHAT_LOCATION.isEqualTo(ChatAgentLocation.EditingSession), CONTEXT_RESPONSE)
},
{
id: MenuId.ChatMessageTitle,
group: 'navigation',
order: 2,
when: ContextKeyExpr.and(CONTEXT_CHAT_LOCATION.isEqualTo(ChatAgentLocation.EditingSession), CONTEXT_REQUEST)
}
]
});
}

async run(accessor: ServicesAccessor, ...args: any[]) {
let item: ChatTreeItem | undefined = args[0];
if (!isResponseVM(item)) {
const chatWidgetService = accessor.get(IChatWidgetService);
const widget = chatWidgetService.lastFocusedWidget;
item = widget?.getFocus();
}

if (!item) {
return;
}

const chatService = accessor.get(IChatService);
const chatModel = chatService.getSession(item.sessionId);
if (chatModel?.initialLocation !== ChatAgentLocation.EditingSession) {
return;
}

const requestId = isRequestVM(item) ? item.id :
isResponseVM(item) ? item.requestId : undefined;

if (requestId) {
const dialogService = accessor.get(IDialogService);
const chatEditingService = accessor.get(IChatEditingService);
const chatRequests = chatModel.getRequests();
const itemIndex = chatRequests.findIndex(request => request.id === requestId);
const editsToUndo = chatRequests.length - itemIndex;

const confirmation = await dialogService.confirm({
title: editsToUndo === 1
? localize('chat.removeLast.confirmation.title', "Do you want to undo your last edit?")
: localize('chat.remove.confirmation.title', "Do you want to undo {0} edits?", editsToUndo),
message: editsToUndo === 1
? localize('chat.removeLast.confirmation.message', "This will remove your last request and undo the edits it made to your working set.")
: localize('chat.remove.confirmation.message', "This will remove all subsequent requests and undo the edits they made to your working set."),
primaryButton: localize('chat.remove.confirmation.primaryButton', "Yes"),
type: 'info'
});

if (confirmation.confirmed) {
// Restore the snapshot to what it was before the request(s) that we deleted
const snapshotRequestId = chatRequests[itemIndex].id;
await chatEditingService.restoreSnapshot(snapshotRequestId);

// Remove the request and all that come after it
for (const request of chatRequests.slice(itemIndex)) {
await chatService.removeRequest(item.sessionId, request.id);
}
}
}
}
});
Loading

0 comments on commit b38ae21

Please sign in to comment.