Merge pull request #102760 from microsoft/joh/celldocs

Clean up and simplify cell text document world

extensions/vscode-notebook-tests/src/notebook.test.ts

@@ -131,6 +131,46 @@ suite('Notebook API tests', () => {
 		await firstDocumentClose;
 	});
 
+	test('notebook open/close, all cell-documents are ready', async function () {
+		const resource = await createRandomFile('', undefined, 'first', '.vsctestnb');
+
+		const p = getEventOncePromise(vscode.notebook.onDidOpenNotebookDocument).then(notebook => {
+			for (let cell of notebook.cells) {
+				const doc = vscode.workspace.textDocuments.find(doc => doc.uri.toString() === cell.uri.toString());
+				assert.ok(doc);
+				assert.strictEqual(doc === cell.document, true);
+				assert.strictEqual(doc?.languageId, cell.language);
+				assert.strictEqual(doc?.isDirty, false);
+				assert.strictEqual(doc?.isClosed, false);
+			}
+		});
+
+		await vscode.commands.executeCommand('vscode.openWith', resource, 'notebookCoreTest');
+		await p;
+		await vscode.commands.executeCommand('workbench.action.closeAllEditors');
+	});
+
+	test('notebook open/close, notebook ready when cell-document open event is fired', async function () {
+		const resource = await createRandomFile('', undefined, 'first', '.vsctestnb');
+		let didHappen = false;
+		const p = getEventOncePromise(vscode.workspace.onDidOpenTextDocument).then(doc => {
+			if (doc.uri.scheme !== 'vscode-notebook-cell') {
+				return;
+			}
+			const notebook = vscode.notebook.notebookDocuments.find(notebook => {
+				const cell = notebook.cells.find(cell => cell.document === doc);
+				return Boolean(cell);
+			});
+			assert.ok(notebook, `notebook for cell ${doc.uri} NOT found`);
+			didHappen = true;
+		});
+
+		await vscode.commands.executeCommand('vscode.openWith', resource, 'notebookCoreTest');
+		await p;
+		assert.strictEqual(didHappen, true);
+		await vscode.commands.executeCommand('workbench.action.closeAllEditors');
+	});
+
 	test('shared document in notebook editors', async function () {
 		assertInitalState();
 

src/vs/base/common/network.ts

@@ -58,6 +58,8 @@ export namespace Schemas {
 
 	export const vscodeNotebook = 'vscode-notebook';
 
+	export const vscodeNotebookCell = 'vscode-notebook-cell';
+
 	export const vscodeSettings = 'vscode-settings';
 
 	export const webviewPanel = 'webview-panel';

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

@@ -203,7 +203,7 @@ export class MainThreadNotebooks extends Disposable implements MainThreadNoteboo
 
 	async removeNotebookTextModel(uri: URI): Promise<void> {
 		// TODO@rebornix, remove cell should use emitDelta as well to ensure document/editor events are sent together
-		await this._proxy.$acceptDocumentAndEditorsDelta({ removedDocuments: [uri] });
+		this._proxy.$acceptDocumentAndEditorsDelta({ removedDocuments: [uri] });
 		let textModelDisposableStore = this._editorEventListenersMapping.get(uri.toString());
 		textModelDisposableStore?.dispose();
 		this._editorEventListenersMapping.delete(URI.from(uri).toString());

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

@@ -1633,7 +1633,7 @@ export interface ExtHostNotebookShape {
 	$acceptModelChanged(uriComponents: UriComponents, event: NotebookCellsChangedEvent): void;
 	$acceptModelSaved(uriComponents: UriComponents): void;
 	$acceptEditorPropertiesChanged(uriComponents: UriComponents, data: INotebookEditorPropertiesChangeData): void;
-	$acceptDocumentAndEditorsDelta(delta: INotebookDocumentsAndEditorsDelta): Promise<void>;
+	$acceptDocumentAndEditorsDelta(delta: INotebookDocumentsAndEditorsDelta): void;
 	$undoNotebook(viewType: string, uri: UriComponents, editId: number, isDirty: boolean): Promise<void>;
 	$redoNotebook(viewType: string, uri: UriComponents, editId: number, isDirty: boolean): Promise<void>;
 

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

@@ -53,7 +53,7 @@ export class ExtHostDocuments implements ExtHostDocumentsShape {
 	}
 
 	public getAllDocumentData(): ExtHostDocumentData[] {
-		return this._documentsAndEditors.allDocuments();
+		return [...this._documentsAndEditors.allDocuments()];
 	}
 
 	public getDocumentData(resource: vscode.Uri): ExtHostDocumentData | undefined {
@@ -69,8 +69,8 @@ export class ExtHostDocuments implements ExtHostDocumentsShape {
 
 	public getDocument(resource: vscode.Uri): vscode.TextDocument {
 		const data = this.getDocumentData(resource);
-		if (!data || !data.document) {
-			throw new Error('Unable to retrieve document from URI');
+		if (!data?.document) {
+			throw new Error(`Unable to retrieve document from URI '${resource}'`);
 		}
 		return data.document;
 	}

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

@@ -15,6 +15,19 @@ import { ExtHostTextEditor } from 'vs/workbench/api/common/extHostTextEditor';
 import * as typeConverters from 'vs/workbench/api/common/extHostTypeConverters';
 import { ILogService } from 'vs/platform/log/common/log';
 import { ResourceMap } from 'vs/base/common/map';
+import { Schemas } from 'vs/base/common/network';
+import { Iterable } from 'vs/base/common/iterator';
+
+class Reference<T> {
+	private _count = 0;
+	constructor(readonly value: T) { }
+	ref() {
+		this._count++;
+	}
+	unref() {
+		return --this._count === 0;
+	}
+}
 
 export class ExtHostDocumentsAndEditors implements ExtHostDocumentsAndEditorsShape {
 
@@ -23,7 +36,7 @@ export class ExtHostDocumentsAndEditors implements ExtHostDocumentsAndEditorsSha
 	private _activeEditorId: string | null = null;
 
 	private readonly _editors = new Map<string, ExtHostTextEditor>();
-	private readonly _documents = new ResourceMap<ExtHostDocumentData>();
+	private readonly _documents = new ResourceMap<Reference<ExtHostDocumentData>>();
 
 	private readonly _onDidAddDocuments = new Emitter<ExtHostDocumentData[]>();
 	private readonly _onDidRemoveDocuments = new Emitter<ExtHostDocumentData[]>();
@@ -50,9 +63,9 @@ export class ExtHostDocumentsAndEditors implements ExtHostDocumentsAndEditorsSha
 			for (const uriComponent of delta.removedDocuments) {
 				const uri = URI.revive(uriComponent);
 				const data = this._documents.get(uri);
-				this._documents.delete(uri);
-				if (data) {
-					removedDocuments.push(data);
+				if (data?.unref()) {
+					this._documents.delete(uri);
+					removedDocuments.push(data.value);
 				}
 			}
 		}
@@ -60,19 +73,30 @@ export class ExtHostDocumentsAndEditors implements ExtHostDocumentsAndEditorsSha
 		if (delta.addedDocuments) {
 			for (const data of delta.addedDocuments) {
 				const resource = URI.revive(data.uri);
-				assert.ok(!this._documents.has(resource), `document '${resource} already exists!'`);
-
-				const documentData = new ExtHostDocumentData(
-					this._extHostRpc.getProxy(MainContext.MainThreadDocuments),
-					resource,
-					data.lines,
-					data.EOL,
-					data.modeId,
-					data.versionId,
-					data.isDirty
-				);
-				this._documents.set(resource, documentData);
-				addedDocuments.push(documentData);
+				let ref = this._documents.get(resource);
+
+				// double check -> only notebook cell documents should be
+				// referenced/opened more than once...
+				if (ref) {
+					if (resource.scheme !== Schemas.vscodeNotebookCell) {
+						throw new Error(`document '${resource} already exists!'`);
+					}
+				}
+				if (!ref) {
+					ref = new Reference(new ExtHostDocumentData(
+						this._extHostRpc.getProxy(MainContext.MainThreadDocuments),
+						resource,
+						data.lines,
+						data.EOL,
+						data.modeId,
+						data.versionId,
+						data.isDirty
+					));
+					this._documents.set(resource, ref);
+					addedDocuments.push(ref.value);
+				}
+
+				ref.ref();
 			}
 		}
 
@@ -92,7 +116,7 @@ export class ExtHostDocumentsAndEditors implements ExtHostDocumentsAndEditorsSha
 				assert.ok(this._documents.has(resource), `document '${resource}' does not exist`);
 				assert.ok(!this._editors.has(data.id), `editor '${data.id}' already exists!`);
 
-				const documentData = this._documents.get(resource)!;
+				const documentData = this._documents.get(resource)!.value;
 				const editor = new ExtHostTextEditor(
 					data.id,
 					this._extHostRpc.getProxy(MainContext.MainThreadTextEditors),
@@ -132,11 +156,11 @@ export class ExtHostDocumentsAndEditors implements ExtHostDocumentsAndEditorsSha
 	}
 
 	getDocument(uri: URI): ExtHostDocumentData | undefined {
-		return this._documents.get(uri);
+		return this._documents.get(uri)?.value;
 	}
 
-	allDocuments(): ExtHostDocumentData[] {
-		return [...this._documents.values()];
+	allDocuments(): Iterable<ExtHostDocumentData> {
+		return Iterable.map(this._documents.values(), ref => ref.value);
 	}
 
 	getEditor(id: string): ExtHostTextEditor | undefined {

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

@@ -7,18 +7,16 @@ import { CancellationToken, CancellationTokenSource } from 'vs/base/common/cance
 import { readonly } from 'vs/base/common/errors';
 import { Emitter, Event } from 'vs/base/common/event';
 import { hash } from 'vs/base/common/hash';
-import { Disposable, DisposableStore, IDisposable } from 'vs/base/common/lifecycle';
+import { Disposable, DisposableStore, dispose, IDisposable } from 'vs/base/common/lifecycle';
 import { Schemas } from 'vs/base/common/network';
 import { joinPath } from 'vs/base/common/resources';
 import { ISplice } from 'vs/base/common/sequence';
-import { NotImplementedProxy } from 'vs/base/common/types';
 import { URI, UriComponents } from 'vs/base/common/uri';
 import * as UUID from 'vs/base/common/uuid';
 import { IExtensionDescription } from 'vs/platform/extensions/common/extensions';
+import { CellKind, ExtHostNotebookShape, IMainContext, IModelAddedData, INotebookDocumentsAndEditorsDelta, INotebookEditorPropertiesChangeData, MainContext, MainThreadNotebookShape, NotebookCellOutputsSplice } from 'vs/workbench/api/common/extHost.protocol';
 import { ILogService } from 'vs/platform/log/common/log';
-import { CellKind, ExtHostNotebookShape, IMainContext, INotebookDocumentsAndEditorsDelta, INotebookEditorPropertiesChangeData, MainContext, MainThreadDocumentsShape, MainThreadNotebookShape, NotebookCellOutputsSplice } from 'vs/workbench/api/common/extHost.protocol';
 import { ExtHostCommands } from 'vs/workbench/api/common/extHostCommands';
-import { ExtHostDocumentData } from 'vs/workbench/api/common/extHostDocumentData';
 import { ExtHostDocumentsAndEditors } from 'vs/workbench/api/common/extHostDocumentsAndEditors';
 import { IExtensionStoragePaths } from 'vs/workbench/api/common/extHostStoragePaths';
 import * as typeConverters from 'vs/workbench/api/common/extHostTypeConverters';
@@ -28,6 +26,7 @@ import { CellEditType, CellOutputKind, diff, ICellDeleteEdit, ICellEditOperation
 import * as vscode from 'vscode';
 import { Cache } from './cache';
 
+
 interface IObservable<T> {
 	proxy: T;
 	onDidChange: Event<void>;
@@ -59,27 +58,18 @@ interface INotebookEventEmitter {
 const addIdToOutput = (output: IRawOutput, id = UUID.generateUuid()): IProcessedOutput => output.outputKind === CellOutputKind.Rich
 	? ({ ...output, outputId: id }) : output;
 
-class DettachedCellDocumentData extends ExtHostDocumentData {
-
-	private static readonly _fakeProxy = new class extends NotImplementedProxy<MainThreadDocumentsShape>('document') {
-		$trySaveDocument() {
-			return Promise.reject('Cell-document cannot be saved');
-		}
-	};
+export class ExtHostCell extends Disposable implements vscode.NotebookCell {
 
-	constructor(cell: IMainCellDto) {
-		super(DettachedCellDocumentData._fakeProxy,
-			URI.revive(cell.uri),
-			cell.source,
-			cell.eol,
-			cell.language,
-			0,
-			false
-		);
+	public static asModelAddData(cell: IMainCellDto): IModelAddedData {
+		return {
+			EOL: cell.eol,
+			lines: cell.source,
+			modeId: cell.language,
+			uri: cell.uri,
+			isDirty: false,
+			versionId: 1
+		};
 	}
-}
-
-export class ExtHostCell extends Disposable implements vscode.NotebookCell {
 
 	private _onDidChangeOutputs = new Emitter<ISplice<IProcessedOutput>[]>();
 	readonly onDidChangeOutputs: Event<ISplice<IProcessedOutput>[]> = this._onDidChangeOutputs.event;
@@ -94,13 +84,6 @@ export class ExtHostCell extends Disposable implements vscode.NotebookCell {
 	readonly uri: URI;
 	readonly cellKind: CellKind;
 
-	// todo@jrieken this is a little fish because we have
-	// vscode.TextDocument for which we never fired an onDidOpen
-	// event and which doesn't appear in the list of documents.
-	// this will change once the "real" document comes along. We
-	// should come up with a better approach here...
-	readonly defaultDocument: DettachedCellDocumentData;
-
 	constructor(
 		private _proxy: MainThreadNotebookShape,
 		readonly notebook: ExtHostNotebookDocument,
@@ -112,7 +95,6 @@ export class ExtHostCell extends Disposable implements vscode.NotebookCell {
 		this.handle = cell.handle;
 		this.uri = URI.revive(cell.uri);
 		this.cellKind = cell.cellKind;
-		this.defaultDocument = new DettachedCellDocumentData(cell);
 
 		this._outputs = cell.outputs;
 		for (const output of this._outputs) {
@@ -128,7 +110,7 @@ export class ExtHostCell extends Disposable implements vscode.NotebookCell {
 	}
 
 	get document(): vscode.TextDocument {
-		return this._extHostDocument.getDocument(this.uri)?.document ?? this.defaultDocument.document;
+		return this._extHostDocument.getDocument(this.uri)!.document;
 	}
 
 	get language(): string {
@@ -348,7 +330,7 @@ export class ExtHostNotebookDocument extends Disposable implements vscode.Notebo
 	dispose() {
 		this._disposed = true;
 		super.dispose();
-		this._cellDisposableMapping.forEach(cell => cell.dispose());
+		dispose(this._cellDisposableMapping.values());
 	}
 
 	get fileName() { return this.uri.fsPath; }
@@ -380,6 +362,7 @@ export class ExtHostNotebookDocument extends Disposable implements vscode.Notebo
 		}
 
 		const contentChangeEvents: vscode.NotebookCellsChangeData[] = [];
+		const addedCellDocuments: IModelAddedData[] = [];
 
 		splices.reverse().forEach(splice => {
 			const cellDtos = splice[2];
@@ -387,6 +370,10 @@ export class ExtHostNotebookDocument extends Disposable implements vscode.Notebo
 
 				const extCell = new ExtHostCell(this._proxy, this, this._documentsAndEditors, cell);
 
+				if (!initialization) {
+					addedCellDocuments.push(ExtHostCell.asModelAddData(cell));
+				}
+
 				if (!this._cellDisposableMapping.has(extCell.handle)) {
 					this._cellDisposableMapping.set(extCell.handle, new DisposableStore());
 				}
@@ -403,21 +390,22 @@ export class ExtHostNotebookDocument extends Disposable implements vscode.Notebo
 			for (let j = splice[0]; j < splice[0] + splice[1]; j++) {
 				this._cellDisposableMapping.get(this.cells[j].handle)?.dispose();
 				this._cellDisposableMapping.delete(this.cells[j].handle);
-
 			}
 
 			const deletedItems = this.cells.splice(splice[0], splice[1], ...newCells);
 
-			const event: vscode.NotebookCellsChangeData = {
+			contentChangeEvents.push({
 				start: splice[0],
 				deletedCount: splice[1],
 				deletedItems,
 				items: newCells
-			};
-
-			contentChangeEvents.push(event);
+			});
 		});
 
+		if (addedCellDocuments) {
+			this._documentsAndEditors.$acceptDocumentsAndEditorsDelta({ addedDocuments: addedCellDocuments });
+		}
+
 		if (!initialization) {
 			this._emitter.emitModelChange({
 				document: this,
@@ -467,7 +455,6 @@ export class ExtHostNotebookDocument extends Disposable implements vscode.Notebo
 
 	private $changeCellLanguage(index: number, language: string): void {
 		const cell = this.cells[index];
-		cell.defaultDocument._acceptLanguageId(language);
 		const event: vscode.NotebookCellLanguageChangeEvent = { document: this, cell, language };
 		this._emitter.emitCellLanguageChange(event);
 	}
@@ -1526,7 +1513,7 @@ export class ExtHostNotebookController implements ExtHostNotebookShape, ExtHostN
 		this._editors.set(editorId, { editor });
 	}
 
-	async $acceptDocumentAndEditorsDelta(delta: INotebookDocumentsAndEditorsDelta) {
+	$acceptDocumentAndEditorsDelta(delta: INotebookDocumentsAndEditorsDelta): void {
 		let editorChanged = false;
 
 		if (delta.removedDocuments) {
@@ -1538,6 +1525,7 @@ export class ExtHostNotebookController implements ExtHostNotebookShape, ExtHostN
 				if (document) {
 					document.dispose();
 					this._documents.delete(revivedUriStr);
+					this._documentsAndEditors.$acceptDocumentsAndEditorsDelta({ removedDocuments: document.cells.map(cell => cell.uri) });
 					this._onDidCloseNotebookDocument.fire(document);
 				}
 
@@ -1552,6 +1540,9 @@ export class ExtHostNotebookController implements ExtHostNotebookShape, ExtHostN
 		}
 
 		if (delta.addedDocuments) {
+
+			const addedCellDocuments: IModelAddedData[] = [];
+
 			delta.addedDocuments.forEach(modelData => {
 				const revivedUri = URI.revive(modelData.uri);
 				const revivedUriStr = revivedUri.toString();
@@ -1598,6 +1589,9 @@ export class ExtHostNotebookController implements ExtHostNotebookShape, ExtHostN
 						]]
 					});
 
+					// add cell document as vscode.TextDocument
+					addedCellDocuments.push(...modelData.cells.map(ExtHostCell.asModelAddData));
+
 					this._documents.get(revivedUriStr)?.dispose();
 					this._documents.set(revivedUriStr, document);
 
@@ -1608,6 +1602,10 @@ export class ExtHostNotebookController implements ExtHostNotebookShape, ExtHostN
 					}
 				}
 
+				this._documentsAndEditors.$acceptDocumentsAndEditorsDelta({
+					addedDocuments: addedCellDocuments
+				});
+
 				const document = this._documents.get(revivedUriStr)!;
 				this._onDidOpenNotebookDocument.fire(document);
 			});

src/vs/workbench/contrib/notebook/common/notebookCommon.ts

@@ -425,7 +425,7 @@ export function getCellUndoRedoComparisonKey(uri: URI) {
 
 export namespace CellUri {
 
-	export const scheme = 'vscode-notebook-cell';
+	export const scheme = Schemas.vscodeNotebookCell;
 	const _regex = /^\d{7,}/;
 
 	export function generate(notebook: URI, handle: number): URI {

src/vs/workbench/test/browser/api/extHostNotebook.test.ts

@@ -4,148 +4,215 @@
  *--------------------------------------------------------------------------------------------*/
 
 import * as assert from 'assert';
+import * as vscode from 'vscode';
 import { ExtHostDocumentsAndEditors } from 'vs/workbench/api/common/extHostDocumentsAndEditors';
 import { TestRPCProtocol } from 'vs/workbench/test/browser/api/testRPCProtocol';
 import { DisposableStore } from 'vs/base/common/lifecycle';
 import { NullLogService } from 'vs/platform/log/common/log';
 import { mock } from 'vs/base/test/common/mock';
-import { MainThreadNotebookShape } from 'vs/workbench/api/common/extHost.protocol';
-import { ExtHostNotebookDocument, ExtHostCell } from 'vs/workbench/api/common/extHostNotebook';
-import { CellKind } from 'vs/workbench/contrib/notebook/common/notebookCommon';
+import { IModelAddedData, MainContext, MainThreadCommandsShape, MainThreadNotebookShape } from 'vs/workbench/api/common/extHost.protocol';
+import { ExtHostNotebookDocument, ExtHostNotebookController } from 'vs/workbench/api/common/extHostNotebook';
+import { CellKind, CellUri, NotebookCellsChangeType } from 'vs/workbench/contrib/notebook/common/notebookCommon';
 import { URI } from 'vs/base/common/uri';
 import { ExtHostDocuments } from 'vs/workbench/api/common/extHostDocuments';
+import { ExtHostCommands } from 'vs/workbench/api/common/extHostCommands';
+import { nullExtensionDescription } from 'vs/workbench/services/extensions/common/extensions';
+import { isEqual } from 'vs/base/common/resources';
+
+suite('NotebookCell#Document', function () {
 
-suite('NotebookCell', function () {
 
 	let rpcProtocol: TestRPCProtocol;
+	let notebook: ExtHostNotebookDocument;
 	let extHostDocumentsAndEditors: ExtHostDocumentsAndEditors;
-
+	let extHostDocuments: ExtHostDocuments;
+	let extHostNotebooks: ExtHostNotebookController;
+	const notebookUri = URI.parse('test:///notebook.file');
 	const disposables = new DisposableStore();
-	const fakeNotebookProxy = new class extends mock<MainThreadNotebookShape>() { };
-	const fakeNotebook = new class extends mock<ExtHostNotebookDocument>() { };
 
 	setup(async function () {
 		disposables.clear();
+
 		rpcProtocol = new TestRPCProtocol();
+		rpcProtocol.set(MainContext.MainThreadCommands, new class extends mock<MainThreadCommandsShape>() {
+			$registerCommand() { }
+		});
+		rpcProtocol.set(MainContext.MainThreadNotebook, new class extends mock<MainThreadNotebookShape>() {
+			async $registerNotebookProvider() { }
+			async $unregisterNotebookProvider() { }
+		});
 		extHostDocumentsAndEditors = new ExtHostDocumentsAndEditors(rpcProtocol, new NullLogService());
-	});
+		extHostDocuments = new ExtHostDocuments(rpcProtocol, extHostDocumentsAndEditors);
+		extHostNotebooks = new ExtHostNotebookController(rpcProtocol, new ExtHostCommands(rpcProtocol, new NullLogService()), extHostDocumentsAndEditors, { isExtensionDevelopmentDebug: false, webviewCspSource: '', webviewResourceRoot: '' }, new NullLogService());
+		let reg = extHostNotebooks.registerNotebookContentProvider(nullExtensionDescription, 'test', new class extends mock<vscode.NotebookContentProvider>() {
+			// async openNotebook() { }
+		});
+		extHostNotebooks.$acceptDocumentAndEditorsDelta({
+			addedDocuments: [{
+				handle: 0,
+				uri: notebookUri,
+				viewType: 'test',
+				versionId: 0,
+				cells: [{
+					handle: 0,
+					uri: CellUri.generate(notebookUri, 0),
+					source: ['### Heading'],
+					eol: '\n',
+					language: 'markdown',
+					cellKind: CellKind.Markdown,
+					outputs: [],
+				}, {
+					handle: 1,
+					uri: CellUri.generate(notebookUri, 1),
+					source: ['console.log("aaa")', 'console.log("bbb")'],
+					eol: '\n',
+					language: 'javascript',
+					cellKind: CellKind.Code,
+					outputs: [],
+				}],
+			}],
+			addedEditors: [{
+				documentUri: notebookUri,
+				id: '_notebook_editor_0',
+				selections: [0]
+			}]
+		});
+		extHostNotebooks.$acceptDocumentAndEditorsDelta({ newActiveEditor: '_notebook_editor_0' });
 
-	test('Document is real', function () {
-
-		const dto = {
-			cellKind: CellKind.Code,
-			eol: '\n',
-			source: ['aaaa', 'bbbb', 'cccc'],
-			handle: 0,
-			language: 'fooLang',
-			outputs: [],
-			uri: URI.parse('test:/path')
-		};
-		const cell = new ExtHostCell(fakeNotebookProxy, fakeNotebook, extHostDocumentsAndEditors, dto);
-
-		assert.ok(cell.document);
-		assert.strictEqual(cell.document.version, 0);
-		assert.strictEqual(cell.document.languageId, dto.language);
-		assert.strictEqual(cell.document.uri.toString(), dto.uri.toString());
-		assert.strictEqual(cell.uri.toString(), dto.uri.toString());
+		notebook = extHostNotebooks.notebookDocuments[0]!;
+
+		disposables.add(reg);
+		disposables.add(notebook);
+		disposables.add(extHostDocuments);
 	});
 
 
-	test('Document is uses actual document when possible', function () {
+	test('cell document is vscode.TextDocument', async function () {
 
-		const dto = {
-			cellKind: CellKind.Code,
-			eol: '\n',
-			source: ['aaaa', 'bbbb', 'cccc'],
-			handle: 0,
-			language: 'fooLang',
-			outputs: [],
-			uri: URI.parse('test:/path')
-		};
-		const cell = new ExtHostCell(fakeNotebookProxy, fakeNotebook, extHostDocumentsAndEditors, dto);
+		assert.strictEqual(notebook.cells.length, 2);
 
-		// this is the "default document" which is used when the real
-		// document isn't open
-		const documentNow = cell.document;
+		const [c1, c2] = notebook.cells;
+		const d1 = extHostDocuments.getDocument(c1.uri);
 
-		extHostDocumentsAndEditors.$acceptDocumentsAndEditorsDelta({
-			addedDocuments: [{
-				isDirty: false,
-				versionId: 12,
-				modeId: dto.language,
-				uri: dto.uri,
-				lines: dto.source,
-				EOL: dto.eol
-			}]
-		});
-
-		// the real document
-		assert.ok(documentNow !== cell.document);
-		assert.strictEqual(cell.document.languageId, dto.language);
-		assert.strictEqual(cell.document.uri.toString(), dto.uri.toString());
-		assert.strictEqual(cell.uri.toString(), dto.uri.toString());
+		assert.ok(d1);
+		assert.equal(d1.languageId, c1.language);
+		assert.equal(d1.version, 1);
 
-		// back to "default document"
-		extHostDocumentsAndEditors.$acceptDocumentsAndEditorsDelta({ removedDocuments: [dto.uri] });
-		assert.ok(documentNow === cell.document);
+		const d2 = extHostDocuments.getDocument(c2.uri);
+		assert.ok(d2);
+		assert.equal(d2.languageId, c2.language);
+		assert.equal(d2.version, 1);
 	});
 
-	test('Document can change language (1/2)', function () {
-
-		const dto = {
-			cellKind: CellKind.Code,
-			eol: '\n',
-			source: ['aaaa', 'bbbb', 'cccc'],
-			handle: 0,
-			language: 'fooLang',
-			outputs: [],
-			uri: URI.parse('test:/path')
-		};
-		const cell = new ExtHostCell(fakeNotebookProxy, fakeNotebook, extHostDocumentsAndEditors, dto);
-
-		assert.strictEqual(cell.document.languageId, dto.language);
-		cell.defaultDocument._acceptLanguageId('barLang');
-		assert.strictEqual(cell.document.languageId, 'barLang');
+	test('cell document goes when notebook closes', async function () {
+		const cellUris: string[] = [];
+		for (let cell of notebook.cells) {
+			assert.ok(extHostDocuments.getDocument(cell.uri));
+			cellUris.push(cell.uri.toString());
+		}
+
+		const removedCellUris: string[] = [];
+		const reg = extHostDocuments.onDidRemoveDocument(doc => {
+			removedCellUris.push(doc.uri.toString());
+		});
+
+		extHostNotebooks.$acceptDocumentAndEditorsDelta({ removedDocuments: [notebook.uri] });
+		reg.dispose();
+
+		assert.strictEqual(removedCellUris.length, 2);
+		assert.deepStrictEqual(removedCellUris.sort(), cellUris.sort());
 	});
 
+	test('cell document is vscode.TextDocument after changing it', async function () {
 
-	test('Document can change language (1/2)', function () {
+		const p = new Promise((resolve, reject) => {
+			extHostNotebooks.onDidChangeNotebookCells(e => {
+				try {
+					assert.strictEqual(e.changes.length, 1);
+					assert.strictEqual(e.changes[0].items.length, 2);
 
+					const [first, second] = e.changes[0].items;
 
-		const dto = {
-			cellKind: CellKind.Code,
-			eol: '\n',
-			source: ['aaaa', 'bbbb', 'cccc'],
-			handle: 0,
-			language: 'fooLang',
-			outputs: [],
-			uri: URI.parse('test:/path')
-		};
+					const doc1 = extHostDocuments.getAllDocumentData().find(data => isEqual(data.document.uri, first.uri));
+					assert.ok(doc1);
+					assert.strictEqual(doc1?.document === first.document, true);
 
-		extHostDocumentsAndEditors.$acceptDocumentsAndEditorsDelta({
-			addedDocuments: [{
-				isDirty: false,
-				versionId: 12,
-				modeId: dto.language,
-				uri: dto.uri,
-				lines: dto.source,
-				EOL: dto.eol
-			}]
-		});
+					const doc2 = extHostDocuments.getAllDocumentData().find(data => isEqual(data.document.uri, second.uri));
+					assert.ok(doc2);
+					assert.strictEqual(doc2?.document === second.document, true);
 
-		const extHostDocuments = new ExtHostDocuments(rpcProtocol, extHostDocumentsAndEditors);
+					resolve();
 
-		const cell = new ExtHostCell(fakeNotebookProxy, fakeNotebook, extHostDocumentsAndEditors, dto);
+				} catch (err) {
+					reject(err);
+				}
+			});
+		});
 
-		// a real document already exists and therefore
-		// the "default document" doesn't count
+		extHostNotebooks.$acceptModelChanged(notebookUri, {
+			kind: NotebookCellsChangeType.ModelChange,
+			versionId: notebook.versionId + 1,
+			changes: [[0, 0, [{
+				handle: 2,
+				uri: CellUri.generate(notebookUri, 2),
+				source: ['Hello', 'World', 'Hello World!'],
+				eol: '\n',
+				language: 'test',
+				cellKind: CellKind.Code,
+				outputs: [],
+			}, {
+				handle: 3,
+				uri: CellUri.generate(notebookUri, 3),
+				source: ['Hallo', 'Welt', 'Hallo Welt!'],
+				eol: '\n',
+				language: 'test',
+				cellKind: CellKind.Code,
+				outputs: [],
+			}]]]
+		});
 
-		assert.strictEqual(cell.document.languageId, dto.language);
-		cell.defaultDocument._acceptLanguageId('barLang');
-		assert.strictEqual(cell.document.languageId, dto.language);
+		await p;
 
-		extHostDocuments.$acceptModelModeChanged(dto.uri, dto.language, 'barLang');
-		assert.strictEqual(cell.document.languageId, 'barLang');
 	});
 
+	test('cell document stays open when notebook is still open', async function () {
+
+		const docs: vscode.TextDocument[] = [];
+		const addData: IModelAddedData[] = [];
+		for (let cell of notebook.cells) {
+			const doc = extHostDocuments.getDocument(cell.uri);
+			assert.ok(doc);
+			assert.equal(extHostDocuments.getDocument(cell.uri).isClosed, false);
+			docs.push(doc);
+			addData.push({
+				EOL: '\n',
+				isDirty: doc.isDirty,
+				lines: doc.getText().split('\n'),
+				modeId: doc.languageId,
+				uri: doc.uri,
+				versionId: doc.version
+			});
+		}
+
+		// this call happens when opening a document on the main side
+		extHostDocumentsAndEditors.$acceptDocumentsAndEditorsDelta({ addedDocuments: addData });
+
+		// this call happens when closing a document from the main side
+		extHostDocumentsAndEditors.$acceptDocumentsAndEditorsDelta({ removedDocuments: docs.map(d => d.uri) });
+
+		// notebook is still open -> cell documents stay open
+		for (let cell of notebook.cells) {
+			assert.ok(extHostDocuments.getDocument(cell.uri));
+			assert.equal(extHostDocuments.getDocument(cell.uri).isClosed, false);
+		}
+
+		// close notebook -> docs are closed
+		extHostNotebooks.$acceptDocumentAndEditorsDelta({ removedDocuments: [notebook.uri] });
+		for (let cell of notebook.cells) {
+			assert.throws(() => extHostDocuments.getDocument(cell.uri));
+		}
+		for (let doc of docs) {
+			assert.equal(doc.isClosed, true);
+		}
+	});
 });

src/vs/workbench/test/browser/api/extHostNotebookConcatDocument.test.ts

@@ -48,7 +48,7 @@ suite('NotebookConcatDocument', function () {
 		let reg = extHostNotebooks.registerNotebookContentProvider(nullExtensionDescription, 'test', new class extends mock<vscode.NotebookContentProvider>() {
 			// async openNotebook() { }
 		});
-		await extHostNotebooks.$acceptDocumentAndEditorsDelta({
+		extHostNotebooks.$acceptDocumentAndEditorsDelta({
 			addedDocuments: [{
 				handle: 0,
 				uri: notebookUri,
@@ -72,7 +72,7 @@ suite('NotebookConcatDocument', function () {
 				}
 			]
 		});
-		await extHostNotebooks.$acceptDocumentAndEditorsDelta({ newActiveEditor: '_notebook_editor_0' });
+		extHostNotebooks.$acceptDocumentAndEditorsDelta({ newActiveEditor: '_notebook_editor_0' });
 
 		notebook = extHostNotebooks.notebookDocuments[0]!;
 
@@ -292,17 +292,6 @@ suite('NotebookConcatDocument', function () {
 		let cell1End = doc.offsetAt(new Position(2, 12));
 		assert.equal(doc.positionAt(cell1End).isEqual(new Position(2, 12)), true);
 
-		extHostDocumentsAndEditors.$acceptDocumentsAndEditorsDelta({
-			addedDocuments: [{
-				uri: notebook.cells[0].uri,
-				versionId: 1,
-				lines: ['Hello', 'World', 'Hello World!'],
-				EOL: '\n',
-				modeId: '',
-				isDirty: false
-			}]
-		});
-
 		extHostDocuments.$acceptModelChanged(notebook.cells[0].uri, {
 			versionId: 0,
 			eol: '\n',