Support for arbitrary notebook kernel execution

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

@@ -15,7 +15,7 @@ import { NotebookDto } from 'vs/workbench/api/browser/mainThreadNotebookDto';
 import { extHostNamedCustomer, IExtHostContext } from 'vs/workbench/services/extensions/common/extHostCustomers';
 import { INotebookEditor } from 'vs/workbench/contrib/notebook/browser/notebookBrowser';
 import { INotebookEditorService } from 'vs/workbench/contrib/notebook/browser/services/notebookEditorService';
-import { INotebookCellExecution, INotebookExecutionStateService } from 'vs/workbench/contrib/notebook/common/notebookExecutionStateService';
+import { INotebookCellExecution, INotebookExecution, INotebookExecutionStateService, NotebookExecutionType } from 'vs/workbench/contrib/notebook/common/notebookExecutionStateService';
 import { IKernelSourceActionProvider, INotebookKernel, INotebookKernelChangeEvent, INotebookKernelDetectionTask, INotebookKernelService } from 'vs/workbench/contrib/notebook/common/notebookKernelService';
 import { SerializableObjectWithBuffers } from 'vs/workbench/services/extensions/common/proxyIdentifier';
 import { ExtHostContext, ExtHostNotebookKernelsShape, ICellExecuteUpdateDto, ICellExecutionCompleteDto, INotebookKernelDto2, MainContext, MainThreadNotebookKernelsShape } from '../common/extHost.protocol';
@@ -114,6 +114,7 @@ export class MainThreadNotebookKernels implements MainThreadNotebookKernelsShape
 	private readonly _proxy: ExtHostNotebookKernelsShape;
 
 	private readonly _executions = new Map<number, INotebookCellExecution>();
+	private readonly _notebookExecutions = new Map<number, INotebookExecution>();
 
 	constructor(
 		extHostContext: IExtHostContext,
@@ -134,10 +135,13 @@ export class MainThreadNotebookKernels implements MainThreadNotebookKernelsShape
 			this._executions.forEach(e => {
 				e.complete({});
 			});
+			this._notebookExecutions.forEach(e => e.complete());
 		}));
 
-		this._disposables.add(this._notebookExecutionStateService.onDidChangeCellExecution(e => {
-			this._proxy.$cellExecutionChanged(e.notebook, e.cellHandle, e.changed?.state);
+		this._disposables.add(this._notebookExecutionStateService.onDidChangeExecution(e => {
+			if (e.type === NotebookExecutionType.cell) {
+				this._proxy.$cellExecutionChanged(e.notebook, e.cellHandle, e.changed?.state);
+			}
 		}));
 	}
 
@@ -257,7 +261,7 @@ export class MainThreadNotebookKernels implements MainThreadNotebookKernelsShape
 		}
 	}
 
-	// --- execution
+	// --- Cell execution
 
 	$createExecution(handle: number, controllerId: string, rawUri: UriComponents, cellHandle: number): void {
 		const uri = URI.revive(rawUri);
@@ -296,6 +300,44 @@ export class MainThreadNotebookKernels implements MainThreadNotebookKernelsShape
 		}
 	}
 
