improve checkpointing and add checkpoint hover on first request (#305572)

* improve checkpointing and add checkpoint hover on first request

* address comments

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

@@ -28,6 +28,7 @@ import { KeybindingWeight } from '../../../../../platform/keybinding/common/keyb
 import { IEditorPane } from '../../../../common/editor.js';
 import { IEditorService } from '../../../../services/editor/common/editorService.js';
 import { IAgentSessionsService } from '../agentSessions/agentSessionsService.js';
+import { IChatRequestVariableEntry, isImplicitVariableEntry, isPromptFileVariableEntry, isPromptTextVariableEntry, isStringVariableEntry, isWorkspaceVariableEntry } from '../../common/attachments/chatVariableEntries.js';
 import { isChatViewTitleActionContext } from '../../common/actions/chatActions.js';
 import { ChatContextKeys } from '../../common/actions/chatContextKeys.js';
 import { applyingChatEditsFailedContextKey, CHAT_EDITING_MULTI_DIFF_SOURCE_RESOLVER_SCHEME, chatEditingResourceContextKey, chatEditingWidgetFileStateContextKey, decidedChatEditingResourceContextKey, hasAppliedChatEditsContextKey, hasUndecidedChatEditingResourceContextKey, IChatEditingService, IChatEditingSession, ModifiedFileEntryState } from '../../common/editing/chatEditingService.js';
@@ -414,7 +415,20 @@ export class ViewAllSessionChangesAction extends Action2 {
 }
 registerAction2(ViewAllSessionChangesAction);
 
-async function restoreSnapshotWithConfirmationByRequestId(accessor: ServicesAccessor, sessionResource: URI, requestId: string): Promise<void> {
+function filterToUserAttachedContext(attachedContext: readonly IChatRequestVariableEntry[] | undefined): IChatRequestVariableEntry[] {
+	if (!attachedContext?.length) {
+		return [];
+	}
+	return attachedContext.filter(a =>
+		!isImplicitVariableEntry(a) &&
+		!isWorkspaceVariableEntry(a) &&
+		!isStringVariableEntry(a) &&
+		!(isPromptFileVariableEntry(a) && a.automaticallyAdded) &&
+		!(isPromptTextVariableEntry(a) && a.automaticallyAdded)
+	);
+}
+
+async function restoreSnapshotWithConfirmationByRequestId(accessor: ServicesAccessor, sessionResource: URI, requestId: string): Promise<boolean> {
 	const configurationService = accessor.get(IConfigurationService);
 	const dialogService = accessor.get(IDialogService);
 	const chatWidgetService = accessor.get(IChatWidgetService);
@@ -422,18 +436,18 @@ async function restoreSnapshotWithConfirmationByRequestId(accessor: ServicesAcce
 	const chatService = accessor.get(IChatService);
 	const chatModel = chatService.getSession(sessionResource);
 	if (!chatModel) {
-		return;
+		return false;
 	}
 
 	const session = chatModel.editingSession;
 	if (!session) {
-		return;
+		return false;
 	}
 
 	const chatRequests = chatModel.getRequests();
 	const itemIndex = chatRequests.findIndex(request => request.id === requestId);
 	if (itemIndex === -1) {
-		return;
+		return false;
 	}
 
 	const editsToUndo = chatRequests.length - itemIndex;
@@ -472,7 +486,7 @@ async function restoreSnapshotWithConfirmationByRequestId(accessor: ServicesAcce
 
 	if (!confirmation.confirmed) {
 		widget?.viewModel?.model.setCheckpoint(undefined);
-		return;
+		return false;
 	}
 
 	if (confirmation.checkboxChecked) {
@@ -482,17 +496,18 @@ async function restoreSnapshotWithConfirmationByRequestId(accessor: ServicesAcce
 	// Restore the snapshot to what it was before the request(s) that we deleted
 	const snapshotRequestId = chatRequests[itemIndex].id;
 	await session.restoreSnapshot(snapshotRequestId, undefined);
+	return true;
 }
 
-async function restoreSnapshotWithConfirmation(accessor: ServicesAccessor, item: ChatTreeItem): Promise<void> {
+async function restoreSnapshotWithConfirmation(accessor: ServicesAccessor, item: ChatTreeItem): Promise<boolean> {
 	const requestId = isRequestVM(item) ? item.id :
 		isResponseVM(item) ? item.requestId : undefined;
 
 	if (!requestId) {
-		return;
+		return false;
 	}
 
-	await restoreSnapshotWithConfirmationByRequestId(accessor, item.sessionResource, requestId);
+	return restoreSnapshotWithConfirmationByRequestId(accessor, item.sessionResource, requestId);
 }
 
 registerAction2(class RemoveAction extends Action2 {
@@ -535,11 +550,15 @@ registerAction2(class RemoveAction extends Action2 {
 			return;
 		}
 
-		await restoreSnapshotWithConfirmation(accessor, item);
+		const confirmed = await restoreSnapshotWithConfirmation(accessor, item);
 
-		if (isRequestVM(item) && configurationService.getValue('chat.undoRequests.restoreInput')) {
+		if (confirmed && isRequestVM(item) && configurationService.getValue('chat.undoRequests.restoreInput')) {
 			widget?.focusInput();
 			widget?.input.setValue(item.messageText, false);
+			const userAttachments = filterToUserAttachedContext(item.attachedContext);
+			if (userAttachments.length) {
+				await widget?.input.restoreAttachments(userAttachments);
+			}
 		}
 	}
 });
@@ -583,13 +602,19 @@ registerAction2(class RestoreCheckpointAction extends Action2 {
 			return;
 		}
 
+		const userAttachments = isRequestVM(item) ? filterToUserAttachedContext(item.attachedContext) : [];
+
 		if (isRequestVM(item)) {
 			widget?.focusInput();
 			widget?.input.setValue(item.messageText, false);
 		}
 
 		widget?.viewModel?.model.setCheckpoint(item.id);
-		await restoreSnapshotWithConfirmation(accessor, item);
+		const confirmed = await restoreSnapshotWithConfirmation(accessor, item);
+
+		if (confirmed && userAttachments.length) {
+			await widget?.input.restoreAttachments(userAttachments);
+		}
 	}
 });
 

src/vs/workbench/contrib/chat/browser/widget/input/chatInputPart.ts

@@ -16,6 +16,7 @@ import { HoverPosition } from '../../../../../../base/browser/ui/hover/hoverWidg
 import { IAction } from '../../../../../../base/common/actions.js';
 import { equals as arraysEqual } from '../../../../../../base/common/arrays.js';
 import { DeferredPromise, RunOnceScheduler } from '../../../../../../base/common/async.js';
+import { isDefined } from '../../../../../../base/common/types.js';
 import { CancellationToken } from '../../../../../../base/common/cancellation.js';
 import { Codicon } from '../../../../../../base/common/codicons.js';
 import { Emitter, Event } from '../../../../../../base/common/event.js';
@@ -1379,16 +1380,15 @@ export class ChatInputPart extends Disposable implements IHistoryNavigationWidge
 		this.navigateHistory(false);
 	}
 
-	private async navigateHistory(previous: boolean): Promise<void> {
+	/**
-		const historyEntry = previous ?
+	 * Restores attachments to the input, re-fetching image binary data as needed.
-			this.history.previous() : this.history.next();
+	 */
+	async restoreAttachments(attachments: readonly IChatRequestVariableEntry[]): Promise<void> {
+		let restored = [...attachments];
 
-		let historyAttachments = historyEntry?.attachments ?? [];
+		if (restored.length > 0) {
-
+			restored = (await Promise.all(restored.map(async (attachment) => {
-		// Check for images in history to restore the value.
+				if (isImageVariableEntry(attachment) && !attachment.value && attachment.references?.length && URI.isUri(attachment.references[0].reference)) {
-		if (historyAttachments.length > 0) {
-			historyAttachments = (await Promise.all(historyAttachments.map(async (attachment) => {
-				if (isImageVariableEntry(attachment) && attachment.references?.length && URI.isUri(attachment.references[0].reference)) {
 					const currReference = attachment.references[0].reference;
 					try {
 						const imageBinary = currReference.toString(true).startsWith('http') ? await this.sharedWebExtracterService.readImage(currReference, CancellationToken.None) : (await this.fileService.readFile(currReference)).value;
@@ -1396,7 +1396,7 @@ export class ChatInputPart extends Disposable implements IHistoryNavigationWidge
 							return undefined;
 						}
 						const newAttachment = { ...attachment };
-						newAttachment.value = (isImageVariableEntry(attachment) && attachment.isPasted) ? imageBinary.buffer : await resizeImage(imageBinary.buffer); // if pasted image, we do not need to resize.
+						newAttachment.value = (isImageVariableEntry(attachment) && attachment.isPasted) ? imageBinary.buffer : await resizeImage(imageBinary.buffer);
 						return newAttachment;
 					} catch (err) {
 						this.logService.error('Failed to fetch and reference.', err);
@@ -1404,10 +1404,17 @@ export class ChatInputPart extends Disposable implements IHistoryNavigationWidge
 					}
 				}
 				return attachment;
-			}))).filter(attachment => attachment !== undefined);
+			}))).filter(isDefined);
 		}
 
-		this._attachmentModel.clearAndSetContext(...historyAttachments);
+		this._attachmentModel.clearAndSetContext(...restored);
+	}
+
+	private async navigateHistory(previous: boolean): Promise<void> {
+		const historyEntry = previous ?
+			this.history.previous() : this.history.next();
+
+		await this.restoreAttachments(historyEntry?.attachments ?? []);
 
 		const inputText = historyEntry?.inputText ?? '';
 		const contribData = historyEntry?.contrib ?? {};

src/vs/workbench/contrib/chat/browser/widget/media/chat.css

@@ -3235,10 +3235,6 @@ have to be updated for changes to the rules above, or to support more deeply nes
 		.request-hover {
 			top: 0px;
 		}
-
-		.checkpoint-container {
-			display: none;
-		}
 	}
 
 	.interactive-list > .monaco-list:focus > .monaco-scrollable-element > .monaco-list-rows > .monaco-list-row.focused.request {

src/vs/workbench/contrib/chat/common/model/chatModel.ts

@@ -1887,7 +1887,7 @@ class InputModel implements IInputModel {
 
 		return {
 			contrib: value.contrib,
-			attachments: persistableAttachments,
+			attachments: persistableAttachments.map(IChatRequestVariableEntry.toExport),
 			mode: value.mode,
 			selectedModel: value.selectedModel ? {
 				identifier: value.selectedModel.identifier,
@@ -2179,7 +2179,7 @@ export class ChatModel extends Disposable implements IChatModel {
 		// Initialize input model from serialized data (undefined for new chats)
 		const serializedInputState = initialModelProps.inputState || (isValidFullData && initialData.inputState ? initialData.inputState : undefined);
 		this.inputModel = new InputModel(serializedInputState && {
-			attachments: serializedInputState.attachments,
+			attachments: (serializedInputState.attachments ?? []).map(IChatRequestVariableEntry.fromExport),
 			mode: serializedInputState.mode,
 			selectedModel: serializedInputState.selectedModel && {
 				identifier: serializedInputState.selectedModel.identifier,

src/vs/workbench/contrib/chat/common/model/chatSessionOperationLog.ts

@@ -9,6 +9,7 @@ import { equals as objectsEqual } from '../../../../../base/common/objects.js';
 import { isEqual as _urisEqual } from '../../../../../base/common/resources.js';
 import { hasKey } from '../../../../../base/common/types.js';
 import { URI, UriComponents } from '../../../../../base/common/uri.js';
+import { IChatRequestVariableEntry } from '../attachments/chatVariableEntries.js';
 import { IChatMarkdownContent, ResponseModelState } from '../chatService/chatService.js';
 import { ModifiedFileEntryState } from '../editing/chatEditingService.js';
 import { IParsedChatRequest } from '../requestParser/chatParserTypes.js';
@@ -116,7 +117,7 @@ const agentEditedFileEventSchema = Adapt.object<IChatAgentEditedFileEvent, IChat
 });
 
 const chatVariableSchema = Adapt.object<IChatRequestVariableData, IChatRequestVariableData>({
-	variables: Adapt.t(v => v.variables, Adapt.array(Adapt.value((a, b) => a.name === b.name))),
+	variables: Adapt.t(v => v.variables.map(IChatRequestVariableEntry.toExport), Adapt.array(Adapt.value((a, b) => a.name === b.name))),
 });
 
 const requestSchema = Adapt.object<IChatRequestModel, ISerializableChatRequestData>({
@@ -155,7 +156,7 @@ const requestSchema = Adapt.object<IChatRequestModel, ISerializableChatRequestDa
 });
 
 const inputStateSchema = Adapt.object<ISerializableChatModelInputState, ISerializableChatModelInputState>({
-	attachments: Adapt.v(i => i.attachments, objectsEqual),
+	attachments: Adapt.v(i => i.attachments.map(IChatRequestVariableEntry.toExport), objectsEqual),
 	mode: Adapt.v(i => i.mode, (a, b) => a.id === b.id),
 	selectedModel: Adapt.v(i => i.selectedModel, (a, b) => a?.identifier === b?.identifier),
 	inputText: Adapt.v(i => i.inputText),

src/vs/workbench/contrib/chat/common/model/chatViewModel.ts

@@ -94,6 +94,7 @@ export interface IChatRequestViewModel {
 	readonly slashCommand: IChatAgentCommand | undefined;
 	readonly agentOrSlashCommandDetected: boolean;
 	readonly shouldBeBlocked: IObservable<boolean>;
+	readonly attachedContext?: readonly IChatRequestVariableEntry[];
 	readonly modelId?: string;
 	readonly timestamp: number;
 	/** The kind of pending request, or undefined if not pending */
@@ -462,6 +463,10 @@ export class ChatRequestViewModel implements IChatRequestViewModel {
 
 	currentRenderedHeight: number | undefined;
 
+	get attachedContext() {
+		return this._model.attachedContext;
+	}
+
 	get modelId() {
 		return this._model.modelId;
 	}