Add icon to restart kernel for VSC Notebooks

package.json

@@ -673,7 +673,11 @@
             {
                 "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"
+                }
             },
             {
                 "command": "python.datascience.notebookeditor.runallcells",
@@ -834,6 +838,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 && python.datascience.notebookeditor.canrestartNotebookkernel"
                 }
             ],
             "explorer/context": [

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,
@@ -103,6 +114,34 @@ export class ActiveEditorContextService implements IExtensionSingleActivationSer
         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.updateMergedContexts();
+        this.updateContextOfActiveNotebookKernel(e);
+    }
+    private updateContextOfActiveNotebookKernel(activeEditor?: INotebookEditor) {
+        if (activeEditor) {
+            this.notebookProvider
+                .getOrCreateNotebook({ identity: activeEditor.file, getOnly: true })
+                .then((nb) => {
+                    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.
+        }
+        this.updateContextOfActiveNotebookKernel(activeEditor);
     }
     private onDidChangeActiveTextEditor(e?: TextEditor) {
         this.isPythonFileActive =

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/notebookEditorProvider.ts

@@ -191,7 +191,13 @@ export class NotebookEditorProvider implements INotebookEditorProvider {
         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);

src/client/datascience/types.ts

@@ -1057,8 +1057,6 @@ type WebViewViewState = {
 };
 export type WebViewViewChangeEventArgs = { current: WebViewViewState; previous: WebViewViewState };
 
-export const INotebookProvider = Symbol('INotebookProvider');
-
 export type GetServerOptions = {
     getOnly?: boolean;
     disableUI?: boolean;
@@ -1079,8 +1077,10 @@ export type GetNotebookOptions = {
     token?: CancellationToken;
 };
 
+export const INotebookProvider = Symbol('INotebookProvider');
 export interface INotebookProvider {
     readonly type: 'raw' | 'jupyter';
+    onSessionStatusChanged: Event<{ status: ServerStatus; notebook: INotebook }>;
     /**
      * Fired when a notebook has been created for a given Uri/Identity
      */