Support restoring edits for background sessions (#279270)

* Support restoring edits for background sessions

* revert change

src/vs/workbench/api/browser/mainThreadChatAgents2.ts

@@ -264,7 +264,7 @@ export class MainThreadChatAgents2 extends Disposable implements MainThreadChatA
 				const response = chatSession?.getRequests().at(-1)?.response;
 				if (chatSession?.editingSession && responsePartHandle !== undefined && response) {
 					const parts = progress.start
-						? await chatSession.editingSession.startExternalEdits(response, responsePartHandle, revive(progress.resources))
+						? await chatSession.editingSession.startExternalEdits(response, responsePartHandle, revive(progress.resources), progress.undoStopId)
 						: await chatSession.editingSession.stopExternalEdits(response, responsePartHandle);
 					chatProgressParts.push(...parts);
 				}

src/vs/workbench/api/browser/mainThreadChatSessions.ts

@@ -120,7 +120,8 @@ export class ObservableChatSession extends Disposable implements IChatSession {
 						prompt: turn.prompt,
 						participant: turn.participant,
 						command: turn.command,
-						variableData: variables ? { variables } : undefined
+						variableData: variables ? { variables } : undefined,
+						id: turn.id
 					};
 				}
 

src/vs/workbench/api/common/extHost.protocol.ts

