Change cell execution/cancel API

Fix #99203

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

@@ -79,7 +79,8 @@ export function smokeTestActivate(context: vscode.ExtensionContext): any {
 				}];
 			}
 		},
-		executeCell: async (_document: vscode.NotebookDocument, _cell: vscode.NotebookCell | undefined, _token: vscode.CancellationToken) => {
+		cancelAllCellsExecution: async () => { },
+		executeCell: async (_document: vscode.NotebookDocument, _cell: vscode.NotebookCell | undefined) => {
 			if (!_cell) {
 				_cell = _document.cells[0];
 			}
@@ -92,6 +93,7 @@ export function smokeTestActivate(context: vscode.ExtensionContext): any {
 			}];
 			return;
 		},
+		cancelCellExecution: async () => { }
 	}));
 
 	context.subscriptions.push(vscode.commands.registerCommand('vscode-notebook-tests.debugAction', async (cell: vscode.NotebookCell) => {

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

@@ -62,7 +62,7 @@ export function activate(context: vscode.ExtensionContext): any {
 
 	context.subscriptions.push(vscode.notebook.registerNotebookKernel('notebookKernelTest', ['*.vsctestnb'], {
 		label: 'Notebook Test Kernel',
-		executeAllCells: async (_document: vscode.NotebookDocument, _token: vscode.CancellationToken) => {
+		executeAllCells: async (_document: vscode.NotebookDocument) => {
 			let cell = _document.cells[0];
 
 			cell.outputs = [{
@@ -73,7 +73,8 @@ export function activate(context: vscode.ExtensionContext): any {
 			}];
 			return;
 		},
-		executeCell: async (document: vscode.NotebookDocument, cell: vscode.NotebookCell | undefined, _token: vscode.CancellationToken) => {
+		cancelAllCellsExecution: async (_document: vscode.NotebookDocument) => { },
+		executeCell: async (document: vscode.NotebookDocument, cell: vscode.NotebookCell | undefined) => {
 			if (!cell) {
 				cell = document.cells[0];
 			}
@@ -113,7 +114,8 @@ export function activate(context: vscode.ExtensionContext): any {
 				}
 			});
 			return;
-		}
+		},
+		cancelCellExecution: async (_document: vscode.NotebookDocument, _cell: vscode.NotebookCell) => { }
 	}));
 
 	const preloadUri = vscode.Uri.file(path.resolve(__dirname, '../src/customRenderer.js'));

src/vs/vscode.proposed.d.ts

@@ -1423,6 +1423,11 @@ declare module 'vscode' {
 		Error = 4
 	}
 
+	export enum NotebookRunState {
+		Running = 1,
+		Idle = 2
+	}
+
 	export interface NotebookCellMetadata {
 		/**
 		 * Controls if the content of a cell is editable or not.
@@ -1525,6 +1530,11 @@ declare module 'vscode' {
 		 * Additional attributes of the document metadata.
 		 */
 		custom?: { [key: string]: any };
+
+		/**
+		 * The document's current run state
+		 */
+		runState?: NotebookRunState;
 	}
 
 	export interface NotebookDocument {
@@ -1832,8 +1842,10 @@ declare module 'vscode' {
 		description?: string;
 		isPreferred?: boolean;
 		preloads?: Uri[];
-		executeCell(document: NotebookDocument, cell: NotebookCell, token: CancellationToken): Promise<void>;
-		executeAllCells(document: NotebookDocument, token: CancellationToken): Promise<void>;
+		executeCell(document: NotebookDocument, cell: NotebookCell): void;
+		cancelCellExecution(document: NotebookDocument, cell: NotebookCell): void;
+		executeAllCells(document: NotebookDocument): void;
+		cancelAllCellsExecution(document: NotebookDocument): void;
 	}
 
 	export interface NotebookDocumentFilter {

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

@@ -467,17 +467,23 @@ export class MainThreadNotebooks extends Disposable implements MainThreadNoteboo
 			resolveNotebookEditor: async (viewType: string, uri: URI, editorId: string) => {
 				await this._proxy.$resolveNotebookEditor(viewType, uri, editorId);
 			},
-			executeNotebookByAttachedKernel: async (viewType: string, uri: URI, token: CancellationToken) => {
-				return this.executeNotebookByAttachedKernel(viewType, uri, token);
+			executeNotebookByAttachedKernel: async (viewType: string, uri: URI) => {
+				return this.executeNotebookByAttachedKernel(viewType, uri);
+			},
+			cancelNotebookByAttachedKernel: async (viewType: string, uri: URI) => {
+				return this.cancelNotebookByAttachedKernel(viewType, uri);
 			},
 			onDidReceiveMessage: (editorId: string, rendererType: string | undefined, message: unknown) => {
 				this._proxy.$onDidReceiveMessage(editorId, rendererType, message);
 			},
 			removeNotebookDocument: async (uri: URI) => {
 				return this.removeNotebookTextModel(uri);
 			},
-			executeNotebookCell: async (uri: URI, handle: number, token: CancellationToken) => {
-				return this._proxy.$executeNotebookByAttachedKernel(_viewType, uri, handle, token);
+			executeNotebookCell: async (uri: URI, handle: number) => {
+				return this._proxy.$executeNotebookByAttachedKernel(_viewType, uri, handle);
+			},
+			cancelNotebookCell: async (uri: URI, handle: number) => {
+				return this._proxy.$cancelNotebookByAttachedKernel(_viewType, uri, handle);
 			},
 			save: async (uri: URI, token: CancellationToken) => {
 				return this._proxy.$saveNotebook(_viewType, uri, token);
@@ -537,8 +543,8 @@ export class MainThreadNotebooks extends Disposable implements MainThreadNoteboo
 			resolveKernel: (editorId: string, uri: URI, kernelId: string, token: CancellationToken) => {
 				return that._proxy.$resolveNotebookKernel(handle, editorId, uri, kernelId, token);
 			},
-			executeNotebook: (uri: URI, kernelId: string, cellHandle: number | undefined, token: CancellationToken) => {
-				return that._proxy.$executeNotebookKernelFromProvider(handle, uri, kernelId, cellHandle, token);
+			executeNotebook: (uri: URI, kernelId: string, cellHandle: number | undefined) => {
+				return that._proxy.$executeNotebookKernelFromProvider(handle, uri, kernelId, cellHandle);
 			}
 		});
 		this._notebookKernelProviders.set(handle, {
@@ -590,8 +596,12 @@ export class MainThreadNotebooks extends Disposable implements MainThreadNoteboo
 		}
 	}
 
-	async executeNotebookByAttachedKernel(viewType: string, uri: URI, token: CancellationToken): Promise<void> {
-		return this._proxy.$executeNotebookByAttachedKernel(viewType, uri, undefined, token);
+	async executeNotebookByAttachedKernel(viewType: string, uri: URI): Promise<void> {
+		return this._proxy.$executeNotebookByAttachedKernel(viewType, uri, undefined);
+	}
+
+	async cancelNotebookByAttachedKernel(viewType: string, uri: URI): Promise<void> {
+		return this._proxy.$cancelNotebookByAttachedKernel(viewType, uri, undefined);
 	}
 
 	async $postMessage(editorId: string, forRendererId: string | undefined, value: any): Promise<boolean> {
@@ -634,8 +644,8 @@ export class MainThreadNotebookKernel implements INotebookKernelInfo {
 	) {
 	}
 
-	async executeNotebook(viewType: string, uri: URI, handle: number | undefined, token: CancellationToken): Promise<void> {
-		return this._proxy.$executeNotebook2(this.id, viewType, uri, handle, token);
+	async executeNotebook(viewType: string, uri: URI, handle: number | undefined): Promise<void> {
+		return this._proxy.$executeNotebook2(this.id, viewType, uri, handle);
 	}
 }
 

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

@@ -1108,7 +1108,8 @@ export function createApiFactoryAndRegisterActors(accessor: ServicesAccessor): I
 			TimelineItem: extHostTypes.TimelineItem,
 			CellKind: extHostTypes.CellKind,
 			CellOutputKind: extHostTypes.CellOutputKind,
-			NotebookCellRunState: extHostTypes.NotebookCellRunState
+			NotebookCellRunState: extHostTypes.NotebookCellRunState,
+			NotebookRunState: extHostTypes.NotebookRunState
 		};
 	};
 }

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

@@ -1616,9 +1616,10 @@ export interface ExtHostNotebookShape {
 	$resolveNotebookEditor(viewType: string, uri: UriComponents, editorId: string): Promise<void>;
 	$provideNotebookKernels(handle: number, uri: UriComponents, token: CancellationToken): Promise<INotebookKernelInfoDto2[]>;
 	$resolveNotebookKernel(handle: number, editorId: string, uri: UriComponents, kernelId: string, token: CancellationToken): Promise<void>;
-	$executeNotebookByAttachedKernel(viewType: string, uri: UriComponents, cellHandle: number | undefined, token: CancellationToken): Promise<void>;
-	$executeNotebookKernelFromProvider(handle: number, uri: UriComponents, kernelId: string, cellHandle: number | undefined, token: CancellationToken): Promise<void>;
-	$executeNotebook2(kernelId: string, viewType: string, uri: UriComponents, cellHandle: number | undefined, token: CancellationToken): Promise<void>;
+	$executeNotebookByAttachedKernel(viewType: string, uri: UriComponents, cellHandle: number | undefined): Promise<void>;
+	$cancelNotebookByAttachedKernel(viewType: string, uri: UriComponents, cellHandle: number | undefined): Promise<void>;
+	$executeNotebookKernelFromProvider(handle: number, uri: UriComponents, kernelId: string, cellHandle: number | undefined): Promise<void>;
+	$executeNotebook2(kernelId: string, viewType: string, uri: UriComponents, cellHandle: number | undefined): Promise<void>;
 	$saveNotebook(viewType: string, uri: UriComponents, token: CancellationToken): Promise<boolean>;
 	$saveNotebookAs(viewType: string, uri: UriComponents, target: UriComponents, token: CancellationToken): Promise<boolean>;
 	$backup(viewType: string, uri: UriComponents, cancellation: CancellationToken): Promise<string | undefined>;

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

@@ -16,7 +16,7 @@ import { ExtHostCommands } from 'vs/workbench/api/common/extHostCommands';
 import { ExtHostDocumentsAndEditors } from 'vs/workbench/api/common/extHostDocumentsAndEditors';
 import { CellEditType, diff, ICellEditOperation, ICellInsertEdit, INotebookDisplayOrder, INotebookEditData, NotebookCellsChangedEvent, NotebookCellsSplice2, ICellDeleteEdit, notebookDocumentMetadataDefaults, NotebookCellsChangeType, NotebookDataDto, IOutputRenderRequest, IOutputRenderResponse, IOutputRenderResponseOutputInfo, IOutputRenderResponseCellInfo, IRawOutput, CellOutputKind, IProcessedOutput, INotebookKernelInfoDto2, IMainCellDto } from 'vs/workbench/contrib/notebook/common/notebookCommon';
 import * as extHostTypes from 'vs/workbench/api/common/extHostTypes';
-import { CancellationToken } from 'vs/base/common/cancellation';
+import { CancellationToken, CancellationTokenSource } from 'vs/base/common/cancellation';
 import { ExtHostDocumentData } from 'vs/workbench/api/common/extHostDocumentData';
 import { NotImplementedProxy } from 'vs/base/common/types';
 import * as typeConverters from 'vs/workbench/api/common/extHostTypeConverters';
@@ -830,19 +830,43 @@ export class ExtHostNotebookKernelProviderAdapter extends Disposable {
 		}
 	}
 
-	async executeNotebook(kernelId: string, document: ExtHostNotebookDocument, cell: ExtHostCell | undefined, token: CancellationToken) {
+	async executeNotebook(kernelId: string, document: ExtHostNotebookDocument, cell: ExtHostCell | undefined) {
 		const kernel = this._idToKernel.get(kernelId);
 
 		if (!kernel) {
 			return;
 		}
 
 		if (cell) {
-			return kernel.executeCell(document, cell, token);
+			return withToken(token => (kernel.executeCell as any)(document, cell, token));
 		} else {
-			return kernel.executeAllCells(document, token);
+			return withToken(token => (kernel.executeAllCells as any)(document, token));
 		}
 	}
+
+	async cancelNotebook(kernelId: string, document: ExtHostNotebookDocument, cell: ExtHostCell | undefined) {
+		const kernel = this._idToKernel.get(kernelId);
+
+		if (!kernel) {
+			return;
+		}
+
+		if (cell) {
+			return kernel.cancelCellExecution(document, cell);
+		} else {
+			return kernel.cancelAllCellsExecution(document);
+		}
+	}
+}
+
+// TODO@roblou remove 'token' passed to all execute APIs once extensions are updated
+async function withToken(cb: (token: CancellationToken) => any) {
+	const source = new CancellationTokenSource();
+	try {
+		await cb(source.token);
+	} finally {
+		source.dispose();
+	}
 }
 
 export class ExtHostNotebookController implements ExtHostNotebookShape, ExtHostNotebookOutputRenderingHandler {
@@ -1220,7 +1244,7 @@ export class ExtHostNotebookController implements ExtHostNotebookShape, ExtHostN
 		}
 	}
 
-	async $executeNotebookByAttachedKernel(viewType: string, uri: UriComponents, cellHandle: number | undefined, token: CancellationToken): Promise<void> {
+	async $executeNotebookByAttachedKernel(viewType: string, uri: UriComponents, cellHandle: number | undefined): Promise<void> {
 		let document = this._documents.get(URI.revive(uri).toString());
 
 		if (!document) {
@@ -1233,23 +1257,44 @@ export class ExtHostNotebookController implements ExtHostNotebookShape, ExtHostN
 
 			if (provider.kernel) {
 				if (cell) {
-					return provider.kernel.executeCell(document, cell, token);
+					return withToken(token => (provider.kernel!.executeCell as any)(document, cell, token));
+				} else {
+					return withToken(token => (provider.kernel!.executeAllCells as any)(document, token));
+				}
+			}
+		}
+	}
+
+	async $cancelNotebookByAttachedKernel(viewType: string, uri: UriComponents, cellHandle: number | undefined): Promise<void> {
+		const document = this._documents.get(URI.revive(uri).toString());
+
+		if (!document) {
+			return;
+		}
+
+		if (this._notebookContentProviders.has(viewType)) {
+			const cell = cellHandle !== undefined ? document.getCell(cellHandle) : undefined;
+			const provider = this._notebookContentProviders.get(viewType)!.provider;
+
+			if (provider.kernel) {
+				if (cell) {
+					return provider.kernel.cancelCellExecution(document, cell);
 				} else {
-					return provider.kernel.executeAllCells(document, token);
+					return provider.kernel.cancelAllCellsExecution(document);
 				}
 			}
 		}
 	}
 
-	async $executeNotebookKernelFromProvider(handle: number, uri: UriComponents, kernelId: string, cellHandle: number | undefined, token: CancellationToken): Promise<void> {
+	async $executeNotebookKernelFromProvider(handle: number, uri: UriComponents, kernelId: string, cellHandle: number | undefined): Promise<void> {
 		await this._withAdapter(handle, uri, async (adapter, document) => {
 			let cell = cellHandle !== undefined ? document.getCell(cellHandle) : undefined;
 
-			return adapter.executeNotebook(kernelId, document, cell, token);
+			return adapter.executeNotebook(kernelId, document, cell);
 		});
 	}
 
-	async $executeNotebook2(kernelId: string, viewType: string, uri: UriComponents, cellHandle: number | undefined, token: CancellationToken): Promise<void> {
+	async $executeNotebook2(kernelId: string, viewType: string, uri: UriComponents, cellHandle: number | undefined): Promise<void> {
 		let document = this._documents.get(URI.revive(uri).toString());
 
 		if (!document || document.viewType !== viewType) {
@@ -1265,9 +1310,9 @@ export class ExtHostNotebookController implements ExtHostNotebookShape, ExtHostN
 		let cell = cellHandle !== undefined ? document.getCell(cellHandle) : undefined;
 
 		if (cell) {
-			return kernelInfo.kernel.executeCell(document, cell, token);
+			return withToken(token => (kernelInfo!.kernel.executeCell as any)(document, cell, token));
 		} else {
-			return kernelInfo.kernel.executeAllCells(document, token);
+			return withToken(token => (kernelInfo!.kernel.executeAllCells as any)(document, token));
 		}
 	}
 

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

@@ -2739,6 +2739,11 @@ export enum NotebookCellRunState {
 	Error = 4
 }
 
+export enum NotebookRunState {
+	Running = 1,
+	Idle = 2
+}
+
 //#endregion
 
 //#region Timeline

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

@@ -8,7 +8,6 @@ import { IListEvent, IListMouseEvent } from 'vs/base/browser/ui/list/list';
 import { IListOptions, IListStyles } from 'vs/base/browser/ui/list/listWidget';
 import { ProgressBar } from 'vs/base/browser/ui/progressbar/progressbar';
 import { ToolBar } from 'vs/base/browser/ui/toolbar/toolbar';
-import { CancellationTokenSource } from 'vs/base/common/cancellation';
 import { Event } from 'vs/base/common/event';
 import { DisposableStore, IDisposable } from 'vs/base/common/lifecycle';
 import { ScrollEvent } from 'vs/base/common/scrollable';
@@ -117,7 +116,6 @@ export interface ICellViewModel {
 	language: string;
 	cellKind: CellKind;
 	editState: CellEditState;
-	currentTokenSource: CancellationTokenSource | undefined;
 	focusMode: CellFocusMode;
 	getText(): string;
 	getTextLength(): number;