Restart kernel initial implementation (#18835)

* Restart kernel initial implementation

* Update notebook extension TestKernel

* PR comments

extensions/notebook/src/jupyter/jupyterKernel.ts

@@ -109,6 +109,10 @@ export class JupyterKernel implements nb.IKernel {
 	interrupt(): Promise<void> {
 		return this.kernelImpl.interrupt();
 	}
+
+	restart(): Promise<void> {
+		return this.kernelImpl.restart();
+	}
 }
 
 export class JupyterFuture implements nb.IFuture {

extensions/notebook/src/test/common.ts

@@ -308,6 +308,9 @@ export class TestKernel implements azdata.nb.IKernel {
 	interrupt(): Thenable<void> {
 		throw new Error('Method not implemented.');
 	}
+	restart(): Thenable<void> {
+		throw new Error('Method not implemented.');
+	}
 }
 //#endregion
 

src/sql/azdata.proposed.d.ts

@@ -128,6 +128,24 @@ declare module 'azdata' {
 		 * An event that is emitted when a [notebook document](#NotebookDocument) is closed.
 		 */
 		export const onDidCloseNotebookDocument: vscode.Event<NotebookDocument>;
+
+		export interface IKernel {
+
+			/**
+			 * Restart a kernel.
+			 *
+			 * #### Notes
+			 * Uses the [Jupyter Notebook API](http://petstore.swagger.io/?url=https://raw.githubusercontent.com/jupyter/notebook/4.x/notebook/services/api/api.yaml#!/kernels).
+			 *
+			 * The promise is fulfilled on a valid response and rejected otherwise.
+			 *
+			 * It is assumed that the API call does not mutate the kernel id or name.
+			 *
+			 * The promise will be rejected if the kernel status is `Dead` or if the
+			 * request fails or the response is invalid.
+			 */
+			restart(): Thenable<void>;
+		}
 	}
 
 	/**

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

@@ -484,6 +484,10 @@ class KernelWrapper implements azdata.nb.IKernel {
 	interrupt(): Thenable<void> {
 		return this._proxy.ext.$interruptKernel(this.kernelDetails.kernelId);
 	}
+
+	restart(): Thenable<void> {
+		return this._proxy.ext.$restartKernel(this.kernelDetails.kernelId);
+	}
 }
 
 

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

@@ -213,6 +213,11 @@ export class ExtHostNotebook implements ExtHostNotebookShape {
 		return kernel.interrupt();
 	}
 
+	$restartKernel(kernelId: number): Thenable<void> {
+		let kernel = this._getAdapter<azdata.nb.IKernel>(kernelId);
+		return kernel.restart();
+	}
+
 	$sendInputReply(futureId: number, content: azdata.nb.IInputReply): void {
 		let future = this._getAdapter<azdata.nb.IFuture>(futureId);
 		return future.sendInputReply(content);

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

@@ -156,6 +156,10 @@ class VSCodeKernel implements azdata.nb.IKernel {
 	public async interrupt(): Promise<void> {
 		return;
 	}
+
+	public async restart(): Promise<void> {
+		return;
+	}
 }
 
 class VSCodeSession implements azdata.nb.ISession {

src/sql/workbench/api/common/sqlExtHost.protocol.ts

@@ -952,6 +952,7 @@ export interface ExtHostNotebookShape {
 	$requestComplete(kernelId: number, content: azdata.nb.ICompleteRequest): Thenable<azdata.nb.ICompleteReplyMsg>;
 	$requestExecute(kernelId: number, content: azdata.nb.IExecuteRequest, disposeOnDone?: boolean): Thenable<INotebookFutureDetails>;
 	$interruptKernel(kernelId: number): Thenable<void>;
+	$restartKernel(kernelId: number): Thenable<void>;
 
 	// Future APIs
 	$sendInputReply(futureId: number, content: azdata.nb.IInputReply): void;

src/sql/workbench/contrib/notebook/browser/notebook.contribution.ts

@@ -239,6 +239,24 @@ CommandsRegistry.registerCommand({
 	}
 });
 
+CommandsRegistry.registerCommand({
+	id: 'notebook.restartKernel',
+	handler: async (accessor: ServicesAccessor) => {
+		const editorService: IEditorService = accessor.get(IEditorService);
+		if (editorService.activeEditor instanceof NotebookInput) {
+			await editorService.activeEditor.notebookModel?.clientSession?.restart();
+		}
+	}
+});
+
+MenuRegistry.appendMenuItem(MenuId.CommandPalette, {
+	command: {
+		id: 'notebook.restartKernel',
+		title: localize('restartNotebookKernel', "Restart Notebook Kernel"),
+	},
+	when: ContextKeyExpr.and(ActiveEditorContext.isEqualTo(NotebookEditor.ID))
+});
+
 MenuRegistry.appendMenuItem(MenuId.CommandPalette, {
 	command: {
 		id: TOGGLE_TAB_FOCUS_COMMAND_ID,

src/sql/workbench/contrib/notebook/test/emptySessionClasses.ts

@@ -168,6 +168,10 @@ class EmptyKernel implements nb.IKernel {
 	interrupt(): Thenable<void> {
 		return Promise.resolve(undefined);
 	}
+
+	restart(): Thenable<void> {
+		return Promise.resolve(undefined);
+	}
 }
 
 export class EmptyFuture implements FutureInternal {

src/sql/workbench/contrib/notebook/test/stubs.ts

@@ -353,7 +353,7 @@ export class ClientSessionStub implements IClientSession {
 	selectKernel(): Promise<void> {
 		throw new Error('Method not implemented.');
 	}
-	restart(): Promise<boolean> {
+	restart(): Promise<void> {
 		throw new Error('Method not implemented.');
 	}
 	setPath(path: string): Promise<void> {
@@ -464,6 +464,9 @@ export class KernelStub implements nb.IKernel {
 	interrupt(): Thenable<void> {
 		throw new Error('Method not implemented.');
 	}
+	restart(): Thenable<void> {
+		throw new Error('Method not implemented.');
+	}
 }
 
 export class FutureStub implements nb.IFuture {

src/sql/workbench/services/notebook/browser/models/clientSession.ts

@@ -323,16 +323,27 @@ export class ClientSession implements IClientSession {
 	/**
 	 * Restart the session.
 	 *
-	 * @returns A promise that resolves with whether the kernel has restarted.
+	 * @returns A promise that resolves when the kernel has restarted.
 	 *
 	 * #### Notes
-	 * If there is a running kernel, present a dialog.
+	 * If there is an existing kernel, restart it and resolve.
-	 * If there is no kernel, we start a kernel with the last run
+	 * If no kernel has been started, this is a no-op, and resolves.
-	 * kernel name and resolves with `true`. If no kernel has been started,
+	 * Reject on error.
-	 * this is a no-op, and resolves with `false`.
 	 */
-	restart(): Promise<boolean> {
+	restart(): Promise<void> {
-		throw new Error('Not implemented');
+		if (!this._session?.kernel) {
+			// no-op if no kernel is present
+			return Promise.resolve();
+		}
+		let restartCompleted = new Deferred<void>();
+		this._session?.kernel?.restart().then(() => {
+			this.options.notificationService.info(localize('kernelRestartedSuccessfully', 'Kernel restarted successfully'));
+			restartCompleted.resolve();
+		}, err => {
+			this.options.notificationService.error(localize('kernelRestartFailed', 'Kernel restart failed: {0}', err));
+			restartCompleted.reject(err);
+		});
+		return restartCompleted.promise;
 	}
 
 	/**

src/sql/workbench/services/notebook/browser/models/modelInterfaces.ts

@@ -179,15 +179,14 @@ export interface IClientSession extends IDisposable {
 	/**
 	 * Restart the session.
 	 *
-	 * @returns A promise that resolves with whether the kernel has restarted.
+	 * @returns A promise that resolves when the kernel has restarted.
 	 *
 	 * #### Notes
-	 * If there is a running kernel, present a dialog.
+	 * If there is an existing kernel, restart it and resolve.
-	 * If there is no kernel, we start a kernel with the last run
+	 * If no kernel has been started, this is a no-op, and resolves.
-	 * kernel name and resolves with `true`. If no kernel has been started,
+	 * Reject on error.
-	 * this is a no-op, and resolves with `false`.
 	 */
-	restart(): Promise<boolean>;
+	restart(): Promise<void>;
 
 	/**
 	 * Change the session path.

src/sql/workbench/services/notebook/browser/sql/sqlSessionManager.ts

@@ -382,6 +382,10 @@ class SqlKernel extends Disposable implements nb.IKernel {
 		});
 	}
 
+	restart(): Thenable<void> {
+		return Promise.reject(localize('SqlKernelRestartNotSupported', 'SQL kernel restart not supported'));
+	}
+
 	private addQueryEventListeners(queryRunner: QueryRunner): void {
 		this._register(queryRunner.onQueryEnd(() => {
 			this.queryComplete().catch(error => {

src/sql/workbench/test/electron-browser/api/mainThreadNotebook.test.ts

@@ -252,6 +252,9 @@ class ExtHostNotebookStub implements ExtHostNotebookShape {
 	$interruptKernel(kernelId: number): Thenable<void> {
 		throw new Error('Method not implemented.');
 	}
+	$restartKernel(kernelId: number): Thenable<void> {
+		throw new Error('Method not implemented.');
+	}
 	$sendInputReply(futureId: number, content: azdata.nb.IInputReply): void {
 		throw new Error('Method not implemented.');
 	}