first pass checkpoints inline in chat (microsoft#254564)

* first pass checkpoints inline in chat:

* add undo redo buttons back

* quieter version

* functionally, we are good here. lots to clean up asdfasdfasdfasdfafdasf

* nice, some cleanup poggersgit add .git add .

* better css

* cleanup

* remove extra instance service

Author: justschen

src/vs/platform/actions/common/actions.ts

@@ -247,6 +247,8 @@ export class MenuId {
 	static readonly ChatInputSymbolAttachmentContext = new MenuId('ChatInputSymbolAttachmentContext');
 	static readonly ChatInlineResourceAnchorContext = new MenuId('ChatInlineResourceAnchorContext');
 	static readonly ChatInlineSymbolAnchorContext = new MenuId('ChatInlineSymbolAnchorContext');
+	static readonly ChatMessageCheckpoint: MenuId = new MenuId('ChatMessageCheckpoint');
+	static readonly ChatMessageRestoreCheckpoint: MenuId = new MenuId('ChatMessageRestoreCheckpoint');
 	static readonly ChatEditingCodeBlockContext = new MenuId('ChatEditingCodeBlockContext');
 	static readonly ChatTitleBarMenu = new MenuId('ChatTitleBarMenu');
 	static readonly ChatAttachmentsContext = new MenuId('ChatAttachmentsContext');

src/vs/workbench/contrib/chat/browser/actions/chatClearActions.ts

@@ -17,7 +17,7 @@ import { ActiveEditorContext } from '../../../../common/contextkeys.js';
 import { ChatContextKeys } from '../../common/chatContextKeys.js';
 import { IChatEditingSession } from '../../common/chatEditingService.js';
 import { ChatModeKind } from '../../common/constants.js';
-import { ChatViewId, IChatWidget } from '../chat.js';
+import { ChatViewId, IChatWidget, IChatWidgetService } from '../chat.js';
 import { EditingSessionAction } from '../chatEditing/chatEditingActions.js';
 import { ChatEditorInput } from '../chatEditorInput.js';
 import { ACTION_ID_NEW_CHAT, ACTION_ID_NEW_EDIT_SESSION, CHAT_CATEGORY, handleCurrentEditingSession } from './chatActions.js';
@@ -148,7 +148,8 @@ export function registerNewChatActions() {
 					id: MenuId.ViewTitle,
 					when: ContextKeyExpr.equals('view', ChatViewId),
 					group: 'navigation',
-					order: -3
+					order: -3,
+					isHiddenByDefault: true
 				}]
 			});
 		}
@@ -167,17 +168,54 @@ export function registerNewChatActions() {
 				icon: Codicon.redo,
 				precondition: ContextKeyExpr.and(ChatContextKeys.chatEditingCanRedo, ChatContextKeys.enabled, ChatContextKeys.editingParticipantRegistered),
 				f1: true,
+				menu: [
+					{
+						id: MenuId.ViewTitle,
+						when: ContextKeyExpr.equals('view', ChatViewId),
+						group: 'navigation',
+						order: -2,
+						isHiddenByDefault: true
+					}
+				]
+			});
+		}
+
+		async runEditingSessionAction(accessor: ServicesAccessor, editingSession: IChatEditingSession) {
+			await editingSession.redoInteraction();
+		}
+	});
+
+	registerAction2(class RedoChatEditInteractionAction2 extends EditingSessionAction {
+		constructor() {
+			super({
+				id: 'workbench.action.chat.redoEdit2',
+				title: localize2('chat.redoEdit.label2', "Redo Last Request"),
+				category: CHAT_CATEGORY,
+				precondition: ContextKeyExpr.and(ChatContextKeys.chatEditingCanRedo, ChatContextKeys.enabled, ChatContextKeys.editingParticipantRegistered),
+				f1: true,
 				menu: [{
-					id: MenuId.ViewTitle,
+					id: MenuId.ChatMessageRestoreCheckpoint,
 					when: ContextKeyExpr.equals('view', ChatViewId),
 					group: 'navigation',
-					order: -2
+					order: -1
 				}]
 			});
 		}
 
 		async runEditingSessionAction(accessor: ServicesAccessor, editingSession: IChatEditingSession) {
-			await editingSession.redoInteraction();
+			const widget = accessor.get(IChatWidgetService);
+
+			const redoUntilDone = async () => {
+				if (!editingSession.canRedo.get()) {
+					return;
+				}
+
+				await editingSession.redoInteraction();
+				setTimeout(() => redoUntilDone(), 0);
+			};
+
+			await redoUntilDone();
+			widget.lastFocusedWidget?.viewModel?.model.setCheckpoint(undefined);
 		}
 	});
 }

