--

Notebook Edit Support

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

@@ -5,6 +5,7 @@
 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 { extHostNamedCustomer, IExtHostContext } from '../../services/extensions/common/extHostCustomers.js';
 import { ExtHostCodeMapperShape, ExtHostContext, ICodeMapperProgressDto, ICodeMapperRequestDto, MainContext, MainThreadCodeMapperShape } from '../common/extHost.protocol.js';
@@ -54,9 +55,25 @@ 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 {
+				response.notebookEdit(resource, edits);
+			}
+		} 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

@@ -1296,7 +1296,12 @@ export interface ICodeMapperTextEdit {
 	edits: languages.TextEdit[];
 }
 
-export type ICodeMapperProgressDto = Dto<ICodeMapperTextEdit>;
+export interface ICodeMapperNotebookEdit {
+	uri: URI;
+	edits: notebookCommon.ICellEditOperation[];
+}
+
+export type ICodeMapperProgressDto = Dto<ICodeMapperTextEdit> | Dto<ICodeMapperNotebookEdit>;
 
 export interface MainThreadCodeMapperShape extends IDisposable {
 	$registerCodeMapperProvider(handle: number): void;
@@ -2132,16 +2137,17 @@ export interface IWorkspaceEditEntryMetadataDto {
 	iconPath?: { id: string } | UriComponents | { light: UriComponents; dark: UriComponents };
 }
 
+export type ICellEditReplaceOperation = {
+	editType: notebookCommon.CellEditType.Replace;
+	index: number;
+	count: number;
+	cells: NotebookCellDataDto[];
+};
 
 export type ICellEditOperationDto =
 	notebookCommon.ICellMetadataEdit
 	| notebookCommon.IDocumentMetadataEdit
-	| {
-		editType: notebookCommon.CellEditType.Replace;
-		index: number;
-		count: number;
-		cells: NotebookCellDataDto[];
-	};
+	| ICellEditReplaceOperation;
 
 export type IWorkspaceCellEditDto = Dto<Omit<notebookCommon.IWorkspaceNotebookCellEdit, 'cellEdit'>> & { cellEdit: ICellEditOperationDto };
 

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,10 @@ 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';
+import { NotebookDto } from '../browser/mainThreadNotebookDto.js';
 
 export class ExtHostCodeMapper implements extHostProtocol.ExtHostCodeMapperShape {
 
@@ -38,6 +40,18 @@ 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.from).filter(isDefined).map(e => ({
+						editType: e.editType,
+						index: e.index,
+						count: e.count,
+						cells: e.cells.map(NotebookDto.fromNotebookCellDataDto)
+					}))
+				});
 			}
 		};
 

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

@@ -38,7 +38,7 @@ import { DEFAULT_EDITOR_ASSOCIATION, SaveReason } from '../../common/editor.js';
 import { IViewBadge } from '../../common/views.js';
 import { ChatAgentLocation, IChatAgentRequest, IChatAgentResult } from '../../contrib/chat/common/chatAgents.js';
 import { IChatRequestVariableEntry } from '../../contrib/chat/common/chatModel.js';
-import { IChatAgentDetection, IChatAgentMarkdownContentWithVulnerability, IChatCodeCitation, IChatCommandButton, IChatConfirmation, IChatContentInlineReference, IChatContentReference, IChatFollowup, IChatMarkdownContent, IChatMoveMessage, IChatProgressMessage, IChatResponseCodeblockUriPart, IChatTaskDto, IChatTaskResult, IChatTextEdit, IChatTreeData, IChatUserActionEvent, IChatWarningMessage } from '../../contrib/chat/common/chatService.js';
+import { IChatAgentDetection, IChatAgentMarkdownContentWithVulnerability, IChatCodeCitation, IChatCommandButton, IChatConfirmation, IChatContentInlineReference, IChatContentReference, IChatFollowup, IChatMarkdownContent, IChatMoveMessage, IChatNotebookEdit, IChatProgressMessage, IChatResponseCodeblockUriPart, IChatTaskDto, IChatTaskResult, IChatTextEdit, IChatTreeData, IChatUserActionEvent, IChatWarningMessage } from '../../contrib/chat/common/chatService.js';
 import { IToolData, IToolResult } from '../../contrib/chat/common/languageModelToolsService.js';
 import * as chatProvider from '../../contrib/chat/common/languageModels.js';
 import { DebugTreeItemCollapsibleState, IDebugVisualizationTreeItem } from '../../contrib/debug/common/debug.js';
@@ -58,6 +58,8 @@ 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';
+import { ICellEditReplaceOperation } from './extHost.protocol.js';
 
 export namespace Command {
 
@@ -2653,6 +2655,43 @@ export namespace ChatResponseTextEditPart {
 
 }
 
+
+export namespace NotebookEdit {
+	export function from(edit: vscode.NotebookEdit): Dto<ICellEditReplaceOperation | 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 to(edit: Dto<ICellEditReplaceOperation>): 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): Dto<IChatNotebookEdit> {
+		return {
+			kind: 'notebookEdit',
+			uri: part.uri,
+			// We are only interested in cell replaces (insertions, deletions, replacements)
+			edits: part.edits.map(e => NotebookEdit.from(e)).filter(isDefined)
+		};
+	}
+
+	export function to(part: Dto<IChatNotebookEdit>): vscode.ChatResponseNotebookEditPart {
+		return new types.ChatResponseNotebookEditPart(URI.revive(part.uri), part.edits.map(NotebookEdit.to));
+	}
+
+}
+
 export namespace ChatResponseReferencePart {
 	export function from(part: types.ChatResponseReferencePart): Dto<IChatContentReference> {
 		const iconPath = ThemeIcon.isThemeIcon(part.iconPath) ? part.iconPath
@@ -2728,6 +2767,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

@@ -24,6 +24,7 @@ import { IFileService } from '../../../../../platform/files/common/files.js';
 import { IInstantiationService } from '../../../../../platform/instantiation/common/instantiation.js';
 import { bindContextKey } from '../../../../../platform/observable/common/platformObservableUtils.js';
 import { IProgressService, ProgressLocation } from '../../../../../platform/progress/common/progress.js';
+import { ICellEditReplaceOperation } from '../../../../api/common/extHost.protocol.js';
 import { EditorInput } from '../../../../common/editor/editorInput.js';
 import { IWorkbenchAssignmentService } from '../../../../services/assignment/common/assignmentService.js';
 import { IDecorationData, IDecorationsProvider, IDecorationsService } from '../../../../services/decorations/common/decorations.js';
@@ -36,7 +37,7 @@ 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 { IChatNotebookEditGroup, IChatResponseModel, IChatTextEditGroup } from '../../common/chatModel.js';
 import { IChatService } from '../../common/chatService.js';
 import { ChatEditingSession } from './chatEditingSession.js';
 import { ChatEditingSnapshotTextModelContentProvider, ChatEditingTextModelContentProvider } from './chatEditingTextModelContentProviders.js';
@@ -250,7 +251,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>>();
@@ -271,7 +272,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) => {
@@ -290,12 +291,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) {
 
@@ -307,9 +315,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);
+									}
 								}
 							}
 						}, { silent: true }).finally(() => {
@@ -370,6 +383,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();