Trust for native notebooks (microsoft#14353)

Author: DonJayamanne

build/ci/scripts/spec_with_pid.js

@@ -24,6 +24,8 @@ var color = Base.color;
 
 exports = module.exports = Spec;
 
+const prefix = process.env.VSC_PYTHON_CI_TEST_PARALLEL ? `${process.pid}   ` : '';
+
 /**
  * Constructs a new `Spec` reporter instance.
  *
@@ -51,7 +53,7 @@ function Spec(runner, options) {
 
     runner.on(EVENT_SUITE_BEGIN, function (suite) {
         ++indents;
-        Base.consoleLog(color('suite', `${process.pid} %s%s`), indent(), suite.title);
+        Base.consoleLog(color('suite', `${prefix}%s%s`), indent(), suite.title);
     });
 
     runner.on(EVENT_SUITE_END, function () {
@@ -62,27 +64,27 @@ function Spec(runner, options) {
     });
 
     runner.on(EVENT_TEST_PENDING, function (test) {
-        var fmt = indent() + color('pending', `${process.pid}   - %s`);
+        var fmt = indent() + color('pending', `${prefix} %s`);
         Base.consoleLog(fmt, test.title);
     });
 
     runner.on(EVENT_TEST_PASS, function (test) {
         var fmt;
         if (test.speed === 'fast') {
-            fmt = indent() + color('checkmark', `${process.pid}   ` + Base.symbols.ok) + color('pass', ' %s');
+            fmt = indent() + color('checkmark', prefix + Base.symbols.ok) + color('pass', ' %s');
             Base.consoleLog(fmt, test.title);
         } else {
             fmt =
                 indent() +
-                color('checkmark', `${process.pid}   ` + Base.symbols.ok) +
+                color('checkmark', prefix + Base.symbols.ok) +
                 color('pass', ' %s') +
                 color(test.speed, ' (%dms)');
             Base.consoleLog(fmt, test.title, test.duration);
         }
     });
 
     runner.on(EVENT_TEST_FAIL, function (test) {
-        Base.consoleLog(indent() + color('fail', `${process.pid}  %d) %s`), ++n, test.title);
+        Base.consoleLog(indent() + color('fail', `${prefix}%d) %s`), ++n, test.title);
     });
 
     runner.once(EVENT_RUN_END, self.epilogue.bind(self));

build/ci/templates/test_phases.yml

@@ -547,17 +547,9 @@ steps:
   - bash: |
       npm install -g vsce
       npm run clean
-      npx tsc -p ./
-      mkdir -p ./tmp/client/logging
-      cp -r ./out/client/logging ./tmp/client
-      npx gulp clean:cleanExceptTests
-      cp -r ./out/test ./tmp/test
       npm run updateBuildNumber -- --buildNumber $BUILD_BUILDID
       npm run package
-      npx gulp clean:cleanExceptTests
-      mkdir -p ./out/client/logging
-      cp -r ./tmp/client/logging ./out/client
-      cp -r ./tmp/test ./out/test
+      npx tsc -p ./
       node --no-force-async-hooks-checks ./out/test/smokeTest.js
     displayName: 'Run Smoke Tests'
     condition: and(succeeded(), contains(variables['TestsToRun'], 'testSmoke'))

package.json

@@ -3474,7 +3474,7 @@
         "test:functional:perf": "node --inspect-brk ./node_modules/mocha/bin/_mocha --require source-map-support/register --config ./build/.mocha.functional.perf.json",
         "test:functional:memleak": "node --inspect-brk ./node_modules/mocha/bin/_mocha --require source-map-support/register --config ./build/.mocha.functional.json",
         "test:functional:cover": "npm run test:functional",
-        "test:functional:parallel": "node ./build/ci/scripts/runFunctionalTests.js",
+        "test:functional:parallel": "cross-env VSC_PYTHON_CI_TEST_PARALLEL=1 node ./build/ci/scripts/runFunctionalTests.js",
         "test:cover:report": "nyc --nycrc-path build/.nycrc  report --reporter=text --reporter=html --reporter=text-summary --reporter=cobertura",
         "testDebugger": "node ./out/test/testBootstrap.js ./out/test/debuggerTest.js",
         "testSingleWorkspace": "node ./out/test/testBootstrap.js ./out/test/standardTest.js",

src/client/common/application/notebook.ts

@@ -16,13 +16,7 @@ import type {
 } from 'vscode-proposed';
 import { UseProposedApi } from '../constants';
 import { IDisposableRegistry } from '../types';
-import {
-    IApplicationEnvironment,
-    IVSCodeNotebook,
-    NotebookCellLanguageChangeEvent,
-    NotebookCellOutputsChangeEvent,
-    NotebookCellsChangeEvent
-} from './types';
+import { IApplicationEnvironment, IVSCodeNotebook, NotebookCellChangedEvent } from './types';
 
 @injectable()
 export class VSCodeNotebook implements IVSCodeNotebook {
@@ -63,14 +57,10 @@ export class VSCodeNotebook implements IVSCodeNotebook {
     public get notebookEditors() {
         return this.canUseNotebookApi ? this.notebook.visibleNotebookEditors : [];
     }
-    public get onDidChangeNotebookDocument(): Event<
-        NotebookCellsChangeEvent | NotebookCellOutputsChangeEvent | NotebookCellLanguageChangeEvent
-    > {
+    public get onDidChangeNotebookDocument(): Event<NotebookCellChangedEvent> {
         return this.canUseNotebookApi
             ? this._onDidChangeNotebookDocument.event
-            : new EventEmitter<
-                  NotebookCellsChangeEvent | NotebookCellOutputsChangeEvent | NotebookCellLanguageChangeEvent
-              >().event;
+            : new EventEmitter<NotebookCellChangedEvent>().event;
     }
     public get activeNotebookEditor(): NotebookEditor | undefined {
         if (!this.useProposedApi) {
@@ -85,9 +75,7 @@ export class VSCodeNotebook implements IVSCodeNotebook {
         }
         return this._notebook!;
     }
-    private readonly _onDidChangeNotebookDocument = new EventEmitter<
-        NotebookCellsChangeEvent | NotebookCellOutputsChangeEvent | NotebookCellLanguageChangeEvent
-    >();
+    private readonly _onDidChangeNotebookDocument = new EventEmitter<NotebookCellChangedEvent>();
     private addedEventHandlers?: boolean;
     private _notebook?: typeof notebook;
     private readonly canUseNotebookApi?: boolean;
@@ -128,6 +116,12 @@ export class VSCodeNotebook implements IVSCodeNotebook {
                 this.notebook.onDidChangeCellLanguage((e) =>
                     this._onDidChangeNotebookDocument.fire({ ...e, type: 'changeCellLanguage' })
                 ),
+                this.notebook.onDidChangeCellMetadata((e) =>
+                    this._onDidChangeNotebookDocument.fire({ ...e, type: 'changeCellMetadata' })
+                ),
+                this.notebook.onDidChangeNotebookDocumentMetadata((e) =>
+                    this._onDidChangeNotebookDocument.fire({ ...e, type: 'changeNotebookMetadata' })
+                ),
                 this.notebook.onDidChangeCellOutputs((e) =>
                     this._onDidChangeNotebookDocument.fire({ ...e, type: 'changeCellOutputs' })
                 ),

src/client/common/application/types.ts

@@ -61,11 +61,13 @@ import {
 import type {
     NotebookCellLanguageChangeEvent as VSCNotebookCellLanguageChangeEvent,
     NotebookCellMetadata,
+    NotebookCellMetadataChangeEvent as VSCNotebookCellMetadataChangeEvent,
     NotebookCellOutputsChangeEvent as VSCNotebookCellOutputsChangeEvent,
     NotebookCellsChangeEvent as VSCNotebookCellsChangeEvent,
     NotebookContentProvider,
     NotebookDocument,
     NotebookDocumentFilter,
+    NotebookDocumentMetadataChangeEvent as VSCNotebookDocumentMetadataChangeEvent,
     NotebookEditor,
     NotebookKernel,
     NotebookKernelProvider
@@ -1526,10 +1528,16 @@ export interface IClipboard {
 
 export type NotebookCellsChangeEvent = { type: 'changeCells' } & VSCNotebookCellsChangeEvent;
 export type NotebookCellOutputsChangeEvent = { type: 'changeCellOutputs' } & VSCNotebookCellOutputsChangeEvent;
+export type NotebookCellMetadataChangeEvent = { type: 'changeCellMetadata' } & VSCNotebookCellMetadataChangeEvent;
 export type NotebookCellLanguageChangeEvent = { type: 'changeCellLanguage' } & VSCNotebookCellLanguageChangeEvent;
+export type NotebookDocumentMetadataChangeEvent = {
+    type: 'changeNotebookMetadata';
+} & VSCNotebookDocumentMetadataChangeEvent;
 export type NotebookCellChangedEvent =
     | NotebookCellsChangeEvent
     | NotebookCellOutputsChangeEvent
+    | NotebookCellMetadataChangeEvent
+    | NotebookDocumentMetadataChangeEvent
     | NotebookCellLanguageChangeEvent;
 export const IVSCodeNotebook = Symbol('IVSCodeNotebook');
 export interface IVSCodeNotebook {

src/client/datascience/context/activeEditorContext.ts

@@ -100,7 +100,7 @@ export class ActiveEditorContextService implements IExtensionSingleActivationSer
             return;
         }
         this.hasNativeNotebookCells
-            .set((this.notebookEditorProvider.activeEditor?.model?.cells?.length || 0) > 0)
+            .set((this.notebookEditorProvider.activeEditor?.model?.cellCount || 0) > 0)
             .ignoreErrors();
     }
     private onDidChangeActiveInteractiveWindow(e?: IInteractiveWindow) {

src/client/datascience/export/exportUtil.ts

@@ -3,7 +3,7 @@ import { inject, injectable } from 'inversify';
 import * as os from 'os';
 import * as path from 'path';
 import * as uuid from 'uuid/v4';
-import { CancellationTokenSource, Uri } from 'vscode';
+import { Uri } from 'vscode';
 import { TemporaryDirectory } from '../../common/platform/types';
 import { sleep } from '../../common/utils/async';
 import { ICell, IDataScienceFileSystem, INotebookExporter, INotebookModel, INotebookStorage } from '../types';
@@ -68,24 +68,15 @@ export class ExportUtil {
 
     public async removeSvgs(source: Uri) {
         const model = await this.notebookStorage.getOrCreateModel({ file: source });
-
-        const newCells: ICell[] = [];
-        for (const cell of model.cells) {
-            const outputs = cell.data.outputs;
-            if (outputs as nbformat.IOutput[]) {
-                this.removeSvgFromOutputs(outputs as nbformat.IOutput[]);
+        const content = JSON.parse(model.getContent()) as nbformat.INotebookContent;
+        for (const cell of content.cells) {
+            const outputs = cell.outputs as nbformat.IOutput[];
+            if (Array.isArray(outputs)) {
+                this.removeSvgFromOutputs(outputs);
             }
-            newCells.push(cell);
         }
-        model.update({
-            kind: 'modify',
-            newCells: newCells,
-            oldCells: model.cells as ICell[],
-            oldDirty: false,
-            newDirty: false,
-            source: 'user'
-        });
-        await this.notebookStorage.save(model, new CancellationTokenSource().token);
+        await this.fs.writeFile(source, JSON.stringify(content, undefined, model.indentAmount));
+        model.dispose(); // We're modifying the JSON in file manually, hence blow away cached model.
     }
 
     private removeSvgFromOutputs(outputs: nbformat.IOutput[]) {

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

@@ -67,7 +67,6 @@ import {
     INotebookExtensibility,
     INotebookImporter,
     INotebookMetadataLive,
-    INotebookModel,
     INotebookProvider,
     IStatusProvider,
     IThemeFinder,
@@ -87,6 +86,7 @@ import { IDataViewerFactory } from '../data-viewing/types';
 import { getCellHashProvider } from '../editor-integration/cellhashprovider';
 import { KernelSelector } from '../jupyter/kernels/kernelSelector';
 import { KernelConnectionMetadata } from '../jupyter/kernels/types';
+import { NativeEditorNotebookModel } from '../notebookStorage/notebookModel';
 
 const nativeEditorDir = path.join(EXTENSION_ROOT_DIR, 'out', 'datascience-ui', 'notebook');
 export class NativeEditor extends InteractiveBase implements INotebookEditor {
@@ -134,7 +134,7 @@ export class NativeEditor extends InteractiveBase implements INotebookEditor {
     public get isDirty(): boolean {
         return this.model ? this.model.isDirty : false;
     }
-    public get model(): Readonly<INotebookModel> {
+    public get model(): Readonly<NativeEditorNotebookModel> {
         return this._model;
     }
     public readonly type: 'old' | 'custom' = 'custom';
@@ -181,7 +181,7 @@ export class NativeEditor extends InteractiveBase implements INotebookEditor {
         notebookProvider: INotebookProvider,
         useCustomEditorApi: boolean,
         private trustService: ITrustService,
-        private _model: INotebookModel,
+        private _model: NativeEditorNotebookModel,
         webviewPanel: WebviewPanel | undefined,
         selector: KernelSelector,
         private nbExtensibility: INotebookExtensibility

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

@@ -29,6 +29,7 @@ import { Commands, Telemetry } from '../constants';
 import { IDataViewerFactory } from '../data-viewing/types';
 import { InteractiveWindowMessages } from '../interactive-common/interactiveWindowTypes';
 import { KernelSelector } from '../jupyter/kernels/kernelSelector';
+import { NativeEditorNotebookModel } from '../notebookStorage/notebookModel';
 import { INotebookStorageProvider } from '../notebookStorage/notebookStorageProvider';
 import {
     ICodeCssGenerator,
@@ -42,7 +43,6 @@ import {
     INotebookExporter,
     INotebookExtensibility,
     INotebookImporter,
-    INotebookModel,
     INotebookProvider,
     IStatusProvider,
     IThemeFinder,
@@ -100,7 +100,7 @@ export class NativeEditorOldWebView extends NativeEditor {
         useCustomEditorApi: boolean,
         private readonly storage: INotebookStorageProvider,
         trustService: ITrustService,
-        model: INotebookModel,
+        model: NativeEditorNotebookModel,
         webviewPanel: WebviewPanel | undefined,
         selector: KernelSelector,
         notebookExtensibility: INotebookExtensibility

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

@@ -34,6 +34,7 @@ import { IDataViewerFactory } from '../data-viewing/types';
 import { NotebookModelChange } from '../interactive-common/interactiveWindowTypes';
 import { KernelSelector } from '../jupyter/kernels/kernelSelector';
 import { NativeEditorProvider } from '../notebookStorage/nativeEditorProvider';
+import { NativeEditorNotebookModel } from '../notebookStorage/notebookModel';
 import { INotebookStorageProvider } from '../notebookStorage/notebookStorageProvider';
 import { VSCodeNotebookModel } from '../notebookStorage/vscNotebookModel';
 import {
@@ -193,7 +194,7 @@ export class NativeEditorProviderOld extends NativeEditorProvider {
         }
     }
 
-    protected createNotebookEditor(model: INotebookModel, panel?: WebviewPanel): NativeEditor {
+    protected createNotebookEditor(model: NativeEditorNotebookModel, panel?: WebviewPanel): NativeEditor {
         const editor = new NativeEditorOldWebView(
             this.serviceContainer.getAll<IInteractiveWindowListener>(IInteractiveWindowListener),
             this.serviceContainer.get<ILiveShareApi>(ILiveShareApi),

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

@@ -6,16 +6,20 @@
 import { inject, injectable } from 'inversify';
 import { commands, Uri } from 'vscode';
 import { IExtensionSingleActivationService } from '../../activation/types';
-import { IApplicationShell, ICommandManager } from '../../common/application/types';
+import { IApplicationShell, ICommandManager, IVSCodeNotebook } from '../../common/application/types';
 import { ContextKey } from '../../common/contextKey';
 import '../../common/extensions';
+import { IFileSystem } from '../../common/platform/types';
 import { IDisposableRegistry } from '../../common/types';
+import { createDeferred } from '../../common/utils/async';
 import { swallowExceptions } from '../../common/utils/decorators';
 import { DataScience } from '../../common/utils/localize';
+import { noop } from '../../common/utils/misc';
 import { sendTelemetryEvent } from '../../telemetry';
 import { Commands, Telemetry } from '../constants';
 import { INotebookStorageProvider } from '../notebookStorage/notebookStorageProvider';
-import { INotebookEditorProvider, ITrustService } from '../types';
+import { VSCodeNotebookModel } from '../notebookStorage/vscNotebookModel';
+import { INotebookEditorProvider, INotebookModel, ITrustService } from '../types';
 
 @injectable()
 export class TrustCommandHandler implements IExtensionSingleActivationService {
@@ -25,7 +29,9 @@ export class TrustCommandHandler implements IExtensionSingleActivationService {
         @inject(INotebookStorageProvider) private readonly storageProvider: INotebookStorageProvider,
         @inject(ICommandManager) private readonly commandManager: ICommandManager,
         @inject(IApplicationShell) private readonly applicationShell: IApplicationShell,
-        @inject(IDisposableRegistry) private readonly disposables: IDisposableRegistry
+        @inject(IDisposableRegistry) private readonly disposables: IDisposableRegistry,
+        @inject(IVSCodeNotebook) private readonly vscNotebook: IVSCodeNotebook,
+        @inject(IFileSystem) private readonly fs: IFileSystem
     ) {}
     public async activate(): Promise<void> {
         this.activateInBackground().ignoreErrors();
@@ -61,11 +67,11 @@ export class TrustCommandHandler implements IExtensionSingleActivationService {
                 sendTelemetryEvent(Telemetry.TrustAllNotebooks);
                 break;
             case DataScience.trustNotebook():
-                // Update model trust
-                model.trust();
-                const contents = model.getContent();
-                await this.trustService.trustNotebook(model.file, contents);
-                sendTelemetryEvent(Telemetry.TrustNotebook);
+                if (model instanceof VSCodeNotebookModel) {
+                    await this.trustNativeNotebook(model);
+                } else {
+                    await this.trustNotebook(model);
+                }
                 break;
             case DataScience.doNotTrustNotebook():
                 sendTelemetryEvent(Telemetry.DoNotTrustNotebook);
@@ -74,4 +80,52 @@ export class TrustCommandHandler implements IExtensionSingleActivationService {
                 break;
         }
     }
+    private async trustNotebook(model: INotebookModel) {
+        // Update model trust
+        model.trust();
+        const contents = model.getContent();
+        await this.trustService.trustNotebook(model.file, contents);
+        sendTelemetryEvent(Telemetry.TrustNotebook);
+    }
+    private async trustNativeNotebook(model: VSCodeNotebookModel) {
+        const trustedNotebookInTrustService = createDeferred<void>();
+        const doc = this.vscNotebook.notebookDocuments.find((item) =>
+            this.fs.arePathsSame(item.uri.fsPath, model.file.fsPath)
+        );
+        let fileReverted = false;
+        const disposable = this.vscNotebook.onDidChangeNotebookDocument((e) => {
+            if (e.document !== doc) {
+                return;
+            }
+            trustedNotebookInTrustService.promise
+                .then(() => {
+                    // Notebook has been trusted, revert the changes in the document so that we re-load the cells.
+                    // This is a hacky solution for trusting native notebooks.
+                    if (!fileReverted) {
+                        fileReverted = true;
+                        return commands.executeCommand('workbench.action.files.revert');
+                    }
+                })
+                .catch(noop);
+            disposable.dispose();
+        });
+        this.disposables.push(disposable);
+
+        try {
+            // Update model trust
+            await model.trustNotebook();
+            // Trust the original contents & contents generated vy using the VSC Notebook model.
+            // When generating JSON from VSC, sometimes contents can get auto formatted based on user settings (hence that JSON in ipynb could be different from original contents).
+            // Thus trust both, the original content & the new content generated by VSC.
+            const originalContents = model.getOriginalContentOnDisc();
+            const contents = model.getContent();
+            await Promise.all([
+                this.trustService.trustNotebook(model.file, contents),
+                this.trustService.trustNotebook(model.file, originalContents)
+            ]);
+            sendTelemetryEvent(Telemetry.TrustNotebook);
+        } finally {
+            trustedNotebookInTrustService.resolve();
+        }
+    }
 }