src/vs/workbench/contrib/chat/browser/actions/chatExecuteActions.ts

@@ -142,6 +142,8 @@ abstract class SubmitAction extends Action2 {
 				const snapshotRequestId = chatRequests[itemIndex].id;
 				await session.restoreSnapshot(snapshotRequestId, undefined);
 			}
+		} else if (widget?.viewModel?.model.checkpoint) {
+			widget.viewModel.model.setCheckpoint(undefined);
 		}
 		widget?.acceptInput(context?.inputValue);
 	}

src/vs/workbench/contrib/chat/browser/chat.contribution.ts

@@ -268,6 +268,12 @@ configurationRegistry.registerConfiguration({
 			description: nls.localize('chat.emptyChatState', "Shows a modified empty chat state with hints in the input placeholder text."),
 			tags: ['experimental', 'onExp'],
 		},
+		'chat.checkpoints.enabled': {
+			type: 'boolean',
+			default: true,
+			description: nls.localize('chat.checkpoints.enabled', "Enables checkpoints in chat. Checkpoints allow you to restore the chat to a previous state."),
+			tags: ['experimental'],
+		},
 		[mcpEnabledSection]: {
 			type: 'boolean',
 			description: nls.localize('chat.mcp.enabled', "Enables integration with Model Context Protocol servers to provide additional tools and functionality."),

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

@@ -305,6 +305,74 @@ export class ChatEditingShowChangesAction extends EditingSessionAction {
 }
 registerAction2(ChatEditingShowChangesAction);
 
+async function restoreSnapshotWithConfirmation(accessor: ServicesAccessor, item: ChatTreeItem): Promise<void> {
+	const configurationService = accessor.get(IConfigurationService);
+	const dialogService = accessor.get(IDialogService);
+	const chatService = accessor.get(IChatService);
+	const chatModel = chatService.getSession(item.sessionId);
+	if (!chatModel) {
+		return;
+	}
+
+	const session = chatModel.editingSession;
+	if (!session) {
+		return;
+	}
+
+	const requestId = isRequestVM(item) ? item.id :
+		isResponseVM(item) ? item.requestId : undefined;
+
+	if (requestId) {
+		const chatRequests = chatModel.getRequests();
+		const itemIndex = chatRequests.findIndex(request => request.id === requestId);
+		const editsToUndo = chatRequests.length - itemIndex;
+
+		const requestsToRemove = chatRequests.slice(itemIndex);
+		const requestIdsToRemove = new Set(requestsToRemove.map(request => request.id));
+		const entriesModifiedInRequestsToRemove = session.entries.get().filter((entry) => requestIdsToRemove.has(entry.lastModifyingRequestId)) ?? [];
+		const shouldPrompt = entriesModifiedInRequestsToRemove.length > 0 && configurationService.getValue('chat.editing.confirmEditRequestRemoval') === true;
+
+		let message: string;
+		if (editsToUndo === 1) {
+			if (entriesModifiedInRequestsToRemove.length === 1) {
+				message = localize('chat.removeLast.confirmation.message2', "This will remove your last request and undo the edits made to {0}. Do you want to proceed?", basename(entriesModifiedInRequestsToRemove[0].modifiedURI));
+			} else {
+				message = localize('chat.removeLast.confirmation.multipleEdits.message', "This will remove your last request and undo edits made to {0} files in your working set. Do you want to proceed?", entriesModifiedInRequestsToRemove.length);
+			}
+		} else {
+			if (entriesModifiedInRequestsToRemove.length === 1) {
+				message = localize('chat.remove.confirmation.message2', "This will remove all subsequent requests and undo edits made to {0}. Do you want to proceed?", basename(entriesModifiedInRequestsToRemove[0].modifiedURI));
+			} else {
+				message = localize('chat.remove.confirmation.multipleEdits.message', "This will remove all subsequent requests and undo edits made to {0} files in your working set. Do you want to proceed?", entriesModifiedInRequestsToRemove.length);
+			}
+		}
+
+		const confirmation = shouldPrompt
+			? 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: message,
+				primaryButton: localize('chat.remove.confirmation.primaryButton', "Yes"),
+				checkbox: { label: localize('chat.remove.confirmation.checkbox', "Don't ask again"), checked: false },
+				type: 'info'
+			})
+			: { confirmed: true };
+
+		if (!confirmation.confirmed) {
+			return;
+		}
+
+		if (confirmation.checkboxChecked) {
+			await configurationService.updateValue('chat.editing.confirmEditRequestRemoval', false);
+		}
+
+		// Restore the snapshot to what it was before the request(s) that we deleted
+		const snapshotRequestId = chatRequests[itemIndex].id;
+		await session.restoreSnapshot(snapshotRequestId, undefined);
+	}
+}
+
 registerAction2(class RemoveAction extends Action2 {
 	constructor() {
 		super({
@@ -326,7 +394,7 @@ registerAction2(class RemoveAction extends Action2 {
 					id: MenuId.ChatMessageTitle,
 					group: 'navigation',
 					order: 2,
-					when: ContextKeyExpr.and(ContextKeyExpr.equals(`config.${ChatConfiguration.EditRequests}`, 'input').negate())
+					when: ContextKeyExpr.and(ContextKeyExpr.equals(`config.${ChatConfiguration.EditRequests}`, 'input').negate(), ContextKeyExpr.equals(`config.${ChatConfiguration.CheckpointsEnabled}`, false)),
 				}
 			]
 		});
@@ -335,6 +403,7 @@ registerAction2(class RemoveAction extends Action2 {
 	async run(accessor: ServicesAccessor, ...args: any[]) {
 		let item: ChatTreeItem | undefined = args[0];
 		const chatWidgetService = accessor.get(IChatWidgetService);
+		const configurationService = accessor.get(IConfigurationService);
 		const widget = chatWidgetService.lastFocusedWidget;
 		if (!isResponseVM(item) && !isRequestVM(item)) {
 			item = widget?.getFocus();
@@ -344,76 +413,56 @@ registerAction2(class RemoveAction extends Action2 {
 			return;
 		}
 
-		const configurationService = accessor.get(IConfigurationService);
-		const dialogService = accessor.get(IDialogService);
-		const chatService = accessor.get(IChatService);
-		const chatModel = chatService.getSession(item.sessionId);
-		if (!chatModel) {
-			return;
-		}
+		await restoreSnapshotWithConfirmation(accessor, item);
 
-		const session = chatModel.editingSession;
-		if (!session) {
-			return;
+		if (isRequestVM(item) && configurationService.getValue('chat.undoRequests.restoreInput')) {
+			widget?.focusInput();
+			widget?.input.setValue(item.messageText, false);
 		}
+	}
+});
 
-		const requestId = isRequestVM(item) ? item.id :
-			isResponseVM(item) ? item.requestId : undefined;
-
-		if (requestId) {
-			const chatRequests = chatModel.getRequests();
-			const itemIndex = chatRequests.findIndex(request => request.id === requestId);
-			const editsToUndo = chatRequests.length - itemIndex;
-
-			const requestsToRemove = chatRequests.slice(itemIndex);
-			const requestIdsToRemove = new Set(requestsToRemove.map(request => request.id));
-			const entriesModifiedInRequestsToRemove = session.entries.get().filter((entry) => requestIdsToRemove.has(entry.lastModifyingRequestId)) ?? [];
-			const shouldPrompt = entriesModifiedInRequestsToRemove.length > 0 && configurationService.getValue('chat.editing.confirmEditRequestRemoval') === true;
-
-			let message: string;
-			if (editsToUndo === 1) {
-				if (entriesModifiedInRequestsToRemove.length === 1) {
-					message = localize('chat.removeLast.confirmation.message2', "This will remove your last request and undo the edits made to {0}. Do you want to proceed?", basename(entriesModifiedInRequestsToRemove[0].modifiedURI));
-				} else {
-					message = localize('chat.removeLast.confirmation.multipleEdits.message', "This will remove your last request and undo edits made to {0} files in your working set. Do you want to proceed?", entriesModifiedInRequestsToRemove.length);
-				}
-			} else {
-				if (entriesModifiedInRequestsToRemove.length === 1) {
-					message = localize('chat.remove.confirmation.message2', "This will remove all subsequent requests and undo edits made to {0}. Do you want to proceed?", basename(entriesModifiedInRequestsToRemove[0].modifiedURI));
-				} else {
-					message = localize('chat.remove.confirmation.multipleEdits.message', "This will remove all subsequent requests and undo edits made to {0} files in your working set. Do you want to proceed?", entriesModifiedInRequestsToRemove.length);
+registerAction2(class RestoreCheckpointAction extends Action2 {
+	constructor() {
+		super({
+			id: 'workbench.action.chat.restoreCheckpoint',
+			title: localize2('chat.restoreCheckpoint.label', "Restore Checkpoint"),
+			tooltip: localize2('chat.restoreCheckpoint.tooltip', "Restore Checkpoint"),
+			f1: false,
+			category: CHAT_CATEGORY,
+			keybinding: {
+				primary: KeyCode.Delete,
+				mac: {
+					primary: KeyMod.CtrlCmd | KeyCode.Backspace,
+				},
+				when: ContextKeyExpr.and(ChatContextKeys.inChatSession, EditorContextKeys.textInputFocus.negate()),
+				weight: KeybindingWeight.WorkbenchContrib,
+			},
+			menu: [
+				{
+					id: MenuId.ChatMessageCheckpoint,
+					group: 'navigation',
+					order: 2,
+					when: ChatContextKeys.isRequest
 				}
-			}
-
-			const confirmation = shouldPrompt
-				? 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: message,
-					primaryButton: localize('chat.remove.confirmation.primaryButton', "Yes"),
-					checkbox: { label: localize('chat.remove.confirmation.checkbox', "Don't ask again"), checked: false },
-					type: 'info'
-				})
-				: { confirmed: true };
-
-			if (!confirmation.confirmed) {
-				return;
-			}
-
-			if (confirmation.checkboxChecked) {
-				await configurationService.updateValue('chat.editing.confirmEditRequestRemoval', false);
-			}
+			]
+		});
+	}
 
-			// Restore the snapshot to what it was before the request(s) that we deleted
-			const snapshotRequestId = chatRequests[itemIndex].id;
-			await session.restoreSnapshot(snapshotRequestId, undefined);
+	async run(accessor: ServicesAccessor, ...args: any[]) {
+		let item: ChatTreeItem | undefined = args[0];
+		const chatWidgetService = accessor.get(IChatWidgetService);
+		const widget = chatWidgetService.lastFocusedWidget;
+		if (!isResponseVM(item) && !isRequestVM(item)) {
+			item = widget?.getFocus();
 		}
 
-		if (isRequestVM(item) && configurationService.getValue('chat.undoRequests.restoreInput')) {
-			widget?.focusInput();
-			widget?.input.setValue(item.messageText, false);
+		if (!item) {
+			return;
 		}
+
+		widget?.viewModel?.model.setCheckpoint(item.id);
+		await restoreSnapshotWithConfirmation(accessor, item);
 	}
 });