From 266e475931645a25cd1601a5dd3b0914911d95ea Mon Sep 17 00:00:00 2001
From: Johannes Rieken <johannes.rieken@gmail.com>
Date: Tue, 16 Feb 2021 11:09:23 +0100
Subject: [PATCH] before removing cell documents capture its API objects, after
 inserting cell documents capture its API objects, fixes
 https://github.com/microsoft/vscode/issues/116711

---
 .../api/common/extHostNotebookDocument.ts     |  8 +--
 .../test/browser/api/extHostNotebook.test.ts  | 50 ++++++++++++++++++-
 2 files changed, 54 insertions(+), 4 deletions(-)

diff --git a/src/vs/workbench/api/common/extHostNotebookDocument.ts b/src/vs/workbench/api/common/extHostNotebookDocument.ts
index 3e9e9e41a85..e9e715885c4 100644
--- a/src/vs/workbench/api/common/extHostNotebookDocument.ts
+++ b/src/vs/workbench/api/common/extHostNotebookDocument.ts
@@ -18,13 +18,13 @@ import * as vscode from 'vscode';
 
 class RawContentChangeEvent {
 
-	constructor(readonly start: number, readonly deletedCount: number, readonly deletedItems: ExtHostCell[], readonly items: ExtHostCell[]) { }
+	constructor(readonly start: number, readonly deletedCount: number, readonly deletedItems: vscode.NotebookCell[], readonly items: ExtHostCell[]) { }
 
 	static asApiEvent(event: RawContentChangeEvent): vscode.NotebookCellsChangeData {
 		return Object.freeze({
 			start: event.start,
 			deletedCount: event.deletedCount,
-			deletedItems: event.deletedItems.map(data => data.cell),
+			deletedItems: event.deletedItems,
 			items: event.items.map(data => data.cell)
 		});
 	}
@@ -252,12 +252,14 @@ export class ExtHostNotebookDocument extends Disposable {
 				this._cellDisposableMapping.delete(this._cells[j].handle);
 			}
 
+			const changeEvent = new RawContentChangeEvent(splice[0], splice[1], [], newCells);
 			const deletedItems = this._cells.splice(splice[0], splice[1], ...newCells);
 			for (let cell of deletedItems) {
 				removedCellDocuments.push(cell.uri);
+				changeEvent.deletedItems.push(cell.cell);
 			}
 
-			contentChangeEvents.push(new RawContentChangeEvent(splice[0], splice[1], deletedItems, newCells));
+			contentChangeEvents.push(changeEvent);
 		});
 
 		this._documentsAndEditors.acceptDocumentsAndEditorsDelta({
diff --git a/src/vs/workbench/test/browser/api/extHostNotebook.test.ts b/src/vs/workbench/test/browser/api/extHostNotebook.test.ts
index 7619a7b337b..a85f77ba805 100644
--- a/src/vs/workbench/test/browser/api/extHostNotebook.test.ts
+++ b/src/vs/workbench/test/browser/api/extHostNotebook.test.ts
@@ -21,6 +21,7 @@ import { nullExtensionDescription } from 'vs/workbench/services/extensions/commo
 import { isEqual } from 'vs/base/common/resources';
 import { IExtensionStoragePaths } from 'vs/workbench/api/common/extHostStoragePaths';
 import { generateUuid } from 'vs/base/common/uuid';
+import { Event } from 'vs/base/common/event';
 
 suite('NotebookCell#Document', function () {
 
@@ -33,9 +34,11 @@ suite('NotebookCell#Document', function () {
 	const notebookUri = URI.parse('test:///notebook.file');
 	const disposables = new DisposableStore();
 
-	setup(async function () {
+	teardown(function () {
 		disposables.clear();
+	});
 
+	setup(async function () {
 		rpcProtocol = new TestRPCProtocol();
 		rpcProtocol.set(MainContext.MainThreadCommands, new class extends mock<MainThreadCommandsShape>() {
 			$registerCommand() { }
@@ -305,4 +308,49 @@ suite('NotebookCell#Document', function () {
 		assert.strictEqual(notebook.notebookDocument.cells.length, 3);
 		assert.strictEqual(second.index, 2);
 	});
+
+	test('ERR MISSING extHostDocument for notebook cell: #116711', async function () {
+
+		const p = Event.toPromise(extHostNotebooks.onDidChangeNotebookCells);
+
+		// DON'T call this, make sure the cell-documents have not been created yet
+		// assert.strictEqual(notebook.notebookDocument.cells.length, 2);
+
+		extHostNotebooks.$acceptModelChanged(notebook.uri, {
+			versionId: 100,
+			rawEvents: [{
+				kind: NotebookCellsChangeType.ModelChange,
+				changes: [[0, 2, [{
+					handle: 3,
+					uri: CellUri.generate(notebookUri, 3),
+					source: ['### Heading'],
+					eol: '\n',
+					language: 'markdown',
+					cellKind: CellKind.Markdown,
+					outputs: [],
+				}, {
+					handle: 4,
+					uri: CellUri.generate(notebookUri, 4),
+					source: ['console.log("aaa")', 'console.log("bbb")'],
+					eol: '\n',
+					language: 'javascript',
+					cellKind: CellKind.Code,
+					outputs: [],
+				}]]]
+			}]
+		}, false);
+
+		assert.strictEqual(notebook.notebookDocument.cells.length, 2);
+
+		const event = await p;
+
+		assert.strictEqual(event.document === notebook.notebookDocument, true);
+		assert.strictEqual(event.changes.length, 1);
+		assert.strictEqual(event.changes[0].deletedCount, 2);
+		assert.strictEqual(event.changes[0].deletedItems[0].document.isClosed, true);
+		assert.strictEqual(event.changes[0].deletedItems[1].document.isClosed, true);
+		assert.strictEqual(event.changes[0].items.length, 2);
+		assert.strictEqual(event.changes[0].items[0].document.isClosed, false);
+		assert.strictEqual(event.changes[0].items[1].document.isClosed, false);
+	});
 });