Enable cell cancellation for Interactive notebooks (#19005) (#19045)

* Also added a check to prevent multiple Interactive cells from executing simultaneously.

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

@@ -173,6 +173,10 @@ export class ExtHostNotebook implements ExtHostNotebookShape {
 	}
 
 	$requestExecute(kernelId: number, content: azdata.nb.IExecuteRequest, disposeOnDone?: boolean): Thenable<INotebookFutureDetails> {
+		// Revive request's URIs to restore functions
+		content.notebookUri = URI.revive(content.notebookUri);
+		content.cellUri = URI.revive(content.cellUri);
+
 		let kernel = this._getAdapter<azdata.nb.IKernel>(kernelId);
 		let future = kernel.requestExecute(content, disposeOnDone);
 		let futureId = this._addNewAdapter(future);
@@ -259,9 +263,17 @@ export class ExtHostNotebook implements ExtHostNotebookShape {
 		return this.registerSerializationProvider(serializationProvider);
 	}
 
-	createNotebookController(extension: IExtensionDescription, id: string, viewType: string, label: string, handler?: (cells: vscode.NotebookCell[], notebook: vscode.NotebookDocument, controller: vscode.NotebookController) => void | Thenable<void>, rendererScripts?: vscode.NotebookRendererScript[]): vscode.NotebookController {
+	createNotebookController(
+		extension: IExtensionDescription,
+		id: string,
+		viewType: string,
+		label: string,
+		getDocHandler: (notebookUri: URI) => azdata.nb.NotebookDocument,
+		execHandler?: (cells: vscode.NotebookCell[], notebook: vscode.NotebookDocument, controller: vscode.NotebookController) => void | Thenable<void>,
+		rendererScripts?: vscode.NotebookRendererScript[]
+	): vscode.NotebookController {
 		let languagesHandler = (languages: string[]) => this._proxy.$updateKernelLanguages(viewType, viewType, languages);
-		let controller = new ADSNotebookController(extension, id, viewType, label, this._extHostNotebookDocumentsAndEditors, languagesHandler, handler, extension.enableProposedApi ? rendererScripts : undefined);
+		let controller = new ADSNotebookController(extension, id, viewType, label, this._extHostNotebookDocumentsAndEditors, languagesHandler, getDocHandler, execHandler, extension.enableProposedApi ? rendererScripts : undefined);
 		let newKernel: azdata.nb.IStandardKernel = {
 			name: viewType,
 			displayName: controller.label,

src/sql/workbench/api/common/notebooks/adsNotebookController.ts

@@ -4,6 +4,7 @@
  *--------------------------------------------------------------------------------------------*/
 
 import type * as vscode from 'vscode';
+import type * as azdata from 'azdata';
 import { IExtensionDescription } from 'vs/platform/extensions/common/extensions';
 import { INotebookKernelDto2 } from 'vs/workbench/api/common/extHost.protocol';
 import { Emitter, Event } from 'vs/base/common/event';
@@ -14,13 +15,14 @@ import { URI } from 'vs/base/common/uri';
 import { NotebookCellExecutionTaskState } from 'vs/workbench/api/common/extHostNotebookKernels';
 import { asArray } from 'vs/base/common/arrays';
 import { convertToADSCellOutput } from 'sql/workbench/api/common/notebooks/notebookUtils';
-import { CancellationToken } from 'vs/base/common/cancellation';
+import { CancellationTokenSource } from 'vs/base/common/cancellation';
 
 type SelectionChangedEvent = { selected: boolean, notebook: vscode.NotebookDocument; };
 type MessageReceivedEvent = { editor: vscode.NotebookEditor, message: any; };
 type ExecutionHandler = (cells: vscode.NotebookCell[], notebook: vscode.NotebookDocument, controller: vscode.NotebookController) => void | Thenable<void>;
 type LanguagesHandler = (languages: string[]) => void;
 type InterruptHandler = (notebook: vscode.NotebookDocument) => void | Promise<void>;
+type GetDocHandler = (notebookUri: URI) => azdata.nb.NotebookDocument;
 
 /**
  * A VS Code Notebook Controller that is used as part of converting VS Code notebook extension APIs into ADS equivalents.
@@ -35,6 +37,8 @@ export class ADSNotebookController implements vscode.NotebookController {
 	private readonly _languagesAdded = new Deferred<void>();
 	private readonly _executionHandlerAdded = new Deferred<void>();
 
+	private readonly _execMap: Map<string, ADSNotebookCellExecution> = new Map();
+
 	constructor(
 		private _extension: IExtensionDescription,
 		private _id: string,
@@ -42,7 +46,8 @@ export class ADSNotebookController implements vscode.NotebookController {
 		private _label: string,
 		private _extHostNotebookDocumentsAndEditors: ExtHostNotebookDocumentsAndEditors,
 		private _languagesHandler: LanguagesHandler,
-		private _handler?: ExecutionHandler,
+		private _getDocHandler: GetDocHandler,
+		private _execHandler?: ExecutionHandler,
 		preloads?: vscode.NotebookRendererScript[]
 	) {
 		this._kernelData = {
@@ -53,7 +58,7 @@ export class ADSNotebookController implements vscode.NotebookController {
 			label: this._label || this._extension.identifier.value,
 			preloads: preloads ? preloads.map(extHostTypeConverters.NotebookRendererScript.from) : []
 		};
-		if (this._handler) {
+		if (this._execHandler) {
 			this._executionHandlerAdded.resolve();
 		}
 	}
@@ -125,11 +130,11 @@ export class ADSNotebookController implements vscode.NotebookController {
 	}
 
 	public get executeHandler(): ExecutionHandler {
-		return this._handler;
+		return this._execHandler;
 	}
 
 	public set executeHandler(value: ExecutionHandler) {
-		this._handler = value;
+		this._execHandler = value;
 		this._executionHandlerAdded.resolve();
 	}
 
@@ -142,8 +147,22 @@ export class ADSNotebookController implements vscode.NotebookController {
 		this._kernelData.supportsInterrupt = Boolean(value);
 	}
 
+	public getCellExecution(cellUri: URI): ADSNotebookCellExecution | undefined {
+		return this._execMap.get(cellUri.toString());
+	}
+
+	public removeCellExecution(cellUri: URI): void {
+		this._execMap.delete(cellUri.toString());
+	}
+
+	public getNotebookDocument(notebookUri: URI): azdata.nb.NotebookDocument {
+		return this._getDocHandler(notebookUri);
+	}
+
 	public createNotebookCellExecution(cell: vscode.NotebookCell): vscode.NotebookCellExecution {
-		return new ADSNotebookCellExecution(cell, this._extHostNotebookDocumentsAndEditors);
+		let exec = new ADSNotebookCellExecution(cell, this._extHostNotebookDocumentsAndEditors);
+		this._execMap.set(cell.document.uri.toString(), exec);
+		return exec;
 	}
 
 	public dispose(): void {
@@ -166,6 +185,7 @@ export class ADSNotebookController implements vscode.NotebookController {
 class ADSNotebookCellExecution implements vscode.NotebookCellExecution {
 	private _executionOrder: number;
 	private _state = NotebookCellExecutionTaskState.Init;
+	private _cancellationSource = new CancellationTokenSource();
 	constructor(private readonly _cell: vscode.NotebookCell, private readonly _extHostNotebookDocumentsAndEditors: ExtHostNotebookDocumentsAndEditors) {
 		this._executionOrder = this._cell.executionSummary?.executionOrder ?? -1;
 	}
@@ -174,8 +194,12 @@ class ADSNotebookCellExecution implements vscode.NotebookCellExecution {
 		return this._cell;
 	}
 
+	public get tokenSource(): vscode.CancellationTokenSource {
+		return this._cancellationSource;
+	}
+
 	public get token(): vscode.CancellationToken {
-		return CancellationToken.None;
+		return this._cancellationSource.token;
 	}
 
 	public get executionOrder(): number {

src/sql/workbench/api/common/notebooks/vscodeExecuteProvider.ts

@@ -9,6 +9,8 @@ import { ADSNotebookController } from 'sql/workbench/api/common/notebooks/adsNot
 import * as nls from 'vs/nls';
 import { convertToVSCodeNotebookCell } from 'sql/workbench/api/common/notebooks/notebookUtils';
 import { CellTypes } from 'sql/workbench/services/notebook/common/contracts';
+import { VSCodeNotebookDocument } from 'sql/workbench/api/common/notebooks/vscodeNotebookDocument';
+import { URI } from 'vs/base/common/uri';
 
 class VSCodeFuture implements azdata.nb.IFuture {
 	private _inProgress = true;
@@ -71,6 +73,7 @@ class VSCodeKernel implements azdata.nb.IKernel {
 	private readonly _name: string;
 	private readonly _info: azdata.nb.IInfoReply;
 	private readonly _kernelSpec: azdata.nb.IKernelSpec;
+	private _activeRequest: azdata.nb.IExecuteRequest;
 
 	constructor(private readonly _controller: ADSNotebookController, private readonly _options: azdata.nb.ISessionOptions) {
 		this._id = this._options.kernelId ?? (VSCodeKernel.kernelId++).toString();
@@ -136,11 +139,20 @@ class VSCodeKernel implements azdata.nb.IKernel {
 		return Promise.resolve(this.spec);
 	}
 
+	private cleanUpActiveExecution(cellUri: URI) {
+		this._activeRequest = undefined;
+		this._controller.removeCellExecution(cellUri);
+	}
+
 	requestExecute(content: azdata.nb.IExecuteRequest, disposeOnDone?: boolean): azdata.nb.IFuture {
+		if (this._activeRequest) {
+			throw new Error(nls.localize('notebookMultipleRequestsError', "Cannot execute code cell. Another cell is currently being executed."));
+		}
 		let executePromise: Promise<void>;
 		if (this._controller.executeHandler) {
 			let cell = convertToVSCodeNotebookCell(CellTypes.Code, content.cellIndex, content.cellUri, content.notebookUri, content.language ?? this._kernelSpec.language, content.code);
-			executePromise = Promise.resolve(this._controller.executeHandler([cell], cell.notebook, this._controller));
+			this._activeRequest = content;
+			executePromise = Promise.resolve(this._controller.executeHandler([cell], cell.notebook, this._controller)).then(() => this.cleanUpActiveExecution(content.cellUri));
 		} else {
 			executePromise = Promise.resolve();
 		}
@@ -154,7 +166,16 @@ class VSCodeKernel implements azdata.nb.IKernel {
 	}
 
 	public async interrupt(): Promise<void> {
-		return;
+		if (this._activeRequest) {
+			if (this._controller.interruptHandler) {
+				let doc = this._controller.getNotebookDocument(this._activeRequest.notebookUri);
+				await this._controller.interruptHandler.call(this._controller, new VSCodeNotebookDocument(doc));
+			} else {
+				let exec = this._controller.getCellExecution(this._activeRequest.cellUri);
+				exec?.tokenSource.cancel();
+			}
+			this.cleanUpActiveExecution(this._activeRequest.cellUri);
+		}
 	}
 
 	public async restart(): Promise<void> {

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

@@ -1158,7 +1158,8 @@ export function createApiFactoryAndRegisterActors(accessor: ServicesAccessor, ex
 		const notebooks: typeof vscode.notebooks = {
 			createNotebookController(id: string, notebookType: string, label: string, handler?, rendererScripts?: vscode.NotebookRendererScript[]) {
 				// {{SQL CARBON EDIT}} Use our own notebooks
-				return extHostNotebook.createNotebookController(extension, id, notebookType, label, handler, extension.enableProposedApi ? rendererScripts : undefined);
+				let getDocHandler = (notebookUri: URI) => extHostNotebookDocumentsAndEditors.getDocument(notebookUri.toString())?.document;
+				return extHostNotebook.createNotebookController(extension, id, notebookType, label, getDocHandler, handler, extension.enableProposedApi ? rendererScripts : undefined);
 			},
 			registerNotebookCellStatusBarItemProvider: (notebookType: string, provider: vscode.NotebookCellStatusBarItemProvider) => {
 				// {{SQL CARBON EDIT}} Disable VS Code notebooks