Add trust feature for VS Code Notebooks

build/ci/vscode-python-pr-validation.yaml

@@ -56,6 +56,7 @@ stages:
             'DataScience':
               TestsToRun: 'testDataScience'
               NeedsPythonTestReqs: true
+              NeedsPythonFunctionalReqs: true
             'Smoke':
               TestsToRun: 'testSmoke'
               NeedsPythonTestReqs: true

package.json

@@ -694,6 +694,25 @@
                 },
                 "enablement": "python.datascience.notebookeditor.canrestartNotebookkernel"
             },
+            {
+                "command": "python.datascience.notebookeditor.trust",
+                "title": "%DataScience.trustNotebookCommandTitle%",
+                "category": "Python",
+                "icon": {
+                    "light": "resources/light/un-trusted.svg",
+                    "dark": "resources/dark/un-trusted.svg"
+                },
+                "enablement": "notebookEditorFocused && !python.datascience.isnotebooktrusted && python.datascience.trustfeatureenabled"
+            },
+            {
+                "command": "python.datascience.notebookeditor.trusted",
+                "title": "%DataScience.notebookIsTrusted%",
+                "category": "Python",
+                "icon": {
+                    "light": "resources/light/trusted.svg",
+                    "dark": "resources/dark/trusted.svg"
+                }
+            },
             {
                 "command": "python.datascience.notebookeditor.runallcells",
                 "title": "%python.command.python.datascience.notebookeditor.runallcells.title%",
@@ -860,6 +879,18 @@
                     "group": "navigation",
                     "when": "notebookEditorFocused"
                 },
+                {
+                    "command": "python.datascience.notebookeditor.trust",
+                    "title": "%DataScience.trustNotebookCommandTitle%",
+                    "group": "navigation@1",
+                    "when": "notebookEditorFocused && !python.datascience.isnotebooktrusted && python.datascience.trustfeatureenabled"
+                },
+                {
+                    "command": "python.datascience.notebookeditor.trusted",
+                    "title": "%DataScience.notebookIsTrusted%",
+                    "group": "navigation@1",
+                    "when": "notebookEditorFocused && python.datascience.isnotebooktrusted && python.datascience.trustfeatureenabled"
+                },
                 {
                     "command": "python.datascience.export",
                     "title": "%DataScience.notebookExportAs%",
@@ -1138,6 +1169,18 @@
                     "category": "Python",
                     "when": "python.datascience.isnativeactive && python.datascience.featureenabled && python.datascience.isnotebooktrusted"
                 },
+                {
+                    "command": "python.datascience.notebookeditor.trust",
+                    "title": "%DataScience.trustNotebookCommandTitle%",
+                    "category": "Python",
+                    "when": "python.datascience.featureenabled && notebookEditorFocused && !python.datascience.isnotebooktrusted && python.datascience.trustfeatureenabled"
+                },
+                {
+                    "command": "python.datascience.notebookeditor.trusted",
+                    "title": "%DataScience.notebookIsTrusted%",
+                    "category": "Python",
+                    "when": "config.noExists"
+                },
                 {
                     "command": "python.datascience.notebookeditor.runallcells",
                     "title": "%python.command.python.datascience.notebookeditor.runallcells.title%",

package.nls.json

@@ -370,6 +370,7 @@
     "DataScience.fetchingDataViewer": "Fetching data ...",
     "DataScience.noRowsInDataViewer": "No rows match current filter",
     "DataScience.jupyterServer": "Jupyter Server",
+    "DataScience.trustNotebookCommandTitle": "Trust notebook",
     "DataScience.notebookIsTrusted": "Trusted",
     "DataScience.notebookIsNotTrusted": "Not Trusted",
     "DataScience.noKernel": "No Kernel",

src/client/common/application/commands.ts

@@ -183,4 +183,6 @@ export interface ICommandNameArgumentTypeMapping extends ICommandNameWithoutArgu
     [DSCommands.OpenNotebookNonCustomEditor]: [Uri];
     [DSCommands.GatherQuality]: [string];
     [DSCommands.EnableLoadingWidgetsFrom3rdPartySource]: [undefined | never];
+    [DSCommands.TrustNotebook]: [undefined | never | Uri];
+    [DSCommands.TrustedNotebook]: [undefined | never];
 }

