diff --git a/src/vs/workbench/api/browser/mainThreadNotebook.ts b/src/vs/workbench/api/browser/mainThreadNotebook.ts
index d927eea3612..4a922ceb819 100644
--- a/src/vs/workbench/api/browser/mainThreadNotebook.ts
+++ b/src/vs/workbench/api/browser/mainThreadNotebook.ts
@@ -255,7 +255,7 @@ export class MainThreadNotebookController implements IMainNotebookController {
 	}
 
 	async executeNotebook(viewType: string, uri: URI, token: CancellationToken): Promise<void> {
-		this._mainThreadNotebook.executeNotebook(viewType, uri, token);
+		return this._mainThreadNotebook.executeNotebook(viewType, uri, token);
 	}
 
 	onDidReceiveMessage(uri: UriComponents, message: any): void {
diff --git a/src/vs/workbench/contrib/notebook/browser/notebookEditorWidget.ts b/src/vs/workbench/contrib/notebook/browser/notebookEditorWidget.ts
index bb4d12af72f..544eb1be04d 100644
--- a/src/vs/workbench/contrib/notebook/browser/notebookEditorWidget.ts
+++ b/src/vs/workbench/contrib/notebook/browser/notebookEditorWidget.ts
@@ -39,7 +39,7 @@ import { CellDragAndDropController, CodeCellRenderer, MarkdownCellRenderer, Note
 import { CodeCellViewModel } from 'vs/workbench/contrib/notebook/browser/viewModel/codeCellViewModel';
 import { NotebookEventDispatcher, NotebookLayoutChangedEvent } from 'vs/workbench/contrib/notebook/browser/viewModel/eventDispatcher';
 import { CellViewModel, IModelDecorationsChangeAccessor, INotebookEditorViewState, NotebookViewModel } from 'vs/workbench/contrib/notebook/browser/viewModel/notebookViewModel';
-import { CellKind, CellUri, IOutput } from 'vs/workbench/contrib/notebook/common/notebookCommon';
+import { CellKind, IOutput } from 'vs/workbench/contrib/notebook/common/notebookCommon';
 import { NotebookEditorModel } from 'vs/workbench/contrib/notebook/common/notebookEditorModel';
 import { INotebookService } from 'vs/workbench/contrib/notebook/common/notebookService';
 import { Webview } from 'vs/workbench/contrib/webview/browser/webview';
@@ -979,7 +979,6 @@ export class NotebookEditorWidget extends Disposable implements INotebookEditor
 			return;
 		}
 
-		// return this.progressService.showWhile(this._executeNotebook());
 		return this._executeNotebook();
 	}
 
@@ -993,10 +992,11 @@ export class NotebookEditorWidget extends Disposable implements INotebookEditor
 			this.editorExecutingNotebook!.set(true);
 			this.notebookViewModel!.currentTokenSource = tokenSource;
 
-			for (let cell of this.notebookViewModel!.viewCells) {
-				if (cell.cellKind === CellKind.Code) {
-					await this._executeNotebookCell(cell, tokenSource);
-				}
+			const provider = this.notebookService.getContributedNotebookProviders(this.viewModel!.uri)[0];
+			if (provider) {
+				const viewType = provider.id;
+				const notebookUri = this.notebookViewModel!.uri;
+				return await this.notebookService.executeNotebook(viewType, notebookUri, tokenSource.token);
 			}
 		} finally {
 			this.editorExecutingNotebook!.set(false);
@@ -1030,13 +1030,12 @@ export class NotebookEditorWidget extends Disposable implements INotebookEditor
 	private async _executeNotebookCell(cell: ICellViewModel, tokenSource: CancellationTokenSource): Promise<void> {
 		try {
 			cell.currentTokenSource = tokenSource;
+
 			const provider = this.notebookService.getContributedNotebookProviders(this.viewModel!.uri)[0];
 			if (provider) {
 				const viewType = provider.id;
-				const notebookUri = CellUri.parse(cell.uri)?.notebook;
-				if (notebookUri) {
-					return await this.notebookService.executeNotebookCell(viewType, notebookUri, cell.handle, tokenSource.token);
-				}
+				const notebookUri = this.notebookViewModel!.uri;
+				return await this.notebookService.executeNotebookCell(viewType, notebookUri, cell.handle, tokenSource.token);
 			}
 		} finally {
 			cell.currentTokenSource = undefined;
diff --git a/src/vs/workbench/contrib/notebook/browser/notebookServiceImpl.ts b/src/vs/workbench/contrib/notebook/browser/notebookServiceImpl.ts
index ca530378bbc..5b6e38f822c 100644
--- a/src/vs/workbench/contrib/notebook/browser/notebookServiceImpl.ts
+++ b/src/vs/workbench/contrib/notebook/browser/notebookServiceImpl.ts
@@ -16,7 +16,7 @@ import { IExtensionService } from 'vs/workbench/services/extensions/common/exten
 import { NotebookOutputRendererInfo } from 'vs/workbench/contrib/notebook/common/notebookOutputRenderer';
 import { Iterable } from 'vs/base/common/iterator';
 import { NotebookTextModel } from 'vs/workbench/contrib/notebook/common/model/notebookTextModel';
-import { CancellationToken, CancellationTokenSource } from 'vs/base/common/cancellation';
+import { CancellationToken } from 'vs/base/common/cancellation';
 import { IEditorService, ICustomEditorViewTypesHandler, ICustomEditorInfo } from 'vs/workbench/services/editor/common/editorService';
 import { NotebookCellTextModel } from 'vs/workbench/contrib/notebook/common/model/notebookCellTextModel';
 import { NotebookEditorModelManager } from 'vs/workbench/contrib/notebook/common/notebookEditorModel';
@@ -256,11 +256,11 @@ export class NotebookService extends Disposable implements INotebookService, ICu
 		return modelData.model;
 	}
 
-	async executeNotebook(viewType: string, uri: URI): Promise<void> {
+	async executeNotebook(viewType: string, uri: URI, token: CancellationToken): Promise<void> {
 		let provider = this._notebookProviders.get(viewType);
 
 		if (provider) {
-			return provider.controller.executeNotebook(viewType, uri, new CancellationTokenSource().token); // Cancellation for notebooks - TODO
+			return provider.controller.executeNotebook(viewType, uri, token);
 		}
 
 		return;