add TextDocument#notebook api proposal, https://github.com/microsoft/vscode/issues/100890

src/vs/vscode.proposed.d.ts

@@ -2024,4 +2024,12 @@ declare module 'vscode' {
 	}
 
 	//#endregion
+
+
+	//#region https://github.com/microsoft/vscode/issues/102091
+
+	export interface TextDocument {
+		notebook: NotebookDocument | undefined;
+	}
+	//#endregion
 }

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

@@ -29,19 +29,17 @@ export function getWordDefinitionFor(modeId: string): RegExp | undefined {
 
 export class ExtHostDocumentData extends MirrorTextModel {
 
-	private _proxy: MainThreadDocumentsShape;
-	private _languageId: string;
-	private _isDirty: boolean;
 	private _document?: vscode.TextDocument;
 	private _isDisposed: boolean = false;
 
-	constructor(proxy: MainThreadDocumentsShape, uri: URI, lines: string[], eol: string,
-		languageId: string, versionId: number, isDirty: boolean
+	constructor(
+		private readonly _proxy: MainThreadDocumentsShape,
+		uri: URI, lines: string[], eol: string, versionId: number,
+		private _languageId: string,
+		private _isDirty: boolean,
+		private readonly _notebook?: vscode.NotebookDocument | undefined
 	) {
 		super(uri, lines, eol, versionId);
-		this._proxy = proxy;
-		this._languageId = languageId;
-		this._isDirty = isDirty;
 	}
 
 	dispose(): void {
@@ -59,25 +57,26 @@ export class ExtHostDocumentData extends MirrorTextModel {
 
 	get document(): vscode.TextDocument {
 		if (!this._document) {
-			const data = this;
+			const that = this;
 			this._document = {
-				get uri() { return data._uri; },
-				get fileName() { return data._uri.fsPath; },
-				get isUntitled() { return data._uri.scheme === Schemas.untitled; },
-				get languageId() { return data._languageId; },
-				get version() { return data._versionId; },
-				get isClosed() { return data._isDisposed; },
-				get isDirty() { return data._isDirty; },
-				save() { return data._save(); },
-				getText(range?) { return range ? data._getTextInRange(range) : data.getText(); },
-				get eol() { return data._eol === '\n' ? EndOfLine.LF : EndOfLine.CRLF; },
-				get lineCount() { return data._lines.length; },
-				lineAt(lineOrPos: number | vscode.Position) { return data._lineAt(lineOrPos); },
-				offsetAt(pos) { return data._offsetAt(pos); },
-				positionAt(offset) { return data._positionAt(offset); },
-				validateRange(ran) { return data._validateRange(ran); },
-				validatePosition(pos) { return data._validatePosition(pos); },
-				getWordRangeAtPosition(pos, regexp?) { return data._getWordRangeAtPosition(pos, regexp); }
+				get uri() { return that._uri; },
+				get fileName() { return that._uri.fsPath; },
+				get isUntitled() { return that._uri.scheme === Schemas.untitled; },
+				get languageId() { return that._languageId; },
+				get version() { return that._versionId; },
+				get isClosed() { return that._isDisposed; },
+				get isDirty() { return that._isDirty; },
+				get notebook() { return that._notebook; },
+				save() { return that._save(); },
+				getText(range?) { return range ? that._getTextInRange(range) : that.getText(); },
+				get eol() { return that._eol === '\n' ? EndOfLine.LF : EndOfLine.CRLF; },
+				get lineCount() { return that._lines.length; },
+				lineAt(lineOrPos: number | vscode.Position) { return that._lineAt(lineOrPos); },
+				offsetAt(pos) { return that._offsetAt(pos); },
+				positionAt(offset) { return that._positionAt(offset); },
+				validateRange(ran) { return that._validateRange(ran); },
+				validatePosition(pos) { return that._validatePosition(pos); },
+				getWordRangeAtPosition(pos, regexp?) { return that._getWordRangeAtPosition(pos, regexp); },
 			};
 		}
 		return Object.freeze(this._document);

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

@@ -4,11 +4,12 @@
  *--------------------------------------------------------------------------------------------*/
 
 import * as assert from 'vs/base/common/assert';
+import * as vscode from 'vscode';
 import { Emitter, Event } from 'vs/base/common/event';
 import { dispose } from 'vs/base/common/lifecycle';
 import { URI } from 'vs/base/common/uri';
 import { createDecorator } from 'vs/platform/instantiation/common/instantiation';
-import { ExtHostDocumentsAndEditorsShape, IDocumentsAndEditorsDelta, MainContext } from 'vs/workbench/api/common/extHost.protocol';
+import { ExtHostDocumentsAndEditorsShape, IDocumentsAndEditorsDelta, IModelAddedData, MainContext } from 'vs/workbench/api/common/extHost.protocol';
 import { ExtHostDocumentData } from 'vs/workbench/api/common/extHostDocumentData';
 import { IExtHostRpcService } from 'vs/workbench/api/common/extHostRpcService';
 import { ExtHostTextEditor } from 'vs/workbench/api/common/extHostTextEditor';
@@ -29,6 +30,14 @@ class Reference<T> {
 	}
 }
 
+export interface IExtHostModelAddedData extends IModelAddedData {
+	notebook?: vscode.NotebookDocument;
+}
+
+export interface IExtHostDocumentsAndEditorsDelta extends IDocumentsAndEditorsDelta {
+	addedDocuments?: IExtHostModelAddedData[];
+}
+
 export class ExtHostDocumentsAndEditors implements ExtHostDocumentsAndEditorsShape {
 
 	readonly _serviceBrand: undefined;
@@ -54,6 +63,10 @@ export class ExtHostDocumentsAndEditors implements ExtHostDocumentsAndEditorsSha
 	) { }
 
 	$acceptDocumentsAndEditorsDelta(delta: IDocumentsAndEditorsDelta): void {
+		this.acceptDocumentsAndEditorsDelta(delta);
+	}
+
+	acceptDocumentsAndEditorsDelta(delta: IExtHostDocumentsAndEditorsDelta): void {
 
 		const removedDocuments: ExtHostDocumentData[] = [];
 		const addedDocuments: ExtHostDocumentData[] = [];
@@ -88,9 +101,10 @@ export class ExtHostDocumentsAndEditors implements ExtHostDocumentsAndEditorsSha
 						resource,
 						data.lines,
 						data.EOL,
-						data.modeId,
 						data.versionId,
-						data.isDirty
+						data.modeId,
+						data.isDirty,
+						data.notebook
 					));
 					this._documents.set(resource, ref);
 					addedDocuments.push(ref.value);

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

@@ -17,7 +17,7 @@ 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 { ExtHostCommands } from 'vs/workbench/api/common/extHostCommands';
-import { ExtHostDocumentsAndEditors } from 'vs/workbench/api/common/extHostDocumentsAndEditors';
+import { ExtHostDocumentsAndEditors, IExtHostModelAddedData } from 'vs/workbench/api/common/extHostDocumentsAndEditors';
 import { IExtensionStoragePaths } from 'vs/workbench/api/common/extHostStoragePaths';
 import * as typeConverters from 'vs/workbench/api/common/extHostTypeConverters';
 import * as extHostTypes from 'vs/workbench/api/common/extHostTypes';
@@ -61,14 +61,15 @@ const addIdToOutput = (output: IRawOutput, id = UUID.generateUuid()): IProcessed
 
 export class ExtHostCell extends Disposable implements vscode.NotebookCell {
 
-	public static asModelAddData(cell: IMainCellDto): IModelAddedData {
+	public static asModelAddData(notebook: ExtHostNotebookDocument, cell: IMainCellDto): IExtHostModelAddedData {
 		return {
 			EOL: cell.eol,
 			lines: cell.source,
 			modeId: cell.language,
 			uri: cell.uri,
 			isDirty: false,
-			versionId: 1
+			versionId: 1,
+			notebook
 		};
 	}
 
@@ -363,7 +364,7 @@ export class ExtHostNotebookDocument extends Disposable implements vscode.Notebo
 		}
 
 		const contentChangeEvents: vscode.NotebookCellsChangeData[] = [];
-		const addedCellDocuments: IModelAddedData[] = [];
+		const addedCellDocuments: IExtHostModelAddedData[] = [];
 
 		splices.reverse().forEach(splice => {
 			const cellDtos = splice[2];
@@ -372,7 +373,7 @@ 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));
+					addedCellDocuments.push(ExtHostCell.asModelAddData(this, cell));
 				}
 
 				if (!this._cellDisposableMapping.has(extCell.handle)) {
@@ -404,7 +405,7 @@ export class ExtHostNotebookDocument extends Disposable implements vscode.Notebo
 		});
 
 		if (addedCellDocuments) {
-			this._documentsAndEditors.$acceptDocumentsAndEditorsDelta({ addedDocuments: addedCellDocuments });
+			this._documentsAndEditors.acceptDocumentsAndEditorsDelta({ addedDocuments: addedCellDocuments });
 		}
 
 		if (!initialization) {
@@ -1588,7 +1589,7 @@ export class ExtHostNotebookController implements ExtHostNotebookShape, ExtHostN
 					});
 
 					// add cell document as vscode.TextDocument
-					addedCellDocuments.push(...modelData.cells.map(ExtHostCell.asModelAddData));
+					addedCellDocuments.push(...modelData.cells.map(cell => ExtHostCell.asModelAddData(document, cell)));
 
 					this._documents.get(revivedUri)?.dispose();
 					this._documents.set(revivedUri, document);
@@ -1600,9 +1601,7 @@ export class ExtHostNotebookController implements ExtHostNotebookShape, ExtHostN
 					}
 				}
 