src/client/datascience/constants.ts

@@ -90,6 +90,8 @@ export namespace Commands {
     export const SaveAsNotebookNonCustomEditor = 'python.datascience.notebookeditor.saveAs';
     export const OpenNotebookNonCustomEditor = 'python.datascience.notebookeditor.open';
     export const GatherQuality = 'python.datascience.gatherquality';
+    export const TrustNotebook = 'python.datascience.notebookeditor.trust';
+    export const TrustedNotebook = 'python.datascience.notebookeditor.trusted';
     export const EnableLoadingWidgetsFrom3rdPartySource =
         'python.datascience.enableLoadingWidgetScriptsFromThirdPartySource';
 }

src/client/datascience/context/activeEditorContext.ts

@@ -122,7 +122,7 @@ export class ActiveEditorContextService implements IExtensionSingleActivationSer
                 .getOrCreateNotebook({ identity: activeEditor.file, getOnly: true })
                 .then((nb) => {
                     if (activeEditor === this.notebookEditorProvider.activeEditor) {
-                        const canStart = nb && nb.status !== ServerStatus.NotStarted;
+                        const canStart = nb && nb.status !== ServerStatus.NotStarted && activeEditor.model?.isTrusted;
                         this.canRestartNotebookKernelContext.set(!!canStart).ignoreErrors();
                     }
                 })

src/client/datascience/export/exportUtil.ts

