WIP: Notebook Edit Support

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

@@ -28,11 +28,12 @@ import { ChatAgentLocation, IChatAgentHistoryEntry, IChatAgentImplementation, IC
 import { IChatEditingService, IChatRelatedFileProviderMetadata } from '../../contrib/chat/common/chatEditingService.js';
 import { ChatRequestAgentPart } from '../../contrib/chat/common/chatParserTypes.js';
 import { ChatRequestParser } from '../../contrib/chat/common/chatRequestParser.js';
-import { IChatContentInlineReference, IChatContentReference, IChatFollowup, IChatProgress, IChatService, IChatTask, IChatWarningMessage } from '../../contrib/chat/common/chatService.js';
+import { IChatContentInlineReference, IChatContentReference, IChatFollowup, IChatNotebookEdit, IChatProgress, IChatService, IChatTask, IChatWarningMessage } from '../../contrib/chat/common/chatService.js';
 import { IExtHostContext, extHostNamedCustomer } from '../../services/extensions/common/extHostCustomers.js';
 import { IExtensionService } from '../../services/extensions/common/extensions.js';
 import { Dto } from '../../services/extensions/common/proxyIdentifier.js';
-import { ExtHostChatAgentsShape2, ExtHostContext, IChatParticipantMetadata, IChatProgressDto, IDynamicChatAgentProps, IExtensionChatAgentMetadata, MainContext, MainThreadChatAgentsShape2 } from '../common/extHost.protocol.js';
+import { ExtHostChatAgentsShape2, ExtHostContext, IChatNotebookEditDto, IChatParticipantMetadata, IChatProgressDto, IDynamicChatAgentProps, IExtensionChatAgentMetadata, MainContext, MainThreadChatAgentsShape2 } from '../common/extHost.protocol.js';
+import { NotebookDto } from './mainThreadNotebookDto.js';
 
 interface AgentData {
 	dispose: () => void;
@@ -222,7 +223,7 @@ export class MainThreadChatAgents2 extends Disposable implements MainThreadChatA
 	}
 
 	async $handleProgressChunk(requestId: string, progress: IChatProgressDto, responsePartHandle?: number): Promise<number | void> {
-		const revivedProgress = revive(progress) as IChatProgress;
+		const revivedProgress = progress.kind === 'notebookEdit' ? ChatNotebookEdit.fromChatEdit(revive(progress)) : revive(progress) as IChatProgress;
 		if (revivedProgress.kind === 'progressTask') {
 			const handle = ++this._responsePartHandlePool;
 			const responsePartId = `${requestId}_${handle}`;
@@ -382,3 +383,35 @@ function computeCompletionRanges(model: ITextModel, position: Position, reg: Reg
 
 	return { insert, replace };
 }
+
+namespace ChatNotebookEdit {
+	export function fromChatEdit(part: IChatNotebookEditDto): IChatNotebookEdit {
+		return {
+			kind: 'notebookEdit',
+			uri: part.uri,
+			edits: part.edits.map(e => {
+				return {
+					count: e.count,
+					editType: e.editType,
+					index: e.index,
+					cells: e.cells.map(NotebookDto.fromNotebookCellDataDto)
+				};
+			})
+		};
+	}
+	export function toChatEdit(part: IChatNotebookEdit): IChatNotebookEditDto {
+		return {
+			kind: 'notebookEdit',
+			uri: URI.revive(part.uri),
+			edits: part.edits.map(e => {
+				return {
+					count: e.count,
+					editType: e.editType,
+					index: e.index,
+					cells: e.cells.map(NotebookDto.toNotebookCellDataDto)
+				};
+			})
+		};
+	}
+
+}

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

@@ -5,9 +5,12 @@
 import { CancellationToken } from '../../../base/common/cancellation.js';
 import { Disposable, DisposableMap, IDisposable } from '../../../base/common/lifecycle.js';
 import { URI } from '../../../base/common/uri.js';
+import { TextEdit } from '../../../editor/common/languages.js';
 import { ICodeMapperProvider, ICodeMapperRequest, ICodeMapperResponse, ICodeMapperService } from '../../contrib/chat/common/chatCodeMapperService.js';
+import { CellEditType, ICellEditOperation } from '../../contrib/notebook/common/notebookCommon.js';
 import { extHostNamedCustomer, IExtHostContext } from '../../services/extensions/common/extHostCustomers.js';
 import { ExtHostCodeMapperShape, ExtHostContext, ICodeMapperProgressDto, ICodeMapperRequestDto, MainContext, MainThreadCodeMapperShape } from '../common/extHost.protocol.js';
+import { NotebookDto } from './mainThreadNotebookDto.js';
 
 @extHostNamedCustomer(MainContext.MainThreadCodeMapper)
 export class MainThreadChatCodemapper extends Disposable implements MainThreadCodeMapperShape {
@@ -54,9 +57,36 @@ export class MainThreadChatCodemapper extends Disposable implements MainThreadCo
 	$handleProgress(requestId: string, data: ICodeMapperProgressDto): Promise<void> {
 		const response = this._responseMap.get(requestId);
 		if (response) {
+			const edits = data.edits;
 			const resource = URI.revive(data.uri);
-			response.textEdit(resource, data.edits);
+			if (!edits.length) {
+				response.textEdit(resource, []);
+			} else if (areTextEdits(edits)) {
+				response.textEdit(resource, edits);
+			} else {
+				const cellEdits: ICellEditOperation[] = [];
+				edits.forEach(dto => {
+					if (dto.editType === CellEditType.Replace) {
+						cellEdits.push({
+							editType: dto.editType,
+							index: dto.index,
+							count: dto.count,
+							cells: dto.cells.map(NotebookDto.fromNotebookCellDataDto)
+						});
+					}
+				});
+				response.notebookEdit(resource, cellEdits);
+			}
+		} else {
 		}
 		return Promise.resolve();
 	}
 }
+
+function areTextEdits(edits: ICodeMapperProgressDto['edits']): edits is TextEdit[] {
+	if (edits.some(e => 'range' in e && 'text' in e)) {
+		return true;
+	} else {
+		return false;
+	}
+}

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

@@ -1776,6 +1776,7 @@ export function createApiFactoryAndRegisterActors(accessor: ServicesAccessor): I
 			ChatResponseCodeblockUriPart: extHostTypes.ChatResponseCodeblockUriPart,
 			ChatResponseWarningPart: extHostTypes.ChatResponseWarningPart,
 			ChatResponseTextEditPart: extHostTypes.ChatResponseTextEditPart,
+			ChatResponseNotebookEditPart: extHostTypes.ChatResponseNotebookEditPart,
 			ChatResponseMarkdownWithVulnerabilitiesPart: extHostTypes.ChatResponseMarkdownWithVulnerabilitiesPart,
 			ChatResponseCommandButtonPart: extHostTypes.ChatResponseCommandButtonPart,
 			ChatResponseDetectedParticipantPart: extHostTypes.ChatResponseDetectedParticipantPart,

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

@@ -54,7 +54,7 @@ import { ChatAgentLocation, IChatAgentMetadata, IChatAgentRequest, IChatAgentRes
 import { ICodeMapperRequest, ICodeMapperResult } from '../../contrib/chat/common/chatCodeMapperService.js';
 import { IChatRelatedFile, IChatRelatedFileProviderMetadata as IChatRelatedFilesProviderMetadata, IChatRequestDraft } from '../../contrib/chat/common/chatEditingService.js';
 import { IChatProgressHistoryResponseContent } from '../../contrib/chat/common/chatModel.js';
-import { IChatContentInlineReference, IChatFollowup, IChatProgress, IChatResponseErrorDetails, IChatTask, IChatTaskDto, IChatUserActionEvent, IChatVoteAction } from '../../contrib/chat/common/chatService.js';
+import { IChatContentInlineReference, IChatFollowup, IChatNotebookEdit, IChatProgress, IChatResponseErrorDetails, IChatTask, IChatTaskDto, IChatUserActionEvent, IChatVoteAction } from '../../contrib/chat/common/chatService.js';
 import { IChatRequestVariableValue, IChatVariableData, IChatVariableResolverProgress } from '../../contrib/chat/common/chatVariables.js';
 import { IChatMessage, IChatResponseFragment, ILanguageModelChatMetadata, ILanguageModelChatSelector, ILanguageModelsChangeEvent } from '../../contrib/chat/common/languageModels.js';
 import { IPreparedToolInvocation, IToolData, IToolInvocation, IToolResult } from '../../contrib/chat/common/languageModelToolsService.js';
@@ -1296,7 +1296,12 @@ export interface ICodeMapperTextEdit {
 	edits: languages.TextEdit[];
 }
 
-export type ICodeMapperProgressDto = Dto<ICodeMapperTextEdit>;
+export interface ICodeMapperNotebookEditDto {
+	uri: URI;
+	edits: ICellEditOperationDto[];
+}
+
+export type ICodeMapperProgressDto = Dto<ICodeMapperTextEdit> | ICodeMapperNotebookEditDto;
 
 export interface MainThreadCodeMapperShape extends IDisposable {
 	$registerCodeMapperProvider(handle: number): void;
@@ -1418,8 +1423,9 @@ export type IDocumentContextDto = {
 };
 
 export type IChatProgressDto =
-	| Dto<Exclude<IChatProgress, IChatTask>>
-	| IChatTaskDto;
+	| Dto<Exclude<IChatProgress, IChatTask | IChatNotebookEdit>>
+	| IChatTaskDto
+	| IChatNotebookEditDto;
 
 export interface ExtHostUrlsShape {
 	$handleExternalUri(handle: number, uri: UriComponents): Promise<void>;
@@ -2132,6 +2138,18 @@ export interface IWorkspaceEditEntryMetadataDto {
 	iconPath?: { id: string } | UriComponents | { light: UriComponents; dark: UriComponents };
 }
 
+export interface IChatNotebookEditDto {
+	uri: URI;
+	edits: ICellEditReplaceOperationDto[];
+	kind: 'notebookEdit';
+}
+
+export type ICellEditReplaceOperationDto = {
+	editType: notebookCommon.CellEditType.Replace;
+	index: number;
+	count: number;
+	cells: NotebookCellDataDto[];
+};
 
 export type ICellEditOperationDto =
 	notebookCommon.ICellMetadataEdit

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

@@ -224,6 +224,15 @@ class ChatAgentResponseStream {
 					_report(dto);
 					return this;
 				},
+				notebookEdit(target, edits) {
+					throwIfDone(this.notebookEdit);
+					checkProposedApiEnabled(that._extension, 'chatParticipantAdditions');
+
+					const part = new extHostTypes.ChatResponseNotebookEditPart(target, Array.isArray(edits) ? edits : []);
+					const dto = typeConvert.ChatResponseNotebookEditPart.from(part);
+					_report(dto);
+					return this;
+				},
 				detectedParticipant(participant, command) {
 					throwIfDone(this.detectedParticipant);
 					checkProposedApiEnabled(that._extension, 'chatParticipantAdditions');
@@ -247,6 +256,7 @@ class ChatAgentResponseStream {
 
 					if (
 						part instanceof extHostTypes.ChatResponseTextEditPart ||
+						part instanceof extHostTypes.ChatResponseNotebookEditPart ||
 						part instanceof extHostTypes.ChatResponseMarkdownWithVulnerabilitiesPart ||
 						part instanceof extHostTypes.ChatResponseDetectedParticipantPart ||
 						part instanceof extHostTypes.ChatResponseWarningPart ||

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

@@ -8,8 +8,9 @@ import { CancellationToken } from '../../../base/common/cancellation.js';
 import { IExtensionDescription } from '../../../platform/extensions/common/extensions.js';
 import { ICodeMapperResult } from '../../contrib/chat/common/chatCodeMapperService.js';
 import * as extHostProtocol from './extHost.protocol.js';
-import { ChatAgentResult, DocumentContextItem, TextEdit } from './extHostTypeConverters.js';
+import { ChatAgentResult, DocumentContextItem, NotebookEdit, TextEdit } from './extHostTypeConverters.js';
 import { URI } from '../../../base/common/uri.js';
+import { isDefined } from '../../../base/common/types.js';
 
 export class ExtHostCodeMapper implements extHostProtocol.ExtHostCodeMapperShape {
 
@@ -38,6 +39,13 @@ export class ExtHostCodeMapper implements extHostProtocol.ExtHostCodeMapperShape
 					uri: target,
 					edits: edits.map(TextEdit.from)
 				});
+			},
+			notebookEdit: (target: vscode.Uri, edits: vscode.NotebookEdit | vscode.NotebookEdit[]) => {
+				edits = (Array.isArray(edits) ? edits : [edits]);
+				this._proxy.$handleProgress(internalRequest.requestId, {
+					uri: target,
+					edits: edits.map(NotebookEdit.toEditReplaceOperation).filter(isDefined)
+				});
 			}
 		};
 

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

@@ -58,6 +58,7 @@ import { IChatResponseTextPart, IChatResponsePromptTsxPart } from '../../contrib
 import { LanguageModelTextPart, LanguageModelPromptTsxPart } from './extHostTypes.js';
 import { MarshalledId } from '../../../base/common/marshallingIds.js';
 import { IChatRequestDraft } from '../../contrib/chat/common/chatEditingService.js';
+import { CellEditType } from '../../contrib/notebook/common/notebookCommon.js';
 
 export namespace Command {
 
@@ -2653,6 +2654,43 @@ export namespace ChatResponseTextEditPart {
 
 }
 
+
+export namespace NotebookEdit {
+	export function toEditReplaceOperation(edit: vscode.NotebookEdit): Dto<extHostProtocol.ICellEditReplaceOperationDto | undefined> {
+		// We are only interested in cell replaces (insertions, deletions, replacements)
+		if (!edit.newCellMetadata && !edit.newNotebookMetadata) {
+			return {
+				editType: CellEditType.Replace,
+				index: edit.range.start,
+				count: edit.range.end - edit.range.start,
+				cells: edit.newCells.map(NotebookCellData.from)
+			};
+		}
+		return undefined;
+	}
+
+	export function fromEditReplaceOperation(edit: Dto<extHostProtocol.ICellEditReplaceOperationDto>): vscode.NotebookEdit {
+		return new types.NotebookEdit(new types.NotebookRange(edit.index, edit.index + edit.count), edit.cells.map(NotebookCellData.to));
+	}
+}
+
+
+export namespace ChatResponseNotebookEditPart {
+	export function from(part: vscode.ChatResponseNotebookEditPart): extHostProtocol.IChatNotebookEditDto {
+		return {
+			kind: 'notebookEdit',
+			uri: URI.revive(part.uri),
+			// We are only interested in cell replaces (insertions, deletions, replacements)
+			edits: part.edits.map(e => NotebookEdit.toEditReplaceOperation(e)).filter(isDefined)
+		};
+	}
+
+	export function to(part: extHostProtocol.IChatNotebookEditDto): vscode.ChatResponseNotebookEditPart {
+		return new types.ChatResponseNotebookEditPart(URI.revive(part.uri), part.edits.map(NotebookEdit.fromEditReplaceOperation));
+	}
+}
+
+
 export namespace ChatResponseReferencePart {
 	export function from(part: types.ChatResponseReferencePart): Dto<IChatContentReference> {
 		const iconPath = ThemeIcon.isThemeIcon(part.iconPath) ? part.iconPath
@@ -2728,6 +2766,8 @@ export namespace ChatResponsePart {
 			return ChatResponseCommandButtonPart.from(part, commandsConverter, commandDisposables);
 		} else if (part instanceof types.ChatResponseTextEditPart) {
 			return ChatResponseTextEditPart.from(part);
+		} else if (part instanceof types.ChatResponseNotebookEditPart) {
+			return ChatResponseNotebookEditPart.from(part);
 		} else if (part instanceof types.ChatResponseMarkdownWithVulnerabilitiesPart) {
 			return ChatResponseMarkdownWithVulnerabilitiesPart.from(part);
 		} else if (part instanceof types.ChatResponseCodeblockUriPart) {

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

@@ -4614,6 +4614,15 @@ export class ChatResponseTextEditPart implements vscode.ChatResponseTextEditPart
 	}
 }
 
+export class ChatResponseNotebookEditPart implements vscode.ChatResponseNotebookEditPart {
+	uri: vscode.Uri;
+	edits: vscode.NotebookEdit[];
+	constructor(uri: vscode.Uri, editsOrDone: vscode.NotebookEdit | vscode.NotebookEdit[]) {
+		this.uri = uri;
+		this.edits = Array.isArray(editsOrDone) ? editsOrDone : [editsOrDone];
+	}
+}
+
 export class ChatRequestTurn implements vscode.ChatRequestTurn {
 	constructor(
 		readonly prompt: string,

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

@@ -31,8 +31,8 @@ import { IMultiDiffSourceResolver, IMultiDiffSourceResolverService, IResolvedMul
 import { ChatAgentLocation, IChatAgentService } from '../../common/chatAgents.js';
 import { ChatContextKeys } from '../../common/chatContextKeys.js';
 import { applyingChatEditsContextKey, applyingChatEditsFailedContextKey, CHAT_EDITING_MULTI_DIFF_SOURCE_RESOLVER_SCHEME, chatEditingMaxFileAssignmentName, chatEditingResourceContextKey, ChatEditingSessionState, decidedChatEditingResourceContextKey, defaultChatEditingMaxFileLimit, hasAppliedChatEditsContextKey, hasUndecidedChatEditingResourceContextKey, IChatEditingService, IChatEditingSession, IChatEditingSessionStream, IChatRelatedFile, IChatRelatedFilesProvider, IModifiedFileEntry, inChatEditingSessionContextKey, WorkingSetEntryState } from '../../common/chatEditingService.js';
-import { IChatResponseModel, IChatTextEditGroup } from '../../common/chatModel.js';
-import { IChatService } from '../../common/chatService.js';
+import { IChatNotebookEditGroup, IChatResponseModel, IChatTextEditGroup } from '../../common/chatModel.js';
+import { ICellEditReplaceOperation, IChatService } from '../../common/chatService.js';
 import { ChatEditingSession } from './chatEditingSession.js';
 import { ChatEditingSnapshotTextModelContentProvider, ChatEditingTextModelContentProvider } from './chatEditingTextModelContentProviders.js';
 
@@ -238,7 +238,7 @@ export class ChatEditingService extends Disposable implements IChatEditingServic
 
 		const observerDisposables = new DisposableStore();
 
-		let editsSource: AsyncIterableSource<IChatTextEditGroup> | undefined;
+		let editsSource: AsyncIterableSource<IChatTextEditGroup | IChatNotebookEditGroup> | undefined;
 		let editsPromise: Promise<void> | undefined;
 		const editsSeen = new ResourceMap<{ seen: number }>();
 		const editedFilesExist = new ResourceMap<Promise<boolean>>();
@@ -259,7 +259,7 @@ export class ChatEditingService extends Disposable implements IChatEditingServic
 
 		const handleResponseParts = async (responseModel: IChatResponseModel) => {
 			for (const part of responseModel.response.value) {
-				if (part.kind === 'codeblockUri' || part.kind === 'textEditGroup') {
+				if (part.kind === 'codeblockUri' || part.kind === 'textEditGroup' || part.kind === 'notebookEditGroup') {
 					// ensure editor is open asap
 					if (!editedFilesExist.get(part.uri)) {
 						editedFilesExist.set(part.uri, this._fileService.exists(part.uri).then((e) => {
@@ -278,12 +278,19 @@ export class ChatEditingService extends Disposable implements IChatEditingServic
 						editsSeen.set(part.uri, entry);
 					}
 
-					const allEdits: TextEdit[][] = part.kind === 'textEditGroup' ? part.edits : [];
-					const newEdits = allEdits.slice(entry.seen);
-					entry.seen += newEdits.length;
-
 					editsSource ??= new AsyncIterableSource();
-					editsSource.emitOne({ uri: part.uri, edits: newEdits, kind: 'textEditGroup', done: part.kind === 'textEditGroup' && part.done });
+
+					if (part.kind === 'notebookEditGroup') {
+						const allEdits: ICellEditReplaceOperation[][] = part.edits;
+						const newEdits = allEdits.slice(entry.seen);
+						entry.seen += newEdits.length;
+						editsSource.emitOne({ uri: part.uri, edits: newEdits, kind: 'notebookEditGroup' });
+					} else {
+						const allEdits: TextEdit[][] = part.kind === 'textEditGroup' ? part.edits : [];
+						const newEdits = allEdits.slice(entry.seen);
+						entry.seen += newEdits.length;
+						editsSource.emitOne({ uri: part.uri, edits: newEdits, kind: 'textEditGroup', done: part.kind === 'textEditGroup' && part.done });
+					}
 
 					if (first) {
 
@@ -295,9 +302,14 @@ export class ChatEditingService extends Disposable implements IChatEditingServic
 									break;
 								}
 								for (let i = 0; i < item.edits.length; i++) {
-									const group = item.edits[i];
 									const isLastGroup = i === item.edits.length - 1;
-									builder.textEdits(item.uri, group, isLastGroup && (item.done ?? false), responseModel);
+									if (item.kind === 'notebookEditGroup') {
+										const group = item.edits[i];
+										builder.notebookEdits(item.uri, group, false, responseModel);
+									} else {
+										const group = item.edits[i];
+										builder.textEdits(item.uri, group, isLastGroup && (item.done ?? false), responseModel);
+									}
 								}
 							}
 						}).finally(() => {
@@ -343,6 +355,9 @@ export class ChatEditingService extends Disposable implements IChatEditingServic
 		const stream: IChatEditingSessionStream = {
 			textEdits: (resource: URI, textEdits: TextEdit[], isDone: boolean, responseModel: IChatResponseModel) => {
 				session.acceptTextEdits(resource, textEdits, isDone, responseModel);
+			},
+			notebookEdits: (resource: URI, edits: ICellEditReplaceOperation[], isDone: boolean, responseModel: IChatResponseModel) => {
+				session.acceptNotebookEdits(resource, edits, isDone, responseModel);
 			}
 		};
 		session.acceptStreamingEditsStart();