+	// --- Notebook execution
+
+	$createNotebookExecution(handle: number, controllerId: string, rawUri: UriComponents): void {
+		const uri = URI.revive(rawUri);
+		const notebook = this._notebookService.getNotebookTextModel(uri);
+		if (!notebook) {
+			throw new Error(`Notebook not found: ${uri.toString()}`);
+		}
+
+		const kernel = this._notebookKernelService.getMatchingKernel(notebook);
+		if (!kernel.selected || kernel.selected.id !== controllerId) {
+			throw new Error(`Kernel is not selected: ${kernel.selected?.id} !== ${controllerId}`);
+		}
+		const execution = this._notebookExecutionStateService.createExecution(uri);
+		execution.confirm();
+		this._notebookExecutions.set(handle, execution);
+	}
+
+	$beginNotebookExecution(handle: number): void {
+		try {
+			const execution = this._notebookExecutions.get(handle);
+			execution?.begin();
+		} catch (e) {
+			onUnexpectedError(e);
+		}
+	}
+
+	$completeNotebookExecution(handle: number): void {
+		try {
+			const execution = this._notebookExecutions.get(handle);
+			execution?.complete();
+		} catch (e) {
+			onUnexpectedError(e);
+		} finally {
+			this._notebookExecutions.delete(handle);
+		}
+	}
+
 	// --- notebook kernel detection task
 	async $addKernelDetectionTask(handle: number, notebookType: string): Promise<void> {
 		const kernelDetectionTask = new MainThreadKernelDetectionTask(notebookType);

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

@@ -1056,6 +1056,10 @@ export interface MainThreadNotebookKernelsShape extends IDisposable {
 	$updateExecution(handle: number, data: SerializableObjectWithBuffers<ICellExecuteUpdateDto[]>): void;
 	$completeExecution(handle: number, data: SerializableObjectWithBuffers<ICellExecutionCompleteDto>): void;
 
+	$createNotebookExecution(handle: number, controllerId: string, uri: UriComponents): void;
+	$beginNotebookExecution(handle: number,): void;
+	$completeNotebookExecution(handle: number): void;
+
 	$addKernelDetectionTask(handle: number, notebookType: string): Promise<void>;
 	$removeKernelDetectionTask(handle: number): void;
 

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

@@ -17,7 +17,7 @@ import { ExtHostNotebookKernelsShape, ICellExecuteUpdateDto, IMainContext, INote
 import { ApiCommand, ApiCommandArgument, ApiCommandResult, ExtHostCommands } from 'vs/workbench/api/common/extHostCommands';
 import { IExtHostInitDataService } from 'vs/workbench/api/common/extHostInitDataService';
 import { ExtHostNotebookController } from 'vs/workbench/api/common/extHostNotebook';
-import { ExtHostCell } from 'vs/workbench/api/common/extHostNotebookDocument';
+import { ExtHostCell, ExtHostNotebookDocument } from 'vs/workbench/api/common/extHostNotebookDocument';
 import * as extHostTypeConverters from 'vs/workbench/api/common/extHostTypeConverters';
 import { NotebookCellExecutionState as ExtHostNotebookCellExecutionState, NotebookCellOutput, NotebookControllerAffinity2 } from 'vs/workbench/api/common/extHostTypes';
 import { asWebviewUri } from 'vs/workbench/contrib/webview/common/webview';
@@ -44,6 +44,7 @@ export class ExtHostNotebookKernels implements ExtHostNotebookKernelsShape {
 
 	private readonly _proxy: MainThreadNotebookKernelsShape;
 	private readonly _activeExecutions = new ResourceMap<NotebookCellExecutionTask>();
+	private readonly _activeNotebookExecutions = new ResourceMap<[NotebookExecutionTask, IDisposable]>();
 
 	private _kernelDetectionTask = new Map<number, vscode.NotebookControllerDetectionTask>();
 	private _kernelDetectionTaskHandlePool: number = 0;
@@ -219,6 +220,17 @@ export class ExtHostNotebookKernels implements ExtHostNotebookKernelsShape {
 				}
 				return that._createNotebookCellExecution(cell, createKernelId(extension.identifier, this.id));
 			},
+			createNotebookExecution(notebook) {
+				checkProposedApiEnabled(extension, 'notebookExecution');
+				if (isDisposed) {
+					throw new Error('notebook controller is DISPOSED');
+				}
+				if (!associatedNotebooks.has(notebook.uri)) {
+					that._logService.trace(`NotebookController[${handle}] NOT associated to notebook, associated to THESE notebooks:`, Array.from(associatedNotebooks.keys()).map(u => u.toString()));
+					throw new Error(`notebook controller is NOT associated to notebook: ${notebook.uri.toString()}`);
+				}
+				return that._createNotebookExecution(notebook, createKernelId(extension.identifier, this.id));
+			},
 			dispose: () => {
 				if (!isDisposed) {
 					this._logService.trace(`NotebookController[${handle}], DISPOSED`);
@@ -387,6 +399,14 @@ export class ExtHostNotebookKernels implements ExtHostNotebookKernelsShape {
 				}
 			}
 		}
+
+		if (obj.controller.interruptHandler) {
+			// If we're interrupting all cells, we also need to cancel the notebook level execution.
+			const items = this._activeNotebookExecutions.get(document.uri);
+			if (handles.length && Array.isArray(items) && items.length) {
+				items.forEach(d => d.dispose());
+			}
+		}
 	}
 
 	$acceptKernelMessageFromRenderer(handle: number, editorId: string, message: any): void {
@@ -439,6 +459,32 @@ export class ExtHostNotebookKernels implements ExtHostNotebookKernelsShape {
 		});
 		return execution.asApiObject();
 	}
+
+	// ---
+
+	_createNotebookExecution(nb: vscode.NotebookDocument, controllerId: string): vscode.NotebookExecution {
+		const notebook = this._extHostNotebook.getNotebookDocument(nb.uri);
+		const runningCell = nb.getCells().find(cell => {
+			const apiCell = notebook.getCellFromApiCell(cell);
+			return apiCell && this._activeExecutions.has(apiCell.uri);
+		});
+		if (runningCell) {
+			throw new Error(`duplicate cell execution for ${runningCell.document.uri}`);
+		}
+		if (this._activeNotebookExecutions.has(notebook.uri)) {
+			throw new Error(`duplicate notebook execution for ${notebook.uri}`);
+		}
+		const execution = new NotebookExecutionTask(controllerId, notebook, this._proxy);
+		const listener = execution.onDidChangeState(() => {
+			if (execution.state === NotebookExecutionTaskState.Resolved) {
+				execution.dispose();
+				listener.dispose();
+				this._activeNotebookExecutions.delete(notebook.uri);
+			}
+		});
+		this._activeNotebookExecutions.set(notebook.uri, [execution, listener]);
+		return execution.asApiObject();
+	}
 }
 
 