@@ -3252,6 +3252,7 @@ export interface MainThreadChatStatusShape {
 }
 
 export type IChatSessionHistoryItemDto = {
+	id?: string;
 	type: 'request';
 	prompt: string;
 	participant: string;

src/vs/workbench/api/common/extHostChatAgents2.ts

@@ -279,12 +279,13 @@ export class ChatAgentResponseStream {
 					throwIfDone(this.externalEdit);
 					const resources = Array.isArray(target) ? target : [target];
 					const operationId = taskHandlePool++;
-
-					await send({ kind: 'externalEdits', start: true, resources }, operationId);
+					const undoStopId = generateUuid();
+					await send({ kind: 'externalEdits', start: true, resources, undoStopId }, operationId);
 					try {
-						return await callback();
+						await callback();
+						return undoStopId;
 					} finally {
-						await send({ kind: 'externalEdits', start: false, resources }, operationId);
+						await send({ kind: 'externalEdits', start: false, resources, undoStopId }, operationId);
 					}
 				},
 				confirmation(title, message, data, buttons) {
@@ -359,7 +360,7 @@ export class ChatAgentResponseStream {
 						return this;
 					} else if (part instanceof extHostTypes.ChatResponseExternalEditPart) {
 						const p = this.externalEdit(part.uris, part.callback);
-						p.then(() => part.didGetApplied());
+						p.then((value) => part.didGetApplied(value));
 						return this;
 					} else {
 						const dto = typeConvert.ChatResponsePart.from(part, that._commandsConverter, that._sessionDisposables);
@@ -702,7 +703,7 @@ export class ExtHostChatAgents2 extends Disposable implements ExtHostChatAgentsS
 			}
 
 			const editedFileEvents = isProposedApiEnabled(extension, 'chatParticipantPrivate') ? h.request.editedFileEvents : undefined;
-			const turn = new extHostTypes.ChatRequestTurn(h.request.message, h.request.command, varsWithoutTools, h.request.agentId, toolReferences, editedFileEvents);
+			const turn = new extHostTypes.ChatRequestTurn(h.request.message, h.request.command, varsWithoutTools, h.request.agentId, toolReferences, editedFileEvents, h.request.requestId);
 			res.push(turn);
 
 			// RESPONSE turn

src/vs/workbench/api/common/extHostChatSessions.ts

@@ -398,6 +398,7 @@ export class ExtHostChatSessions extends Disposable implements ExtHostChatSessio
 		const variables = turn.references.map(ref => this.convertReferenceToVariable(ref));
 		return {
 			type: 'request' as const,
+			id: turn.id,
 			prompt: turn.prompt,
 			participant: turn.participant,
 			command: turn.command,

src/vs/workbench/api/common/extHostTypeConverters.ts

@@ -2608,10 +2608,11 @@ export namespace ChatResponseCodeblockUriPart {
 			kind: 'codeblockUri',
 			uri: part.value,
 			isEdit: part.isEdit,
+			undoStopId: part.undoStopId
 		};
 	}
 	export function to(part: Dto<IChatResponseCodeblockUriPart>): vscode.ChatResponseCodeblockUriPart {
-		return new types.ChatResponseCodeblockUriPart(URI.revive(part.uri), part.isEdit);
+		return new types.ChatResponseCodeblockUriPart(URI.revive(part.uri), part.isEdit, part.undoStopId);
 	}
 }
 

src/vs/workbench/api/common/extHostTypes.ts

@@ -3166,14 +3166,14 @@ export class ChatResponseMultiDiffPart {
 }
 
 export class ChatResponseExternalEditPart {
-	applied: Thenable<void>;
-	didGetApplied!: () => void;
+	applied: Thenable<string>;
+	didGetApplied!: (value: string) => void;
 
 	constructor(
 		public uris: vscode.Uri[],
 		public callback: () => Thenable<unknown>,
 	) {
-		this.applied = new Promise<void>((resolve) => {
+		this.applied = new Promise<string>((resolve) => {
 			this.didGetApplied = resolve;
 		});
 	}
@@ -3252,10 +3252,12 @@ export class ChatResponseReferencePart {
 
 export class ChatResponseCodeblockUriPart {
 	isEdit?: boolean;
+	undoStopId?: string;
 	value: vscode.Uri;
-	constructor(value: vscode.Uri, isEdit?: boolean) {
+	constructor(value: vscode.Uri, isEdit?: boolean, undoStopId?: string) {
 		this.value = value;
 		this.isEdit = isEdit;
+		this.undoStopId = undoStopId;
 	}
 }
 
@@ -3385,7 +3387,8 @@ export class ChatRequestTurn implements vscode.ChatRequestTurn2 {
 		readonly references: vscode.ChatPromptReference[],
 		readonly participant: string,
 		readonly toolReferences: vscode.ChatLanguageModelToolReference[],
-		readonly editedFileEvents?: vscode.ChatRequestEditedFileEvent[]
+		readonly editedFileEvents?: vscode.ChatRequestEditedFileEvent[],
+		readonly id?: string
 	) { }
 }
 

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

@@ -16,7 +16,6 @@ import { derived, IObservable, IReader, ITransaction, observableValue, transacti
 import { isEqual } from '../../../../../base/common/resources.js';
 import { hasKey, Mutable } from '../../../../../base/common/types.js';
 import { URI } from '../../../../../base/common/uri.js';
-import { generateUuid } from '../../../../../base/common/uuid.js';
 import { IBulkEditService } from '../../../../../editor/browser/services/bulkEditService.js';
 import { Range } from '../../../../../editor/common/core/range.js';
 import { TextEdit } from '../../../../../editor/common/languages.js';
@@ -89,7 +88,7 @@ class ThrottledSequencer extends Sequencer {
 	}
 }
 
-function createOpeningEditCodeBlock(uri: URI, isNotebook: boolean): IChatProgress[] {
+function createOpeningEditCodeBlock(uri: URI, isNotebook: boolean, undoStopId: string): IChatProgress[] {
 	return [
 		{
 			kind: 'markdownContent',
@@ -98,7 +97,8 @@ function createOpeningEditCodeBlock(uri: URI, isNotebook: boolean): IChatProgres
 		{
 			kind: 'codeblockUri',
 			uri,
-			isEdit: true
+			isEdit: true,
+			undoStopId
 		},
 		{
 			kind: 'markdownContent',
@@ -532,15 +532,11 @@ export class ChatEditingSession extends Disposable implements IChatEditingSessio
 		};
 	}
 
-	async startExternalEdits(responseModel: IChatResponseModel, operationId: number, resources: URI[]): Promise<IChatProgress[]> {
+	async startExternalEdits(responseModel: IChatResponseModel, operationId: number, resources: URI[], undoStopId: string): Promise<IChatProgress[]> {
 		const snapshots = new ResourceMap<string | undefined>();
 		const acquiredLockPromises: DeferredPromise<void>[] = [];
 		const releaseLockPromises: DeferredPromise<void>[] = [];
-		const undoStopId = generateUuid();
-		const progress: IChatProgress[] = [{
-			kind: 'undoStop',
-			id: undoStopId,
-		}];
+		const progress: IChatProgress[] = [];
 		const telemetryInfo = this._getTelemetryInfoForModel(responseModel);
 
 		await chatEditingSessionIsReady(this);
@@ -566,7 +562,7 @@ export class ChatEditingSession extends Disposable implements IChatEditingSessio
 
 
 				const notebookUri = CellUri.parse(resource)?.notebook || resource;
-				progress.push(...createOpeningEditCodeBlock(resource, this._notebookService.hasSupportedNotebooks(notebookUri)));
+				progress.push(...createOpeningEditCodeBlock(resource, this._notebookService.hasSupportedNotebooks(notebookUri), undoStopId));
 
 				// Save to disk to ensure disk state is current before external edits
 				await entry?.save();

src/vs/workbench/contrib/chat/common/chatEditingService.ts

@@ -134,7 +134,7 @@ export interface IChatEditingSession extends IDisposable {
 	 * agents that make changes on-disk rather than streaming edits through the
 	 * chat session.
 	 */
-	startExternalEdits(responseModel: IChatResponseModel, operationId: number, resources: URI[]): Promise<IChatProgress[]>;
+	startExternalEdits(responseModel: IChatResponseModel, operationId: number, resources: URI[], undoStopId: string): Promise<IChatProgress[]>;
 	stopExternalEdits(responseModel: IChatResponseModel, operationId: number): Promise<IChatProgress[]>;
 
 	/**

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

@@ -1932,10 +1932,11 @@ export class ChatModel extends Disposable implements IChatModel {
 		this._onDidChange.fire({ kind: 'setHidden' });
 	}
 
-	addRequest(message: IParsedChatRequest, variableData: IChatRequestVariableData, attempt: number, modeInfo?: IChatRequestModeInfo, chatAgent?: IChatAgentData, slashCommand?: IChatAgentCommand, confirmation?: string, locationData?: IChatLocationData, attachments?: IChatRequestVariableEntry[], isCompleteAddedRequest?: boolean, modelId?: string, userSelectedTools?: UserSelectedTools): ChatRequestModel {
+	addRequest(message: IParsedChatRequest, variableData: IChatRequestVariableData, attempt: number, modeInfo?: IChatRequestModeInfo, chatAgent?: IChatAgentData, slashCommand?: IChatAgentCommand, confirmation?: string, locationData?: IChatLocationData, attachments?: IChatRequestVariableEntry[], isCompleteAddedRequest?: boolean, modelId?: string, userSelectedTools?: UserSelectedTools, id?: string): ChatRequestModel {
 		const editedFileEvents = [...this.currentEditedFileEvents.values()];
 		this.currentEditedFileEvents.clear();
 		const request = new ChatRequestModel({
+			restoredId: id,
 			session: this,
 			message,
 			variableData,
@@ -2016,7 +2017,7 @@ export class ChatModel extends Disposable implements IChatModel {
 		} else if (progress.kind === 'move') {
 			this._onDidChange.fire({ kind: 'move', target: progress.uri, range: progress.range });
 		} else if (progress.kind === 'codeblockUri' && progress.isEdit) {
-			request.response.addUndoStop({ id: generateUuid(), kind: 'undoStop' });
+			request.response.addUndoStop({ id: progress.undoStopId ?? generateUuid(), kind: 'undoStop' });
 			request.response.updateContent(progress, quiet);
 		} else if (progress.kind === 'progressTaskResult') {
 			// Should have been handled upstream, not sent to model
@@ -2061,7 +2062,6 @@ export class ChatModel extends Disposable implements IChatModel {
 			// Maybe something went wrong?
 			return;
 		}
-
 		request.response.setFollowups(followups);
 	}
 

src/vs/workbench/contrib/chat/common/chatService.ts

@@ -194,6 +194,7 @@ export interface IChatUndoStop {
 
 export interface IChatExternalEditsDto {
 	kind: 'externalEdits';
+	undoStopId: string;
 	start: boolean; /** true=start, false=stop */
 	resources: UriComponents[];
 }
@@ -228,6 +229,7 @@ export interface IChatResponseCodeblockUriPart {
 	kind: 'codeblockUri';
 	uri: URI;
 	isEdit?: boolean;
+	undoStopId?: string;
 }
 
 export interface IChatAgentMarkdownContentWithVulnerability {

src/vs/workbench/contrib/chat/common/chatServiceImpl.ts

@@ -653,7 +653,10 @@ export class ChatService extends Disposable implements IChatService {
 					undefined, // confirmation
 					undefined, // locationData
 					undefined, // attachments
-					true // isCompleteAddedRequest - this indicates it's a complete request, not user input
+					false, // Do not treat as requests completed, else edit pills won't show.
+					undefined,
+					undefined,
+					message.id
 				);
 			} else {
 				// response

src/vs/workbench/contrib/chat/common/chatSessionsService.ts

@@ -81,6 +81,7 @@ export interface IChatSessionItem {
 }
 
 export type IChatSessionHistoryItem = {
+	id?: string;
 	type: 'request';
 	prompt: string;
 	participant: string;

src/vscode-dts/vscode.proposed.chatParticipantAdditions.d.ts

@@ -32,7 +32,8 @@ declare module 'vscode' {
 	export class ChatResponseCodeblockUriPart {
 		isEdit?: boolean;
 		value: Uri;
-		constructor(value: Uri, isEdit?: boolean);
+		undoStopId?: string;
+		constructor(value: Uri, isEdit?: boolean, undoStopId?: string);
 	}
 
 	/**
@@ -170,7 +171,7 @@ declare module 'vscode' {
 	export class ChatResponseExternalEditPart {
 		uris: Uri[];
 		callback: () => Thenable<unknown>;
-		applied: Thenable<void>;
+		applied: Thenable<string>;
 		constructor(uris: Uri[], callback: () => Thenable<unknown>);
 	}
 
@@ -314,7 +315,7 @@ declare module 'vscode' {
 		 * tracked as agent edits. This can be used to track edits made from
 		 * external tools that don't generate simple {@link textEdit textEdits}.
 		 */
-		externalEdit<T>(target: Uri | Uri[], callback: () => Thenable<T>): Thenable<T>;
+		externalEdit(target: Uri | Uri[], callback: () => Thenable<unknown>): Thenable<string>;
 
 		markdownWithVulnerabilities(value: string | MarkdownString, vulnerabilities: ChatVulnerability[]): void;
 		codeblockUri(uri: Uri, isEdit?: boolean): void;

src/vscode-dts/vscode.proposed.chatParticipantPrivate.d.ts

@@ -148,6 +148,11 @@ declare module 'vscode' {
 	}
 
 	export class ChatResponseTurn2 {
+		/**
+		 * The id of the chat response. Used to identity an interaction with any of the chat surfaces.
+		 */
+		readonly id?: string;
+
 		/**
 		 * The content that was received from the chat participant. Only the stream parts that represent actual content (not metadata) are represented.
 		 */