Merge pull request #96001 from microsoft/rebornix/integration-test

Notebook integration test and undo/redo

.vscode/launch.json

@@ -159,6 +159,24 @@
 				"order": 5
 			}
 		},
+		{
+			"type": "extensionHost",
+			"request": "launch",
+			"name": "VS Code Notebook Tests",
+			"runtimeExecutable": "${execPath}",
+			"args": [
+				"${workspaceFolder}/extensions/vscode-notebook-tests/test",
+				"--extensionDevelopmentPath=${workspaceFolder}/extensions/vscode-notebook-tests",
+				"--extensionTestsPath=${workspaceFolder}/extensions/vscode-notebook-tests/out"
+			],
+			"outFiles": [
+				"${workspaceFolder}/out/**/*.js"
+			],
+			"presentation": {
+				"group": "6_tests",
+				"order": 6
+			}
+		},
 		{
 			"type": "chrome",
 			"request": "attach",

build/lib/extensions.js

@@ -186,6 +186,7 @@ const excludedExtensions = [
     'vscode-test-resolver',
     'ms-vscode.node-debug',
     'ms-vscode.node-debug2',
+    'vscode-notebook-tests'
 ];
 const builtInExtensions = JSON.parse(fs.readFileSync(path.join(__dirname, '../../product.json'), 'utf8')).builtInExtensions;
 function packageLocalExtensionsStream() {

build/lib/extensions.ts

@@ -220,6 +220,7 @@ const excludedExtensions = [
 	'vscode-test-resolver',
 	'ms-vscode.node-debug',
 	'ms-vscode.node-debug2',
+	'vscode-notebook-tests'
 ];
 
 interface IBuiltInExtension {

extensions/vscode-notebook-tests/.vscode/launch.json

+// A launch configuration that compiles the extension and then opens it inside a new window
+{
+	"version": "0.1.0",
+	"configurations": [
+		{
+			"name": "Launch Tests",
+			"type": "extensionHost",
+			"request": "launch",
+			"runtimeExecutable": "${execPath}",
+			"args": ["${workspaceFolder}/../../", "${workspaceFolder}/test", "--extensionDevelopmentPath=${workspaceFolder}", "--extensionTestsPath=${workspaceFolder}/out" ],
+			"stopOnEntry": false,
+			"sourceMaps": true,
+			"outDir": "${workspaceFolder}/out",
+			"preLaunchTask": "npm"
+		}
+	]
+}
\ No newline at end of file

extensions/vscode-notebook-tests/.vscode/tasks.json

+{
+	"version": "2.0.0",
+	"command": "npm",
+	"type": "shell",
+	"presentation": {
+		"reveal": "silent"
+	},
+	"args": ["run", "compile"],
+	"isBackground": true,
+	"problemMatcher": "$tsc-watch"
+}

extensions/vscode-notebook-tests/package.json

+{
+  "name": "vscode-notebook-tests",
+  "description": "Notebook tests for VS Code",
+  "version": "0.0.1",
+  "publisher": "vscode",
+  "license": "MIT",
+  "private": true,
+  "activationEvents": [
+    "*"
+  ],
+  "main": "./out/notebookTestMain",
+  "enableProposedApi": true,
+  "engines": {
+    "vscode": "^1.25.0"
+  },
+  "scripts": {
+    "compile": "node ./node_modules/vscode/bin/compile -watch -p ./",
+    "vscode:prepublish": "node ../../node_modules/gulp/bin/gulp.js --gulpfile ../../build/gulpfile.extensions.js compile-extension:vscode-notebook-tests ./tsconfig.json"
+  },
+  "dependencies": {},
+  "devDependencies": {
+    "typescript": "^3.8.3",
+    "@types/node": "^12.11.7",
+    "mocha-junit-reporter": "^1.17.0",
+    "mocha-multi-reporters": "^1.1.7",
+    "vscode": "~1.1.36",
+    "mocha": "^2.3.3"
+  },
+  "contributes": {
+    "notebookProvider": [
+      {
+        "viewType": "notebookTest",
+        "displayName": "Notebook Test",
+        "selector": [
+          {
+            "filenamePattern": "*.ipynb",
+            "excludeFileNamePattern": "*.test.ipynb"
+          }
+        ]
+      }
+    ]
+  }
+}

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

+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+const path = require('path');
+const testRunner = require('vscode/lib/testrunner');
+
+const suite = 'Integration Notebook Tests';
+
+const options: any = {
+	ui: 'tdd',
+	useColors: (!process.env.BUILD_ARTIFACTSTAGINGDIRECTORY && process.platform !== 'win32'),
+	timeout: 60000
+};
+
+if (process.env.BUILD_ARTIFACTSTAGINGDIRECTORY) {
+	options.reporter = 'mocha-multi-reporters';
+	options.reporterOptions = {
+		reporterEnabled: 'spec, mocha-junit-reporter',
+		mochaJunitReporterReporterOptions: {
+			testsuitesTitle: `${suite} ${process.platform}`,
+			mochaFile: path.join(process.env.BUILD_ARTIFACTSTAGINGDIRECTORY, `test-results/${process.platform}-${suite.toLowerCase().replace(/[^\w]/g, '-')}-results.xml`)
+		}
+	};
+}
+
+testRunner.configure(options);
+
+export = testRunner;

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

+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import 'mocha';
+import * as assert from 'assert';
+import * as vscode from 'vscode';
+import { join } from 'path';
+
+function waitFor(ms: number): Promise<void> {
+	let resolveFunc: () => void;
+
+	const promise = new Promise<void>(resolve => {
+		resolveFunc = resolve;
+	});
+	setTimeout(() => {
+		resolveFunc!();
+	}, ms);
+
+	return promise;
+}
+
+suite('notebook workflow', () => {
+	test('notebook open', async function () {
+		const resource = vscode.Uri.parse(join(vscode.workspace.rootPath || '', './first.ipynb'));
+		await vscode.commands.executeCommand('vscode.openWith', resource, 'notebookTest');
+
+		await waitFor(500);
+		assert.equal(vscode.notebook.activeNotebookEditor !== undefined, true, 'notebook first');
+		assert.equal(vscode.notebook.activeNotebookEditor!.selection?.source, 'test');
+		assert.equal(vscode.notebook.activeNotebookEditor!.selection?.language, 'typescript');
+
+		await vscode.commands.executeCommand('workbench.notebook.code.insertCellBelow');
+		assert.equal(vscode.notebook.activeNotebookEditor!.selection?.source, '');
+
+		await vscode.commands.executeCommand('workbench.notebook.code.insertCellAbove');
+		const activeCell = vscode.notebook.activeNotebookEditor!.selection;
+		assert.notEqual(vscode.notebook.activeNotebookEditor!.selection, undefined);
+		assert.equal(activeCell!.source, '');
+		assert.equal(vscode.notebook.activeNotebookEditor!.document.cells.length, 3);
+		assert.equal(vscode.notebook.activeNotebookEditor!.document.cells.indexOf(activeCell!), 1);
+
+		await vscode.commands.executeCommand('workbench.action.files.save');
+		await vscode.commands.executeCommand('workbench.action.closeActiveEditor');
+	});
+
+	test('notebook cell actions', async function () {
+		const resource = vscode.Uri.parse(join(vscode.workspace.rootPath || '', './second.ipynb'));
+		await vscode.commands.executeCommand('vscode.openWith', resource, 'notebookTest');
+
+		await waitFor(500);
+		assert.equal(vscode.notebook.activeNotebookEditor !== undefined, true, 'notebook first');
+		assert.equal(vscode.notebook.activeNotebookEditor!.selection?.source, 'test');
+		assert.equal(vscode.notebook.activeNotebookEditor!.selection?.language, 'typescript');
+
+		// ---- insert cell below and focus ---- //
+		await vscode.commands.executeCommand('workbench.notebook.code.insertCellBelow');
+		assert.equal(vscode.notebook.activeNotebookEditor!.selection?.source, '');
+
+		// ---- insert cell above and focus ---- //
+		await vscode.commands.executeCommand('workbench.notebook.code.insertCellAbove');
+		let activeCell = vscode.notebook.activeNotebookEditor!.selection;
+		assert.notEqual(vscode.notebook.activeNotebookEditor!.selection, undefined);
+		assert.equal(activeCell!.source, '');
+		assert.equal(vscode.notebook.activeNotebookEditor!.document.cells.length, 3);
+		assert.equal(vscode.notebook.activeNotebookEditor!.document.cells.indexOf(activeCell!), 1);
+
+		// ---- focus bottom ---- //
+		await vscode.commands.executeCommand('workbench.action.notebook.focusBottom');
+		activeCell = vscode.notebook.activeNotebookEditor!.selection;
+		assert.equal(vscode.notebook.activeNotebookEditor!.document.cells.indexOf(activeCell!), 2);
+
+		// ---- focus top and then copy down ---- //
+		await vscode.commands.executeCommand('workbench.action.notebook.focusTop');
+		activeCell = vscode.notebook.activeNotebookEditor!.selection;
+		assert.equal(vscode.notebook.activeNotebookEditor!.document.cells.indexOf(activeCell!), 0);
+
+		await vscode.commands.executeCommand('workbench.notebook.cell.copyDown');
+		activeCell = vscode.notebook.activeNotebookEditor!.selection;
+		assert.equal(vscode.notebook.activeNotebookEditor!.document.cells.indexOf(activeCell!), 1);
+		assert.equal(activeCell?.source, 'test');
+
+		await vscode.commands.executeCommand('workbench.notebook.cell.delete');
+		activeCell = vscode.notebook.activeNotebookEditor!.selection;
+		assert.equal(vscode.notebook.activeNotebookEditor!.document.cells.indexOf(activeCell!), 1);
+		assert.equal(activeCell?.source, '');
+
+		// ---- focus top and then copy up ---- //
+		await vscode.commands.executeCommand('workbench.action.notebook.focusTop');
+		await vscode.commands.executeCommand('workbench.notebook.cell.copyUp');
+		assert.equal(vscode.notebook.activeNotebookEditor!.document.cells.length, 4);
+		assert.equal(vscode.notebook.activeNotebookEditor!.document.cells[0].source, 'test');
+		assert.equal(vscode.notebook.activeNotebookEditor!.document.cells[1].source, 'test');
+		assert.equal(vscode.notebook.activeNotebookEditor!.document.cells[2].source, '');
+		assert.equal(vscode.notebook.activeNotebookEditor!.document.cells[3].source, '');
+		activeCell = vscode.notebook.activeNotebookEditor!.selection;
+		assert.equal(vscode.notebook.activeNotebookEditor!.document.cells.indexOf(activeCell!), 0);
+
+
+		// ---- move up and down ---- //
+
+		// await vscode.commands.executeCommand('workbench.notebook.cell.moveDown');
+		// await vscode.commands.executeCommand('workbench.notebook.cell.moveDown');
+		// activeCell = vscode.notebook.activeNotebookEditor!.selection;
+
+		// assert.equal(vscode.notebook.activeNotebookEditor!.document.cells.indexOf(activeCell!), 2);
+		// assert.equal(vscode.notebook.activeNotebookEditor!.document.cells[0].source, 'test');
+		// assert.equal(vscode.notebook.activeNotebookEditor!.document.cells[1].source, '');
+		// assert.equal(vscode.notebook.activeNotebookEditor!.document.cells[2].source, 'test');
+		// assert.equal(vscode.notebook.activeNotebookEditor!.document.cells[3].source, '');
+
+
+		// ---- ---- //
+
+		await vscode.commands.executeCommand('workbench.action.closeActiveEditor');
+	});
+});

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

+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import * as vscode from 'vscode';
+
+export function activate(context: vscode.ExtensionContext): any {
+	context.subscriptions.push(vscode.notebook.registerNotebookProvider('notebookTest', {
+		resolveNotebook: async (editor: vscode.NotebookEditor) => {
+			await editor.edit(eb => {
+				eb.insert(0, 'test', 'typescript', vscode.CellKind.Code, [], {});
+			});
+			return;
+		},
+		executeCell: async (_document: vscode.NotebookDocument, _cell: vscode.NotebookCell | undefined, _token: vscode.CancellationToken) => {
+			return;
+		},
+		save: async (_document: vscode.NotebookDocument) => {
+			return true;
+		}
+	}));
+}

extensions/vscode-notebook-tests/src/typings/ref.d.ts

+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+/// <reference path="../../../../src/vs/vscode.d.ts" />
+/// <reference path="../../../../src/vs/vscode.proposed.d.ts" />
+/// <reference types='@types/node'/>
+

extensions/vscode-notebook-tests/test/first.ipynb

+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/

extensions/vscode-notebook-tests/test/second.ipynb

+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/

extensions/vscode-notebook-tests/tsconfig.json

+{
+	"extends": "../shared.tsconfig.json",
+	"compilerOptions": {
+		"outDir": "./out"
+	},
+	"include": [
+		"src/**/*"
+	]
+}
\ No newline at end of file

scripts/test-integration.sh

@@ -25,6 +25,7 @@ else
 	# and the build bundles extensions into .build webpacked
 	yarn gulp 	compile-extension:vscode-api-tests \
 				compile-extension:vscode-colorize-tests \
+				compile-extension:vscode-notebook-tests \
 				compile-extension:markdown-language-features \
 				compile-extension:emmet \
 				compile-extension:css-language-features-server \
@@ -44,6 +45,7 @@ fi
 ./scripts/test.sh --runGlob **/*.integrationTest.js "$@"
 
 # Tests in the extension host
+"$INTEGRATION_TEST_ELECTRON_PATH" $LINUX_NO_SANDBOX $ROOT/extensions/vscode-notebook-tests/test --enable-proposed-api=vscode.vscode-notebook-tests --extensionDevelopmentPath=$ROOT/extensions/vscode-notebook-tests --extensionTestsPath=$ROOT/extensions/vscode-notebook-tests/out/ --disable-telemetry --disable-crash-reporter --disable-updates --disable-extensions --skip-getting-started --user-data-dir=$VSCODEUSERDATADIR
 "$INTEGRATION_TEST_ELECTRON_PATH" $LINUX_NO_SANDBOX $ROOT/extensions/vscode-api-tests/testWorkspace --enable-proposed-api=vscode.vscode-api-tests --extensionDevelopmentPath=$ROOT/extensions/vscode-api-tests --extensionTestsPath=$ROOT/extensions/vscode-api-tests/out/singlefolder-tests --disable-telemetry --disable-crash-reporter --disable-updates --disable-extensions --skip-getting-started --user-data-dir=$VSCODEUSERDATADIR
 "$INTEGRATION_TEST_ELECTRON_PATH" $LINUX_NO_SANDBOX $ROOT/extensions/vscode-api-tests/testworkspace.code-workspace --enable-proposed-api=vscode.vscode-api-tests --extensionDevelopmentPath=$ROOT/extensions/vscode-api-tests --extensionTestsPath=$ROOT/extensions/vscode-api-tests/out/workspace-tests --disable-telemetry --disable-crash-reporter --disable-updates --disable-extensions --skip-getting-started --user-data-dir=$VSCODEUSERDATADIR
 "$INTEGRATION_TEST_ELECTRON_PATH" $LINUX_NO_SANDBOX $ROOT/extensions/vscode-colorize-tests/test --extensionDevelopmentPath=$ROOT/extensions/vscode-colorize-tests --extensionTestsPath=$ROOT/extensions/vscode-colorize-tests/out --disable-telemetry --disable-crash-reporter --disable-updates --disable-extensions --skip-getting-started --user-data-dir=$VSCODEUSERDATADIR

src/vs/workbench/contrib/notebook/browser/contrib/coreActions.ts

@@ -687,8 +687,8 @@ registerAction2(class extends Action2 {
 
 async function moveCell(context: INotebookCellActionContext, direction: 'up' | 'down'): Promise<void> {
 	const result = direction === 'up' ?
-		context.notebookEditor.moveCellUp(context.cell) :
-		context.notebookEditor.moveCellDown(context.cell);
+		await context.notebookEditor.moveCellUp(context.cell) :
+		await context.notebookEditor.moveCellDown(context.cell);
 
 	if (result) {
 		// move cell command only works when the cell container has focus

src/vs/workbench/contrib/notebook/browser/view/notebookCellList.ts

@@ -193,6 +193,14 @@ export class NotebookCellList extends WorkbenchList<CellViewModel> implements ID
 			}
 		}));
 
+		this._viewModelStore.add(model.onDidChangeSelection(() => {
+			// convert model selections to view selections
+			const viewSelections = model.selectionHandles.map(handle => {
+				return model.getCellByHandle(handle);
+			}).filter(cell => !!cell).map(cell => this._getViewIndexUpperBound(cell!));
+			this.setFocus(viewSelections);
+		}));
+
 		const hiddenRanges = model.getHiddenRanges();
 		this.setHiddenAreas(hiddenRanges, false);
 		const newRanges = reduceCellRanges(hiddenRanges);
@@ -356,7 +364,7 @@ export class NotebookCellList extends WorkbenchList<CellViewModel> implements ID
 
 	setFocus(indexes: number[], browserEvent?: UIEvent): void {
 		if (this._viewModel) {
-			this._viewModel.selections = indexes.map(index => this.element(index)).map(cell => cell.handle);
+			this._viewModel.selectionHandles = indexes.map(index => this.element(index)).map(cell => cell.handle);
 		}
 
 		super.setFocus(indexes, browserEvent);

src/vs/workbench/contrib/notebook/browser/viewModel/cellEdit.ts

@@ -17,6 +17,7 @@ export interface ICellEditingDelegate {
 	deleteCell?(index: number): void;
 	moveCell?(fromIndex: number, toIndex: number): void;
 	createCellViewModel?(cell: ICell): BaseCellViewModel;
+	setSelections(selections: number[]): void;
 }
 
 export class InsertCellEdit implements IResourceUndoRedoElement {
@@ -26,7 +27,9 @@ export class InsertCellEdit implements IResourceUndoRedoElement {
 		public resource: URI,
 		private insertIndex: number,
 		private cell: BaseCellViewModel,
-		private editingDelegate: ICellEditingDelegate
+		private editingDelegate: ICellEditingDelegate,
+		private beforedSelections: number[],
+		private endSelections: number[]
 	) {
 	}
 
@@ -36,6 +39,7 @@ export class InsertCellEdit implements IResourceUndoRedoElement {
 		}
 
 		this.editingDelegate.deleteCell(this.insertIndex);
+		this.editingDelegate.setSelections(this.beforedSelections);
 	}
 	redo(): void | Promise<void> {
 		if (!this.editingDelegate.insertCell) {
@@ -43,6 +47,7 @@ export class InsertCellEdit implements IResourceUndoRedoElement {
 		}
 
 		this.editingDelegate.insertCell(this.insertIndex, this.cell);
+		this.editingDelegate.setSelections(this.endSelections);
 	}
 }
 
@@ -55,7 +60,9 @@ export class DeleteCellEdit implements IResourceUndoRedoElement {
 		public resource: URI,
 		private insertIndex: number,
 		cell: BaseCellViewModel,
-		private editingDelegate: ICellEditingDelegate
+		private editingDelegate: ICellEditingDelegate,
+		private beforedSelections: number[],
+		private endSelections: number[]
 	) {
 		this._rawCell = cell.model;
 
@@ -70,6 +77,7 @@ export class DeleteCellEdit implements IResourceUndoRedoElement {
 
 		const cell = this.editingDelegate.createCellViewModel(this._rawCell);
 		this.editingDelegate.insertCell(this.insertIndex, cell);
+		this.editingDelegate.setSelections(this.beforedSelections);
 	}
 
 	redo(): void | Promise<void> {
@@ -78,6 +86,7 @@ export class DeleteCellEdit implements IResourceUndoRedoElement {
 		}
 
 		this.editingDelegate.deleteCell(this.insertIndex);
+		this.editingDelegate.setSelections(this.endSelections);
 	}
 }
 
@@ -89,7 +98,9 @@ export class MoveCellEdit implements IResourceUndoRedoElement {
 		public resource: URI,
 		private fromIndex: number,
 		private toIndex: number,
-		private editingDelegate: ICellEditingDelegate
+		private editingDelegate: ICellEditingDelegate,
+		private beforedSelections: number[],
+		private endSelections: number[]
 	) {
 	}
 
@@ -99,6 +110,7 @@ export class MoveCellEdit implements IResourceUndoRedoElement {
 		}
 
 		this.editingDelegate.moveCell(this.toIndex, this.fromIndex);
+		this.editingDelegate.setSelections(this.beforedSelections);
 	}
 
 	redo(): void | Promise<void> {
@@ -107,6 +119,7 @@ export class MoveCellEdit implements IResourceUndoRedoElement {
 		}
 
 		this.editingDelegate.moveCell(this.fromIndex, this.toIndex);
+		this.editingDelegate.setSelections(this.endSelections);
 	}
 }
 
@@ -116,7 +129,9 @@ export class SpliceCellsEdit implements IResourceUndoRedoElement {
 	constructor(
 		public resource: URI,
 		private diffs: [number, CellViewModel[], CellViewModel[]][],
-		private editingDelegate: ICellEditingDelegate
+		private editingDelegate: ICellEditingDelegate,
+		private beforeHandles: number[],
+		private endHandles: number[]
 	) {
 	}
 
@@ -134,6 +149,7 @@ export class SpliceCellsEdit implements IResourceUndoRedoElement {
 				this.editingDelegate.insertCell!(diff[0], cell);
 			});
 		});
+		this.editingDelegate.setSelections(this.beforeHandles);
 	}
 
 	redo(): void | Promise<void> {
@@ -151,5 +167,6 @@ export class SpliceCellsEdit implements IResourceUndoRedoElement {
 			});
 		});
 
+		this.editingDelegate.setSelections(this.endHandles);
 	}
 }

src/vs/workbench/contrib/notebook/browser/viewModel/notebookViewModel.ts

@@ -127,6 +127,20 @@ function _normalizeOptions(options: IModelDecorationOptions): ModelDecorationOpt
 	return ModelDecorationOptions.createDynamic(options);
 }
 
+function selectionsEqual(a: number[], b: number[]) {
+	if (a.length !== b.length) {
+		return false;
+	}
+
+	for (let i = 0; i < a.length; i++) {
+		if (a[i] !== b[i]) {
+			return false;
+		}
+	}
+
+	return true;
+}
+
 let MODEL_ID = 0;
 
 
@@ -197,15 +211,24 @@ export class NotebookViewModel extends Disposable implements EditorFoldingStateD
 		return this._layoutInfo;
 	}
 
+	private readonly _onDidChangeSelection = new Emitter<void>();
+	get onDidChangeSelection(): Event<void> { return this._onDidChangeSelection.event; }
+
 	private _selections: number[] = [];
 
-	get selections() {
+	get selectionHandles() {
 		return this._selections;
 	}
 
-	set selections(selections: number[]) {
+	set selectionHandles(selections: number[]) {
+		selections = selections.sort();
+		if (selectionsEqual(selections, this.selectionHandles)) {
+			return;
+		}
+
 		this._selections = selections;
 		this._model.notebook.selections = selections;
+		this._onDidChangeSelection.fire();
 	}
 
 	private _decorationsTree = new DecorationsTree();
@@ -257,10 +280,39 @@ export class NotebookViewModel extends Disposable implements EditorFoldingStateD
 				splices: diffs
 			});
 
+			let endSelectionHandles: number[] = [];
+			if (this.selectionHandles.length) {
+				const primaryHandle = this.selectionHandles[0];
+				const primarySelectionIndex = this._viewCells.indexOf(this.getCellByHandle(primaryHandle)!);
+				endSelectionHandles = [primaryHandle];
+				let delta = 0;
+
+				for (let i = 0; i < diffs.length; i++) {
+					const diff = diffs[0];
+					if (diff[0] + diff[1] <= primarySelectionIndex) {
+						delta += diff[2].length - diff[1];
+						continue;
+					}
+
+					if (diff[0] > primarySelectionIndex) {
+						endSelectionHandles = [primaryHandle];
+						break;
+					}
+
+					if (diff[0] + diff[1] > primaryHandle) {
+						endSelectionHandles = [this._viewCells[diff[0] + delta].handle];
+						break;
+					}
+				}
+			}
+
 			this.undoService.pushElement(new SpliceCellsEdit(this.uri, undoDiff, {
 				insertCell: this._insertCellDelegate.bind(this),
-				deleteCell: this._deleteCellDelegate.bind(this)
-			}));
+				deleteCell: this._deleteCellDelegate.bind(this),
+				setSelections: this._setSelectionsDelegate.bind(this)
+			}, this.selectionHandles, endSelectionHandles));
+
+			this.selectionHandles = endSelectionHandles;
 		}));
 
 		this._register(this._model.notebook.onDidChangeMetadata(e => {
@@ -379,6 +431,10 @@ export class NotebookViewModel extends Disposable implements EditorFoldingStateD
 		});
 	}
 
+	getCellByHandle(handle: number) {
+		return this._handleToViewCellMapping.get(handle);
+	}
+
 	getCellIndex(cell: ICellViewModel) {
 		return this._viewCells.indexOf(cell as CellViewModel);
 	}
@@ -517,6 +573,10 @@ export class NotebookViewModel extends Disposable implements EditorFoldingStateD
 		this._onDidChangeViewCells.fire({ synchronous: true, splices: [[deleteIndex, 1, []]] });
 	}
 
+	private _setSelectionsDelegate(selections: number[]) {
+		this.selectionHandles = selections;
+	}
+
 	createCell(index: number, source: string[], language: string, type: CellKind, synchronous: boolean) {
 		const cell = this._model.notebook.createCellTextModel(source, language, type, [], undefined);
 		let newCell: CellViewModel = createCellViewModel(this.instantiationService, this, cell);
@@ -524,10 +584,12 @@ export class NotebookViewModel extends Disposable implements EditorFoldingStateD
 		this._handleToViewCellMapping.set(newCell.handle, newCell);
 		this._model.insertCell(cell, index);
 		this._localStore.add(newCell);
+
 		this.undoService.pushElement(new InsertCellEdit(this.uri, index, newCell, {
 			insertCell: this._insertCellDelegate.bind(this),
-			deleteCell: this._deleteCellDelegate.bind(this)
-		}));
+			deleteCell: this._deleteCellDelegate.bind(this),
+			setSelections: this._setSelectionsDelegate.bind(this)
+		}, this.selectionHandles, this.selectionHandles));
 
 		this._decorationsTree.acceptReplace(index, 0, 1, true);
 		this._onDidChangeViewCells.fire({ synchronous: synchronous, splices: [[index, 0, [newCell]]] });
@@ -543,8 +605,9 @@ export class NotebookViewModel extends Disposable implements EditorFoldingStateD
 		this._localStore.add(newCell);
 		this.undoService.pushElement(new InsertCellEdit(this.uri, index, newCell, {
 			insertCell: this._insertCellDelegate.bind(this),
-			deleteCell: this._deleteCellDelegate.bind(this)
-		}));
+			deleteCell: this._deleteCellDelegate.bind(this),
+			setSelections: this._setSelectionsDelegate.bind(this)
+		}, this.selectionHandles, this.selectionHandles));
 
 		this._decorationsTree.acceptReplace(index, 0, 1, true);
 		this._onDidChangeViewCells.fire({ synchronous: synchronous, splices: [[index, 0, [newCell]]] });
@@ -552,20 +615,41 @@ export class NotebookViewModel extends Disposable implements EditorFoldingStateD
 	}
 
 	deleteCell(index: number, synchronous: boolean) {
+		const primarySelectionIndex = this.selectionHandles.length ? this._viewCells.indexOf(this.getCellByHandle(this.selectionHandles[0])!) : null;
+
 		let viewCell = this._viewCells[index];
 		this._viewCells.splice(index, 1);
 		this._handleToViewCellMapping.delete(viewCell.handle);
 
 		this._model.deleteCell(index);
 
+		let endSelections: number[] = [];
+		if (this.selectionHandles.length) {
+			const primarySelectionHandle = this.selectionHandles[0];
+
+			if (index === primarySelectionIndex) {
+				if (primarySelectionIndex < this.length - 1) {
+					endSelections = [this._viewCells[primarySelectionIndex + 1].handle];
+				} else if (primarySelectionIndex === this.length - 1 && this.length > 1) {
+					endSelections = [this._viewCells[primarySelectionIndex - 1].handle];
+				} else {
+					endSelections = [];
+				}
+			} else {
+				endSelections = [primarySelectionHandle];
+			}
+		}
+
 		this.undoService.pushElement(new DeleteCellEdit(this.uri, index, viewCell, {
 			insertCell: this._insertCellDelegate.bind(this),
 			deleteCell: this._deleteCellDelegate.bind(this),
 			createCellViewModel: (cell: NotebookCellTextModel) => {
 				return createCellViewModel(this.instantiationService, this, cell);
-			}
-		}));
+			},
+			setSelections: this._setSelectionsDelegate.bind(this)
+		}, this.selectionHandles, endSelections));
 
+		this.selectionHandles = endSelections;
 
 		this._decorationsTree.acceptReplace(index, 1, 0, true);
 
@@ -589,10 +673,13 @@ export class NotebookViewModel extends Disposable implements EditorFoldingStateD
 			this.undoService.pushElement(new MoveCellEdit(this.uri, index, newIdx, {
 				moveCell: (fromIndex: number, toIndex: number) => {
 					this.moveCellToIdx(fromIndex, toIndex, true, false);
-				}
-			}));
+				},
+				setSelections: this._setSelectionsDelegate.bind(this)
+			}, this.selectionHandles, this.selectionHandles));
 		}
 
+		this.selectionHandles = this.selectionHandles;
+
 		this._onDidChangeViewCells.fire({ synchronous: synchronous, splices: [[index, 1, []]] });
 		this._onDidChangeViewCells.fire({ synchronous: synchronous, splices: [[newIdx, 0, [viewCell]]] });