@@ -63,7 +63,7 @@ export class ExportUtil {
             await this.jupyterExporter.exportToFile(cells, tempFile.filePath, false);
             const newPath = path.join(tempDir.path, '.ipynb');
             await this.fileSystem.copyFile(tempFile.filePath, newPath);
-            model = await this.notebookStorage.load(Uri.file(newPath));
+            model = await this.notebookStorage.get(Uri.file(newPath));
         } finally {
             tempFile.dispose();
             tempDir.dispose();
@@ -73,7 +73,7 @@ export class ExportUtil {
     }
 
     public async removeSvgs(source: Uri) {
-        const model = await this.notebookStorage.load(source);
+        const model = await this.notebookStorage.get(source);
 
         const newCells: ICell[] = [];
         for (const cell of model.cells) {

src/client/datascience/interactive-ipynb/nativeEditor.ts

@@ -6,7 +6,6 @@ import * as path from 'path';
 import {
     CancellationToken,
     CancellationTokenSource,
-    commands,
     Event,
     EventEmitter,
     Memento,
@@ -98,7 +97,6 @@ import { KernelSwitcher } from '../jupyter/kernels/kernelSwitcher';
 const nativeEditorDir = path.join(EXTENSION_ROOT_DIR, 'out', 'datascience-ui', 'notebook');
 @injectable()
 export class NativeEditor extends InteractiveBase implements INotebookEditor {
-    public readonly type: 'old' | 'custom' = 'custom';
     public get onDidChangeViewState(): Event<void> {
         return this._onDidChangeViewState.event;
     }
@@ -140,6 +138,7 @@ export class NativeEditor extends InteractiveBase implements INotebookEditor {
     public get isDirty(): boolean {
         return this.model ? this.model.isDirty : false;
     }
+    public readonly type: 'old' | 'custom' = 'custom';
     public model: Readonly<INotebookModel> | undefined;
     protected savedEvent: EventEmitter<INotebookEditor> = new EventEmitter<INotebookEditor>();
     protected closedEvent: EventEmitter<INotebookEditor> = new EventEmitter<INotebookEditor>();
@@ -152,6 +151,7 @@ export class NativeEditor extends InteractiveBase implements INotebookEditor {
     private startupTimer: StopWatch = new StopWatch();
     private loadedAllCells: boolean = false;
     private executeCancelTokens = new Set<CancellationTokenSource>();
+    private previouslyNotTrusted: boolean = false;
 
     constructor(
         @multiInject(IInteractiveWindowListener) listeners: IInteractiveWindowListener[],
@@ -227,6 +227,7 @@ export class NativeEditor extends InteractiveBase implements INotebookEditor {
         );
         asyncRegistry.push(this);
 
+        asyncRegistry.push(this.trustService.onDidSetNotebookTrust(this.monitorChangesToTrust, this));
         this.synchronizer.subscribeToUserActions(this, this.postMessage.bind(this));
     }
 
@@ -251,6 +252,7 @@ export class NativeEditor extends InteractiveBase implements INotebookEditor {
 
         // Sign up for dirty events
         model.changed(this.modelChanged.bind(this));
+        this.previouslyNotTrusted = model.isTrusted;
     }
 
     // tslint:disable-next-line: no-any
@@ -596,7 +598,13 @@ export class NativeEditor extends InteractiveBase implements INotebookEditor {
             }
         }
     }
-
+    private async monitorChangesToTrust() {
+        if (this.previouslyNotTrusted && this.model?.isTrusted) {
+            this.previouslyNotTrusted = false;
+            // Tell UI to update main state
+            this.postMessage(InteractiveWindowMessages.TrustNotebookComplete).ignoreErrors();
+        }
+    }
     private renameVariableExplorerHeights(name: string, updatedName: string) {
         // Updates the workspace storage to reflect the updated name of the notebook
         // should be called if the name of the notebook changes
@@ -612,38 +620,8 @@ export class NativeEditor extends InteractiveBase implements INotebookEditor {
     }
 
     private async launchNotebookTrustPrompt() {
-        const prompts = [
-            localize.DataScience.trustNotebook(),
-            localize.DataScience.doNotTrustNotebook(),
-            localize.DataScience.trustAllNotebooks()
-        ];
-        const selection = await this.applicationShell.showErrorMessage(
-            localize.DataScience.launchNotebookTrustPrompt(),
-            ...prompts
-        );
-        if (!selection) {
-            return;
-        }
-        if (this.model && selection === localize.DataScience.trustNotebook() && !this.model.isTrusted) {
-            try {
-                const contents = this.model.getContent();
-                await this.trustService.trustNotebook(this.model.file, contents);
-                // Update model trust
-                this.model.update({
-                    source: 'user',
-                    kind: 'updateTrust',
-                    oldDirty: this.model.isDirty,
-                    newDirty: this.model.isDirty,
-                    isNotebookTrusted: true
-                });
-                // Tell UI to update main state
-                await this.postMessage(InteractiveWindowMessages.TrustNotebookComplete);
-            } catch (err) {
-                traceError(err);
-            }
-        } else if (selection === localize.DataScience.trustAllNotebooks()) {
-            // Take the user to the settings UI where they can manually turn on the alwaysTrustNotebooks setting
-            commands.executeCommand('workbench.action.openSettings', 'python.dataScience.alwaysTrustNotebooks');
+        if (this.model && !this.model.isTrusted) {
+            await this.commandManager.executeCommand(Commands.TrustNotebook, this.model.file);
         }
     }
 

src/client/datascience/interactive-ipynb/nativeEditorProvider.ts

@@ -184,7 +184,7 @@ export class NativeEditorProvider implements INotebookEditorProvider, CustomEdit
         this.untitledCounter = getNextUntitledCounter(file, this.untitledCounter);
 
         // Load our model from our storage object.
-        const model = await this.storage.load(file, contents, options);
+        const model = await this.storage.get(file, contents, options);
 
         // Make sure to listen to events on the model
         this.trackModel(model);

src/client/datascience/interactive-ipynb/nativeEditorStorage.ts

@@ -74,20 +74,20 @@ export class NativeEditorStorage implements INotebookStorage {
         return `${path.basename(model.file.fsPath)}-${uuid()}`;
     }
 
-    public load(
+    public get(
         file: Uri,
         possibleContents?: string,
         backupId?: string,
         forVSCodeNotebook?: boolean
     ): Promise<INotebookModel>;
-    public load(
+    public get(
         file: Uri,
         possibleContents?: string,
         // tslint:disable-next-line: unified-signatures
         skipDirtyContents?: boolean,
         forVSCodeNotebook?: boolean
     ): Promise<INotebookModel>;
-    public load(
+    public get(
         file: Uri,
         possibleContents?: string,
         // tslint:disable-next-line: no-any