WIP

TODO.md

@@ -0,0 +1,30 @@
+TODO:
+* When updating the diff we get errors about model having been disposed
+(possible user is scrolling throught the notebook, and the cells move out hence editos is disposed)
+* We might want to hold a ref to all text models of all cells, so that we can update the diff
+Do this in chat controller or modified file entry, when opening notebook, ensure all cell are resolved and we get the
+text ref models of all cells (original and modified)
+* Make changes to a notebook to delete some cells
+Scroll those deleted cells into view.
+Undo/Redo chat session edits.
+Notice the errors
+
+0.17/out/client_renderer/preload.js.map'
+2
+  ERR Unable to resolve text model content for resource vscode-notebook-cell:/Users/donjayamanne/demo/pyarrow_data/sample.ipynb#X11sZmlsZQ%3D%3D: Error: Unable to resolve text model content for resource vscode-notebook-cell:/Users/donjayamanne/demo/pyarrow_data/sample.ipynb#X11sZmlsZQ%3D%3D
+    at ResourceModelCollection.resolveTextModelContent (vscode-file://vscode-app/Users/donjayamanne/Development/vsc/vscode/out/vs/workbench/services/textmodelResolver/common/textModelResolverService.js:167:15)
+    at async ResourceModelCollection.doCreateReferencedObject (vscode-file://vscode-app/Users/donjayamanne/Development/vsc/vscode/out/vs/workbench/services/textmodelResolver/common/textModelResolverService.js:73:13)
+log.ts:445
+  ERR Model is disposed!: Error: Model is disposed!
+    at TextModel._assertNotDisposed (vscode-file://vscode-app/Users/donjayamanne/Development/vsc/vscode/out/vs/editor/common/model/textModel.js:296:19)
+    at TextModel.getVersionId (vscode-file://vscode-app/Users/donjayamanne/Development/vsc/vscode/out/vs/editor/common/model/textModel.js:509:14)
+    at ChatEditingModifiedNotebookFileEntry._updateDiffInfo (vscode-file://vscode-app/Users/donjayamanne/Development/vsc/vscode/out/vs/workbench/contrib/chat/browser/chatEditing/chatEditingModifiedNotebookFileEntry.js:450:52)
+    at async vscode-file://vscode-app/Users/donjayamanne/Development/vsc/vscode/out/vs/workbench/contrib/chat/browser/chatEditing/chatEditingModifiedNotebookFileEntry.js:432:21
+log.ts:445
+  ERR Unable to resolve text model content for resource vscode-notebook-cell:/Users/donjayamanne/demo/pyarrow_data/sample.ipynb#X11sZmlsZQ%3D%3D: Error: Unable to resolve text model content for resource vscode-notebook-cell:/Users/donjayamanne/demo/pyarrow_data/sample.ipynb#X11sZmlsZQ%3D%3D
+    at ResourceModelCollection.resolveTextModelContent (vscode-file://vscode-app/Users/donjayamanne/Development/vsc/vscode/out/vs/workbench/services/textmodelResolver/common/textModelResolverService.js:167:15)
+    at async ResourceModelCollection.doCreateReferencedObject (vscode-file://vscode-app/Users/donjayamanne/Development/vsc/vscode/out/vs/workbench/services/textmodelResolver/common/textModelResolverService.js:73:13)
+    at async AsyncReferenceCollection.acquire (vscode-file://vscode-app/Users/donjayamanne/Development/vsc/vscode/out/vs/base/common/lifecycle.js:531:28)
+    at async TextModelResolverService.createModelReference (vscode-file://vscode-app/Users/donjayamanne/Development/vsc/vscode/out/vs/workbench/services/textmodelResolver/common/textModelResolverService.js:205:16)
+    at async WordHighlighter._run (vscode-file://vscode-app/Users/donjayamanne/Development/vsc/vscode/out/vs/editor/contrib/wordHighlighter/browser/wordHighlighter.js:589:35)
+

extensions/ipynb/src/notebookModelStoreSync.ts

@@ -123,6 +123,7 @@ function trackAndUpdateCellMetadata(notebook: NotebookDocument, updates: { cell:
 	promise.then(clean, clean);
 }
 
+const pendingCellUpdates = new WeakSet<NotebookCell>();
 function onDidChangeNotebookCells(e: NotebookDocumentChangeEventEx) {
 	if (!isSupportedNotebook(e.notebook)) {
 		return;
@@ -155,11 +156,15 @@ function onDidChangeNotebookCells(e: NotebookDocumentChangeEventEx) {
 			// Thus this is a change in metadata, which we will need to update in the model.
 			metadata.execution_count = null;
 			metadataUpdated = true;
+			// Note: We will get another event for this, see below for the check.
+			// track the fact that we're expecting an update for this cell.
+			pendingCellUpdates.add(e.cell);
 		} else if ((!e.executionSummary || (!e.executionSummary?.executionOrder && !e.executionSummary?.success && !e.executionSummary?.timing))
-			&& !e.metadata && !e.outputs && currentMetadata.execution_count) {
-			// This is a result of the previous cell being cleared.
+			&& !e.metadata && !e.outputs && currentMetadata.execution_count && pendingCellUpdates.has(e.cell)) {
+			// This is a result of the cell being cleared (i.e. we perfomed an update request and this is now the update event).
 			metadata.execution_count = null;
 			metadataUpdated = true;
+			pendingCellUpdates.delete(e.cell);
 		}
 
 		if (e.document?.languageId && e.document?.languageId !== preferredCellLanguage && e.document?.languageId !== languageIdInMetadata) {

src/vs/base/common/buffer.ts

@@ -288,6 +288,10 @@ export interface VSBufferReadable extends streams.Readable<VSBuffer> { }
 
 export interface VSBufferReadableStream extends streams.ReadableStream<VSBuffer> { }
 
+export function bufferToReadableStream(buffer: VSBuffer): VSBufferReadableStream {
+	return streams.toStream<VSBuffer>(buffer, chunks => VSBuffer.concat(chunks));
+}
+
 export interface VSBufferWriteableStream extends streams.WriteableStream<VSBuffer> { }
 
 export interface VSBufferReadableBufferedStream extends streams.ReadableBufferedStream<VSBuffer> { }

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,37 @@ 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,
+			done: part.done,
+			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',
+			done: part.done,
+			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

@@ -1775,6 +1775,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,19 @@ export interface IWorkspaceEditEntryMetadataDto {
 	iconPath?: { id: string } | UriComponents | { light: UriComponents; dark: UriComponents };
 }
 
+export interface IChatNotebookEditDto {
+	uri: URI;
+	edits: ICellEditReplaceOperationDto[];
+	kind: 'notebookEdit';
+	done?: boolean;
+}
+
+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 {
 
@@ -2675,6 +2676,48 @@ 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),
+			done: part.isDone
+		};
+	}
+
+	export function to(part: extHostProtocol.IChatNotebookEditDto): vscode.ChatResponseNotebookEditPart {
+		if (part.done) {
+			return new types.ChatResponseNotebookEditPart(URI.revive(part.uri), true);
+		} else {
+			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
@@ -2750,6 +2793,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

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