@@ -622,6 +668,68 @@ class NotebookCellExecutionTask extends Disposable {
 	}
 }
 
+
+enum NotebookExecutionTaskState {
+	Init,
+	Started,
+	Resolved
+}
+
+
+class NotebookExecutionTask extends Disposable {
+	private static HANDLE = 0;
+	private _handle = NotebookExecutionTask.HANDLE++;
+
+	private _onDidChangeState = new Emitter<void>();
+	readonly onDidChangeState = this._onDidChangeState.event;
+
+	private _state = NotebookExecutionTaskState.Init;
+	get state(): NotebookExecutionTaskState { return this._state; }
+
+	private readonly _tokenSource = this._register(new CancellationTokenSource());
+
+	constructor(
+		controllerId: string,
+		private readonly _notebook: ExtHostNotebookDocument,
+		private readonly _proxy: MainThreadNotebookKernelsShape
+	) {
+		super();
+
+		this._proxy.$createNotebookExecution(this._handle, controllerId, this._notebook.uri);
+	}
+
+	cancel(): void {
+		this._tokenSource.cancel();
+	}
+	asApiObject(): vscode.NotebookExecution {
+		const result: vscode.NotebookExecution = {
+			start: () => {
+				if (this._state === NotebookExecutionTaskState.Resolved || this._state === NotebookExecutionTaskState.Started) {
+					throw new Error('Cannot call start again');
+				}
+
+				this._state = NotebookExecutionTaskState.Started;
+				this._onDidChangeState.fire();
+
+				this._proxy.$beginNotebookExecution(this._handle);
+			},
+
+			end: () => {
+				if (this._state === NotebookExecutionTaskState.Resolved) {
+					throw new Error('Cannot call resolve twice');
+				}
+
+				this._state = NotebookExecutionTaskState.Resolved;
+				this._onDidChangeState.fire();
+
+				this._proxy.$completeNotebookExecution(this._handle);
+			},
+
+		};
+		return Object.freeze(result);
+	}
+}
+
 class TimeoutBasedCollector<T> {
 	private batch: T[] = [];
 	private startedTimer = Date.now();

src/vs/workbench/contrib/interactive/browser/interactiveEditor.ts

@@ -52,7 +52,7 @@ import { MarkerController } from 'vs/editor/contrib/gotoError/browser/gotoError'
 import { EditorInput } from 'vs/workbench/common/editor/editorInput';
 import { ITextResourceConfigurationService } from 'vs/editor/common/services/textResourceConfiguration';
 import { ITextEditorOptions, TextEditorSelectionSource } from 'vs/platform/editor/common/editor';
-import { INotebookExecutionStateService } from 'vs/workbench/contrib/notebook/common/notebookExecutionStateService';
+import { INotebookExecutionStateService, NotebookExecutionType } from 'vs/workbench/contrib/notebook/common/notebookExecutionStateService';
 import { NOTEBOOK_KERNEL } from 'vs/workbench/contrib/notebook/common/notebookContextKeys';
 import { ICursorPositionChangedEvent } from 'vs/editor/common/cursorEvents';
 import { IExtensionService } from 'vs/workbench/services/extensions/common/extensions';
@@ -153,8 +153,8 @@ export class InteractiveEditor extends EditorPane {
 
 		codeEditorService.registerDecorationType('interactive-decoration', DECORATION_KEY, {});
 		this._register(this.#keybindingService.onDidUpdateKeybindings(this.#updateInputDecoration, this));
-		this._register(this.#notebookExecutionStateService.onDidChangeCellExecution((e) => {
-			if (isEqual(e.notebook, this.#notebookWidget.value?.viewModel?.notebookDocument.uri)) {
+		this._register(this.#notebookExecutionStateService.onDidChangeExecution((e) => {
+			if (e.type === NotebookExecutionType.cell && isEqual(e.notebook, this.#notebookWidget.value?.viewModel?.notebookDocument.uri)) {
 				const cell = this.#notebookWidget.value?.getCellByHandle(e.cellHandle);
 				if (cell && e.changed?.state) {
 					this.#scrollIfNecessary(cell);

src/vs/workbench/contrib/notebook/browser/contrib/cellStatusBar/executionStatusBarItemController.ts

@@ -16,7 +16,7 @@ import { registerNotebookContribution } from 'vs/workbench/contrib/notebook/brow
 import { cellStatusIconError, cellStatusIconSuccess } from 'vs/workbench/contrib/notebook/browser/notebookEditorWidget';
 import { errorStateIcon, executingStateIcon, pendingStateIcon, successStateIcon } from 'vs/workbench/contrib/notebook/browser/notebookIcons';
 import { CellStatusbarAlignment, INotebookCellStatusBarItem, NotebookCellExecutionState, NotebookCellInternalMetadata } from 'vs/workbench/contrib/notebook/common/notebookCommon';
-import { INotebookCellExecution, INotebookExecutionStateService } from 'vs/workbench/contrib/notebook/common/notebookExecutionStateService';
+import { INotebookCellExecution, INotebookExecutionStateService, NotebookExecutionType } from 'vs/workbench/contrib/notebook/common/notebookExecutionStateService';
 import { INotebookService } from 'vs/workbench/contrib/notebook/common/notebookService';
 import { IMarkdownString } from 'vs/base/common/htmlContent';
 
@@ -112,8 +112,8 @@ class ExecutionStateCellStatusBarItem extends Disposable {
 		super();
 
 		this._update();
-		this._register(this._executionStateService.onDidChangeCellExecution(e => {
-			if (e.affectsCell(this._cell.uri)) {
+		this._register(this._executionStateService.onDidChangeExecution(e => {
+			if (e.type === NotebookExecutionType.cell && e.affectsCell(this._cell.uri)) {
 				this._update();
 			}
 		}));

src/vs/workbench/contrib/notebook/browser/contrib/debug/notebookDebugDecorations.ts

@@ -13,7 +13,7 @@ import { INotebookCellDecorationOptions, INotebookDeltaDecoration, INotebookEdit
 import { registerNotebookContribution } from 'vs/workbench/contrib/notebook/browser/notebookEditorExtensions';
 import { runningCellRulerDecorationColor } from 'vs/workbench/contrib/notebook/browser/notebookEditorWidget';
 import { CellUri, NotebookCellExecutionState } from 'vs/workbench/contrib/notebook/common/notebookCommon';
-import { INotebookExecutionStateService } from 'vs/workbench/contrib/notebook/common/notebookExecutionStateService';
+import { INotebookExecutionStateService, NotebookExecutionType } from 'vs/workbench/contrib/notebook/common/notebookExecutionStateService';
 
 interface ICellAndRange {
 	handle: number;
@@ -36,8 +36,8 @@ export class PausedCellDecorationContribution extends Disposable implements INot
 
 		this._register(_debugService.getModel().onDidChangeCallStack(() => this.updateExecutionDecorations()));
 		this._register(_debugService.getViewModel().onDidFocusStackFrame(() => this.updateExecutionDecorations()));
-		this._register(_notebookExecutionStateService.onDidChangeCellExecution(e => {
-			if (this._notebookEditor.textModel && e.affectsNotebook(this._notebookEditor.textModel.uri)) {
+		this._register(_notebookExecutionStateService.onDidChangeExecution(e => {
+			if (e.type === NotebookExecutionType.cell && this._notebookEditor.textModel && e.affectsNotebook(this._notebookEditor.textModel.uri)) {
 				this.updateExecutionDecorations();
 			}
 		}));

src/vs/workbench/contrib/notebook/browser/contrib/execute/executionEditorProgress.ts

@@ -21,7 +21,7 @@ export class ExecutionEditorProgressController extends Disposable implements INo
 
 		this._register(_notebookEditor.onDidScroll(() => this._update()));
 
-		this._register(_notebookExecutionStateService.onDidChangeCellExecution(e => {
+		this._register(_notebookExecutionStateService.onDidChangeExecution(e => {
 			if (e.notebook.toString() !== this._notebookEditor.textModel?.uri.toString()) {
 				return;
 			}
@@ -40,8 +40,9 @@ export class ExecutionEditorProgressController extends Disposable implements INo
 
 		const scrollPadding = this._notebookEditor.notebookOptions.computeTopInsertToolbarHeight(this._notebookEditor.textModel.viewType);
 
-		const executing = this._notebookExecutionStateService.getCellExecutionsForNotebook(this._notebookEditor.textModel?.uri)
+		const cellExecutions = this._notebookExecutionStateService.getCellExecutionsForNotebook(this._notebookEditor.textModel?.uri)
 			.filter(exe => exe.state === NotebookCellExecutionState.Executing);
+		const notebookExecution = this._notebookExecutionStateService.getExecution(this._notebookEditor.textModel?.uri);
 		const executionIsVisible = (exe: INotebookCellExecution) => {
 			for (const range of this._notebookEditor.visibleRanges) {
 				for (const cell of this._notebookEditor.getCellsInRange(range)) {
@@ -56,7 +57,9 @@ export class ExecutionEditorProgressController extends Disposable implements INo
 
 			return false;
 		};
-		if (!executing.length || executing.some(executionIsVisible) || executing.some(e => e.isPaused)) {
+		const noCellsRunning = !cellExecutions.length || cellExecutions.some(executionIsVisible) || cellExecutions.some(e => e.isPaused);
+		const isSomethingRunning = !!notebookExecution || !noCellsRunning;
+		if (isSomethingRunning) {
 			this._notebookEditor.hideProgress();
 		} else {
 			this._notebookEditor.showProgress();

src/vs/workbench/contrib/notebook/browser/contrib/outline/notebookOutline.ts

@@ -35,7 +35,7 @@ import { IdleValue } from 'vs/base/common/async';
 import { IConfigurationService } from 'vs/platform/configuration/common/configuration';
 import { IConfigurationRegistry, Extensions as ConfigurationExtensions } from 'vs/platform/configuration/common/configurationRegistry';
 import { renderMarkdownAsPlaintext } from 'vs/base/browser/markdownRenderer';
-import { INotebookExecutionStateService } from 'vs/workbench/contrib/notebook/common/notebookExecutionStateService';
+import { INotebookExecutionStateService, NotebookExecutionType } from 'vs/workbench/contrib/notebook/common/notebookExecutionStateService';
 import { executingStateIcon } from 'vs/workbench/contrib/notebook/browser/notebookIcons';
 import { URI } from 'vs/base/common/uri';
 import { getMarkdownHeadersInCell } from 'vs/workbench/contrib/notebook/browser/viewModel/foldingModel';
@@ -353,8 +353,8 @@ export class NotebookCellOutline implements IOutline<OutlineEntry> {
 			this._onDidChange.fire({});
 		}));
 
-		this._dispoables.add(_notebookExecutionStateService.onDidChangeCellExecution(e => {
-			if (!!this._editor.textModel && e.affectsNotebook(this._editor.textModel?.uri)) {
+		this._dispoables.add(_notebookExecutionStateService.onDidChangeExecution(e => {
+			if (e.type === NotebookExecutionType.cell && !!this._editor.textModel && e.affectsNotebook(this._editor.textModel?.uri)) {
 				this._recomputeState();
 			}
 		}));