-				this._documentsAndEditors.$acceptDocumentsAndEditorsDelta({
-					addedDocuments: addedCellDocuments
-				});
+				this._documentsAndEditors.$acceptDocumentsAndEditorsDelta({ addedDocuments: addedCellDocuments });
 
 				const document = this._documents.get(revivedUri)!;
 				this._onDidOpenNotebookDocument.fire(document);

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

@@ -35,7 +35,7 @@ suite('ExtHostDocumentData', () => {
 			'and this is line number two', //27
 			'it is followed by #3', //20
 			'and finished with the fourth.', //29
-		], '\n', 'text', 1, false);
+		], '\n', 1, 'text', false);
 	});
 
 	test('readonly-ness', () => {
@@ -55,7 +55,7 @@ suite('ExtHostDocumentData', () => {
 				saved = uri;
 				return Promise.resolve(true);
 			}
-		}, URI.parse('foo:bar'), [], '\n', 'text', 1, true);
+		}, URI.parse('foo:bar'), [], '\n', 1, 'text', true);
 
 		return data.document.save().then(() => {
 			assert.equal(saved.toString(), 'foo:bar');
@@ -242,7 +242,7 @@ suite('ExtHostDocumentData', () => {
 	test('getWordRangeAtPosition', () => {
 		data = new ExtHostDocumentData(undefined!, URI.file(''), [
 			'aaaa bbbb+cccc abc'
-		], '\n', 'text', 1, false);
+		], '\n', 1, 'text', false);
 
 		let range = data.document.getWordRangeAtPosition(new Position(0, 2))!;
 		assert.equal(range.start.line, 0);
@@ -276,7 +276,7 @@ suite('ExtHostDocumentData', () => {
 			'function() {',
 			'	"far boo"',
 			'}'
-		], '\n', 'text', 1, false);
+		], '\n', 1, 'text', false);
 
 		let range = data.document.getWordRangeAtPosition(new Position(0, 0), /\/\*.+\*\//);
 		assert.equal(range, undefined);
@@ -304,7 +304,7 @@ suite('ExtHostDocumentData', () => {
 
 		data = new ExtHostDocumentData(undefined!, URI.file(''), [
 			perfData._$_$_expensive
-		], '\n', 'text', 1, false);
+		], '\n', 1, 'text', false);
 
 		let range = data.document.getWordRangeAtPosition(new Position(0, 1_177_170), regex)!;
 		assert.equal(range, undefined);
@@ -323,7 +323,7 @@ suite('ExtHostDocumentData', () => {
 
 		data = new ExtHostDocumentData(undefined!, URI.file(''), [
 			line
-		], '\n', 'text', 1, false);
+		], '\n', 1, 'text', false);
 
 		let range = data.document.getWordRangeAtPosition(new Position(0, 27), regex)!;
 		assert.equal(range.start.line, 0);
@@ -387,7 +387,7 @@ suite('ExtHostDocumentData updates line mapping', () => {
 	}
 
 	function testLineMappingDirectionAfterEvents(lines: string[], eol: string, direction: AssertDocumentLineMappingDirection, e: IModelChangedEvent): void {
-		let myDocument = new ExtHostDocumentData(undefined!, URI.file(''), lines.slice(0), eol, 'text', 1, false);
+		let myDocument = new ExtHostDocumentData(undefined!, URI.file(''), lines.slice(0), eol, 1, 'text', false);
 		assertDocumentLineMapping(myDocument, direction);
 
 		myDocument.onEvents(e);

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

@@ -97,11 +97,13 @@ suite('NotebookCell#Document', function () {
 		assert.ok(d1);
 		assert.equal(d1.languageId, c1.language);
 		assert.equal(d1.version, 1);
+		assert.ok(d1.notebook === notebook);
 
 		const d2 = extHostDocuments.getDocument(c2.uri);
 		assert.ok(d2);
 		assert.equal(d2.languageId, c2.language);
 		assert.equal(d2.version, 1);
+		assert.ok(d2.notebook === notebook);
 	});
 
 	test('cell document goes when notebook closes', async function () {
@@ -215,4 +217,10 @@ suite('NotebookCell#Document', function () {
 			assert.equal(doc.isClosed, true);
 		}
 	});
+
+	test('cell document knows notebook', function () {
+		for (let cells of notebook.cells) {
+			assert.equal(cells.document.notebook === notebook, true);
+		}
+	});
 });

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

@@ -17,7 +17,7 @@ suite('ExtHostTextEditor', () => {
 	let editor: ExtHostTextEditor;
 	let doc = new ExtHostDocumentData(undefined!, URI.file(''), [
 		'aaaa bbbb+cccc abc'
-	], '\n', 'text', 1, false);
+	], '\n', 1, 'text', false);
 
 	setup(() => {
 		editor = new ExtHostTextEditor('fake', null!, new NullLogService(), doc, [], { cursorStyle: 0, insertSpaces: true, lineNumbers: 1, tabSize: 4, indentSize: 4 }, [], 1);