Add icon to restart kernel for VSC Notebooks (microsoft#12686)

For #10496

Existing icon used for restart icon

package.json

@@ -673,7 +673,12 @@
             {
                 "command": "python.datascience.notebookeditor.restartkernel",
                 "title": "%python.command.python.datascience.restartkernel.title%",
-                "category": "Python"
+                "category": "Python",
+                "icon": {
+                    "light": "resources/light/restart-kernel.svg",
+                    "dark": "resources/dark/restart-kernel.svg"
+                },
+                "enablement": "python.datascience.notebookeditor.canrestartNotebookkernel"
             },
             {
                 "command": "python.datascience.notebookeditor.runallcells",
@@ -834,6 +839,12 @@
                     "title": "%python.command.python.execInTerminal.title%",
                     "group": "navigation",
                     "when": "resourceLangId == python && python.showPlayIcon"
+                },
+                {
+                    "command": "python.datascience.notebookeditor.restartkernel",
+                    "title": "%python.command.python.datascience.restartkernel.title%",
+                    "group": "navigation",
+                    "when": "notebookEditorFocused"
                 }
             ],
             "explorer/context": [

src/client/common/application/notebook.ts

@@ -14,8 +14,7 @@ import type {
     NotebookOutputSelector
 } from 'vscode-proposed';
 import { UseProposedApi } from '../constants';
-import { NotebookEditorSupport } from '../experiments/groups';
-import { IDisposableRegistry, IExperimentsManager } from '../types';
+import { IDisposableRegistry } from '../types';
 import {
     IApplicationEnvironment,
     IVSCodeNotebook,
@@ -79,14 +78,9 @@ export class VSCodeNotebook implements IVSCodeNotebook {
     constructor(
         @inject(UseProposedApi) private readonly useProposedApi: boolean,
         @inject(IDisposableRegistry) private readonly disposables: IDisposableRegistry,
-        @inject(IExperimentsManager) readonly experimentManager: IExperimentsManager,
         @inject(IApplicationEnvironment) readonly env: IApplicationEnvironment
     ) {
-        if (
-            this.useProposedApi &&
-            experimentManager.inExperiment(NotebookEditorSupport.nativeNotebookExperiment) &&
-            this.env.channel === 'insiders'
-        ) {
+        if (this.useProposedApi && this.env.channel === 'insiders') {
             this.addEventHandlers();
             this.canUseNotebookApi = true;
         }

src/client/datascience/constants.ts

@@ -119,6 +119,7 @@ export namespace EditorContexts {
     export const IsPythonOrInteractiveOrNativeActive = 'python.datascience.ispythonorinteractiveornativeeactive';
     export const HaveCellSelected = 'python.datascience.havecellselected';
     export const IsNotebookTrusted = 'python.datascience.isnotebooktrusted';
+    export const CanRestartNotebookKernel = 'python.datascience.notebookeditor.canrestartNotebookkernel';
 }
 
 export namespace RegExpValues {

src/client/datascience/context/activeEditorContext.ts

@@ -5,19 +5,23 @@
 
 import { inject, injectable } from 'inversify';
 import { TextEditor } from 'vscode';
+import { ServerStatus } from '../../../datascience-ui/interactive-common/mainState';
 import { IExtensionSingleActivationService } from '../../activation/types';
 import { ICommandManager, IDocumentManager } from '../../common/application/types';
 import { PYTHON_LANGUAGE } from '../../common/constants';
 import { ContextKey } from '../../common/contextKey';
 import { NotebookEditorSupport } from '../../common/experiments/groups';
+import { traceError } from '../../common/logger';
 import { IDisposable, IDisposableRegistry, IExperimentsManager } from '../../common/types';
 import { setSharedProperty } from '../../telemetry';
 import { EditorContexts } from '../constants';
 import {
     IInteractiveWindow,
     IInteractiveWindowProvider,
+    INotebook,
     INotebookEditor,
     INotebookEditorProvider,
+    INotebookProvider,
     ITrustService
 } from '../types';
 
@@ -30,6 +34,7 @@ export class ActiveEditorContextService implements IExtensionSingleActivationSer
     private pythonOrInteractiveContext: ContextKey;
     private pythonOrNativeContext: ContextKey;
     private pythonOrInteractiveOrNativeContext: ContextKey;
+    private canRestartNotebookKernelContext: ContextKey;
     private hasNativeNotebookCells: ContextKey;
     private isNotebookTrusted: ContextKey;
     private isPythonFileActive: boolean = false;
@@ -40,10 +45,15 @@ export class ActiveEditorContextService implements IExtensionSingleActivationSer
         @inject(ICommandManager) private readonly commandManager: ICommandManager,
         @inject(IDisposableRegistry) disposables: IDisposableRegistry,
         @inject(IExperimentsManager) private readonly experiments: IExperimentsManager,
+        @inject(INotebookProvider) private readonly notebookProvider: INotebookProvider,
         @inject(ITrustService) private readonly trustService: ITrustService
     ) {
         disposables.push(this);
         this.nativeContext = new ContextKey(EditorContexts.IsNativeActive, this.commandManager);
+        this.canRestartNotebookKernelContext = new ContextKey(
+            EditorContexts.CanRestartNotebookKernel,
+            this.commandManager
+        );
         this.interactiveContext = new ContextKey(EditorContexts.IsInteractiveActive, this.commandManager);
         this.interactiveOrNativeContext = new ContextKey(
             EditorContexts.IsInteractiveOrNativeActive,
@@ -66,6 +76,7 @@ export class ActiveEditorContextService implements IExtensionSingleActivationSer
     }
     public async activate(): Promise<void> {
         this.docManager.onDidChangeActiveTextEditor(this.onDidChangeActiveTextEditor, this, this.disposables);
+        this.notebookProvider.onSessionStatusChanged(this.onDidKernelStatusChange, this, this.disposables);
         this.interactiveProvider.onDidChangeActiveInteractiveWindow(
             this.onDidChangeActiveInteractiveWindow,
             this,
@@ -84,7 +95,7 @@ export class ActiveEditorContextService implements IExtensionSingleActivationSer
         }
     }
 
-    private udpateNativeNotebookCellContext() {
+    private updateNativeNotebookCellContext() {
         if (!this.experiments.inExperiment(NotebookEditorSupport.nativeNotebookExperiment)) {
             return;
         }
@@ -101,13 +112,44 @@ export class ActiveEditorContextService implements IExtensionSingleActivationSer
         // This is temporary, and once we ship native editor this needs to be removed.
         setSharedProperty('ds_notebookeditor', e?.type);
         this.nativeContext.set(!!e).ignoreErrors();
-        this.isNotebookTrusted.set(e?.model === undefined ? false : e.model.isTrusted).ignoreErrors(); // Update the currently active notebook's trust state
+        this.isNotebookTrusted.set(e?.model?.isTrusted === true).ignoreErrors();
         this.updateMergedContexts();
+        this.updateContextOfActiveNotebookKernel(e);
+    }
+    private updateContextOfActiveNotebookKernel(activeEditor?: INotebookEditor) {
+        if (activeEditor) {
+            this.notebookProvider
+                .getOrCreateNotebook({ identity: activeEditor.file, getOnly: true })
+                .then((nb) => {
+                    if (activeEditor === this.notebookEditorProvider.activeEditor) {
+                        const canStart = nb && nb.status !== ServerStatus.NotStarted;
+                        this.canRestartNotebookKernelContext.set(!!canStart).ignoreErrors();
+                    }
+                })
+                .catch(
+                    traceError.bind(undefined, 'Failed to determine if a notebook is active for the current editor')
+                );
+        } else {
+            this.canRestartNotebookKernelContext.set(false).ignoreErrors();
+        }
+    }
+    private onDidKernelStatusChange({ notebook }: { status: ServerStatus; notebook: INotebook }) {
+        // Ok, kernel status has changed.
+        const activeEditor = this.notebookEditorProvider.activeEditor;
+        if (!activeEditor) {
+            return;
+        }
+        if (activeEditor.file.toString() !== notebook.identity.toString()) {
+            // Status of a notebook thats not related to active editor has changed.
+            // We can ignore that.
+            return;
+        }
+        this.updateContextOfActiveNotebookKernel(activeEditor);
     }
     private onDidChangeActiveTextEditor(e?: TextEditor) {
         this.isPythonFileActive =
             e?.document.languageId === PYTHON_LANGUAGE && !this.notebookEditorProvider.activeEditor;
-        this.udpateNativeNotebookCellContext();
+        this.updateNativeNotebookCellContext();
         this.updateMergedContexts();
     }
     // When trust service says trust has changed, update context with whether the currently active notebook is trusted

src/client/datascience/interactive-common/notebookProvider.ts

@@ -5,6 +5,7 @@
 
 import { inject, injectable } from 'inversify';
 import { EventEmitter, Uri } from 'vscode';
+import { ServerStatus } from '../../../datascience-ui/interactive-common/mainState';
 import { IWorkspaceService } from '../../common/application/types';
 import { IFileSystem } from '../../common/platform/types';
 import { IDisposableRegistry, Resource } from '../../common/types';
@@ -28,16 +29,20 @@ import {
 export class NotebookProvider implements INotebookProvider {
     private readonly notebooks = new Map<string, Promise<INotebook>>();
     private _notebookCreated = new EventEmitter<{ identity: Uri; notebook: INotebook }>();
+    private readonly _onSessionStatusChanged = new EventEmitter<{ status: ServerStatus; notebook: INotebook }>();
     private _connectionMade = new EventEmitter<void>();
     private _type: 'jupyter' | 'raw' = 'jupyter';
     public get activeNotebooks() {
         return [...this.notebooks.values()];
     }
+    public get onSessionStatusChanged() {
+        return this._onSessionStatusChanged.event;
+    }
     constructor(
         @inject(IFileSystem) private readonly fs: IFileSystem,
         @inject(INotebookEditorProvider) private readonly editorProvider: INotebookEditorProvider,
         @inject(IInteractiveWindowProvider) private readonly interactiveWindowProvider: IInteractiveWindowProvider,
-        @inject(IDisposableRegistry) disposables: IDisposableRegistry,
+        @inject(IDisposableRegistry) private readonly disposables: IDisposableRegistry,
         @inject(IRawNotebookProvider) private readonly rawNotebookProvider: IRawNotebookProvider,
         @inject(IJupyterNotebookProvider) private readonly jupyterNotebookProvider: IJupyterNotebookProvider,
         @inject(IWorkspaceService) private readonly workspaceService: IWorkspaceService,
@@ -92,19 +97,20 @@ export class NotebookProvider implements INotebookProvider {
     public async getOrCreateNotebook(options: GetNotebookOptions): Promise<INotebook | undefined> {
         const rawKernel = await this.rawNotebookProvider.supported();
 
+        // Check our own promise cache
+        if (this.notebooks.get(options.identity.fsPath)) {
+            return this.notebooks.get(options.identity.fsPath)!!;
+        }
+
         // Check to see if our provider already has this notebook
         const notebook = rawKernel
             ? await this.rawNotebookProvider.getNotebook(options.identity, options.token)
             : await this.jupyterNotebookProvider.getNotebook(options);
         if (notebook) {
+            this.cacheNotebookPromise(options.identity, Promise.resolve(notebook));
             return notebook;
         }
 
-        // Next check our own promise cache
-        if (this.notebooks.get(options.identity.fsPath)) {
-            return this.notebooks.get(options.identity.fsPath)!!;
-        }
-
         // If get only, don't create a notebook
         if (options.getOnly) {
             return undefined;
@@ -165,6 +171,11 @@ export class NotebookProvider implements INotebookProvider {
             .then((nb) => {
                 // If the notebook is disposed, remove from cache.
                 nb.onDisposed(removeFromCache);
+                nb.onSessionStatusChanged(
+                    (e) => this._onSessionStatusChanged.fire({ status: e, notebook: nb }),
+                    this,
+                    this.disposables
+                );
                 this._notebookCreated.fire({ identity: identity, notebook: nb });
             })
             .catch(noop);

src/client/datascience/notebook/executionService.ts

@@ -186,7 +186,10 @@ export class NotebookExecutionService implements INotebookExecutionService {
         }
 
         const editor = this.editorProvider.editors.find((e) => e.model === model);
-        if (!(editor instanceof NotebookEditor)) {
+        if (!editor) {
+            throw new Error('No editor for Model');
+        }
+        if (editor && !(editor instanceof NotebookEditor)) {
             throw new Error('Executing Notebook with another Editor');
         }
         // If we need to cancel this execution (from our code, due to kernel restarts or similar, then cancel).

src/client/datascience/notebook/notebookEditor.ts

@@ -166,6 +166,7 @@ export class NotebookEditor implements INotebookEditor {
                 const no = DataScience.restartKernelMessageNo();
                 const v = await this.applicationShell.showInformationMessage(message, yes, no);
                 if (v === yes) {
+                    this.restartingKernel = false;
                     await this.restartKernel();
                 }
             }
@@ -178,13 +179,10 @@ export class NotebookEditor implements INotebookEditor {
         }
     }
 
-    public async restartKernel(internal: boolean = false): Promise<void> {
+    public async restartKernel(): Promise<void> {
         this.executionService.cancelPendingExecutions(this.document);
 
-        // Only log this if it's user requested restart
-        if (!internal) {
-            sendTelemetryEvent(Telemetry.RestartKernelCommand);
-        }
+        sendTelemetryEvent(Telemetry.RestartKernelCommand);
         if (this.restartingKernel) {
             return;
         }
@@ -195,28 +193,22 @@ export class NotebookEditor implements INotebookEditor {
         });
 
         if (notebook && !this.restartingKernel) {
-            this.restartingKernel = true;
-
-            try {
-                if (await this.shouldAskForRestart()) {
-                    // Ask the user if they want us to restart or not.
-                    const message = DataScience.restartKernelMessage();
-                    const yes = DataScience.restartKernelMessageYes();
-                    const dontAskAgain = DataScience.restartKernelMessageDontAskAgain();
-                    const no = DataScience.restartKernelMessageNo();
+            if (await this.shouldAskForRestart()) {
+                // Ask the user if they want us to restart or not.
+                const message = DataScience.restartKernelMessage();
+                const yes = DataScience.restartKernelMessageYes();
+                const dontAskAgain = DataScience.restartKernelMessageDontAskAgain();
+                const no = DataScience.restartKernelMessageNo();
 
-                    const v = await this.applicationShell.showInformationMessage(message, yes, dontAskAgain, no);
-                    if (v === dontAskAgain) {
-                        await this.disableAskForRestart();
-                        await this.restartKernelInternal(notebook);
-                    } else if (v === yes) {
-                        await this.restartKernelInternal(notebook);
-                    }
-                } else {
+                const response = await this.applicationShell.showInformationMessage(message, yes, dontAskAgain, no);
+                if (response === dontAskAgain) {
+                    await this.disableAskForRestart();
+                    await this.restartKernelInternal(notebook);
+                } else if (response === yes) {
                     await this.restartKernelInternal(notebook);
                 }
-            } finally {
-                this.restartingKernel = false;
+            } else {
+                await this.restartKernelInternal(notebook);
             }
         }
     }
@@ -229,7 +221,11 @@ export class NotebookEditor implements INotebookEditor {
         // Set our status
         const status = this.statusProvider.set(DataScience.restartingKernelStatus(), true, undefined, undefined);
 
+        // Disable running cells.
+        const [cellRunnable, runnable] = [this.document.metadata.cellRunnable, this.document.metadata.runnable];
         try {
+            this.document.metadata.cellRunnable = false;
+            this.document.metadata.runnable = false;
             await notebook.restartKernel(
                 this.configurationService.getSettings(this.file).datascience.jupyterInterruptTimeout
             );
@@ -252,6 +248,8 @@ export class NotebookEditor implements INotebookEditor {
         } finally {
             status.dispose();
             this.restartingKernel = false;
+            // Restore previous state.
+            [this.document.metadata.cellRunnable, this.document.metadata.runnable] = [cellRunnable, runnable];
         }
     }
     private async shouldAskForRestart(): Promise<boolean> {

src/client/datascience/notebook/notebookEditorProvider.ts

@@ -149,11 +149,10 @@ export class NotebookEditorProvider implements INotebookEditorProvider {
     }
 
     private closedEditor(editor: INotebookEditor): void {
-        if (!(editor instanceof NotebookEditor)) {
-            throw new Error('Executing Notebook with another Editor');
+        if (this.openedEditors.has(editor)) {
+            this.openedEditors.delete(editor);
+            this._onDidCloseNotebookEditor.fire(editor);
         }
-        this.openedEditors.delete(editor);
-        this._onDidCloseNotebookEditor.fire(editor);
     }
 
     private async onDidOpenNotebookDocument(doc: NotebookDocument): Promise<void> {
@@ -188,10 +187,15 @@ export class NotebookEditorProvider implements INotebookEditorProvider {
         const deferred = this.notebooksWaitingToBeOpenedByUri.get(uri.toString())!;
         deferred.resolve(editor);
         this.notebookEditorsByUri.set(uri.toString(), editor);
-        this.onDidChangeActiveVsCodeNotebookEditor(this.vscodeNotebook.activeNotebookEditor);
     }
     private onDidChangeActiveVsCodeNotebookEditor(editor: VSCodeNotebookEditor | undefined) {
-        if (!editor || this.trackedVSCodeNotebookEditors.has(editor)) {
+        if (!editor) {
+            this._onDidChangeActiveNotebookEditor.fire(undefined);
+            return;
+        }
+        if (this.trackedVSCodeNotebookEditors.has(editor)) {
+            const ourEditor = this.editors.find((item) => item.file.toString() === editor.document.uri.toString());
+            this._onDidChangeActiveNotebookEditor.fire(ourEditor);
             return;
         }
         this.trackedVSCodeNotebookEditors.add(editor);