Fix save on close

news/2 Fixes/11711.md

@@ -0,0 +1 @@
+Fix saving during close and auto backup to actually save a notebook.

src/client/common/application/commands.ts

@@ -195,8 +195,8 @@ export interface ICommandNameArgumentTypeMapping extends ICommandNameWithoutArgu
     [DSCommands.SetJupyterKernel]: [KernelConnectionMetadata, Uri, undefined | Uri];
     [DSCommands.SwitchJupyterKernel]: [ISwitchKernelOptions | undefined];
     [DSCommands.SelectJupyterCommandLine]: [undefined | Uri];
-    [DSCommands.SaveNotebookNonCustomEditor]: [Uri];
-    [DSCommands.SaveAsNotebookNonCustomEditor]: [Uri, Uri];
+    [DSCommands.SaveNotebookNonCustomEditor]: [INotebookModel];
+    [DSCommands.SaveAsNotebookNonCustomEditor]: [INotebookModel, Uri];
     [DSCommands.OpenNotebookNonCustomEditor]: [Uri];
     [DSCommands.GatherQuality]: [string];
     [DSCommands.LatestExtension]: [string];

src/client/datascience/gather/gatherLogger.ts

@@ -3,7 +3,7 @@ import { inject, injectable } from 'inversify';
 import cloneDeep = require('lodash/cloneDeep');
 import { extensions } from 'vscode';
 import { concatMultilineString } from '../../../datascience-ui/common';
-import { traceError } from '../../common/logger';
+import { traceError, traceInfo } from '../../common/logger';
 import { IConfigurationService } from '../../common/types';
 import { noop } from '../../common/utils/misc';
 import { sendTelemetryEvent } from '../../telemetry';
@@ -69,7 +69,11 @@ export class GatherLogger implements IGatherLogger {
                 }
             }
             const api = ext.exports;
-            this.gather = api.getGatherProvider();
+            try {
+                this.gather = api.getGatherProvider();
+            } catch {
+                traceInfo(`Gather not installed`);
+            }
         }
     }
 }

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

@@ -269,7 +269,7 @@ export class NativeEditorOldWebView extends NativeEditor {
         }
         try {
             if (!this.isUntitled) {
-                await this.commandManager.executeCommand(Commands.SaveNotebookNonCustomEditor, this.model?.file);
+                await this.commandManager.executeCommand(Commands.SaveNotebookNonCustomEditor, this.model);
                 this.savedEvent.fire(this);
                 return;
             }
@@ -295,7 +295,7 @@ export class NativeEditorOldWebView extends NativeEditor {
             if (fileToSaveTo) {
                 await this.commandManager.executeCommand(
                     Commands.SaveAsNotebookNonCustomEditor,
-                    this.model.file,
+                    this.model,
                     fileToSaveTo
                 );
                 this.savedEvent.fire(this);

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

@@ -83,7 +83,7 @@ export class NativeEditorProviderOld extends NativeEditorProvider {
         @inject(IWorkspaceService) workspace: IWorkspaceService,
         @inject(IConfigurationService) configuration: IConfigurationService,
         @inject(ICustomEditorService) customEditorService: ICustomEditorService,
-        @inject(IDataScienceFileSystem) private fs: IDataScienceFileSystem,
+        @inject(IDataScienceFileSystem) fs: IDataScienceFileSystem,
         @inject(IDocumentManager) private documentManager: IDocumentManager,
         @inject(ICommandManager) private readonly cmdManager: ICommandManager,
         @inject(IDataScienceErrorHandler) private dataScienceErrorHandler: IDataScienceErrorHandler,
@@ -98,7 +98,8 @@ export class NativeEditorProviderOld extends NativeEditorProvider {
             configuration,
             customEditorService,
             storage,
-            notebookProvider
+            notebookProvider,
+            fs
         );
 
         // No live share sync required as open document from vscode will give us our contents.
@@ -107,21 +108,18 @@ export class NativeEditorProviderOld extends NativeEditorProvider {
             this.documentManager.onDidChangeActiveTextEditor(this.onDidChangeActiveTextEditorHandler.bind(this))
         );
         this.disposables.push(
-            this.cmdManager.registerCommand(Commands.SaveNotebookNonCustomEditor, async (resource: Uri) => {
-                const customDocument = this.customDocuments.get(resource.fsPath);
-                if (customDocument) {
-                    await this.saveCustomDocument(customDocument, new CancellationTokenSource().token);
-                }
+            this.cmdManager.registerCommand(Commands.SaveNotebookNonCustomEditor, async (model: INotebookModel) => {
+                await this.storage.save(model, new CancellationTokenSource().token);
             })
         );
         this.disposables.push(
             this.cmdManager.registerCommand(
                 Commands.SaveAsNotebookNonCustomEditor,
-                async (resource: Uri, targetResource: Uri) => {
-                    const customDocument = this.customDocuments.get(resource.fsPath);
+                async (model: INotebookModel, targetResource: Uri) => {
+                    await this.storage.saveAs(model, targetResource);
+                    const customDocument = this.customDocuments.get(model.file.fsPath);
                     if (customDocument) {
-                        await this.saveCustomDocumentAs(customDocument, targetResource);
-                        this.customDocuments.delete(resource.fsPath);
+                        this.customDocuments.delete(model.file.fsPath);
                         this.customDocuments.set(targetResource.fsPath, { ...customDocument, uri: targetResource });
                     }
                 }

src/client/datascience/notebookStorage/nativeEditorProvider.ts

@@ -108,7 +108,8 @@ export class NativeEditorProvider implements INotebookEditorProvider, CustomEdit
         @inject(IConfigurationService) protected readonly configuration: IConfigurationService,
         @inject(ICustomEditorService) private customEditorService: ICustomEditorService,
         @inject(INotebookStorageProvider) protected readonly storage: INotebookStorageProvider,
-        @inject(INotebookProvider) private readonly notebookProvider: INotebookProvider
+        @inject(INotebookProvider) private readonly notebookProvider: INotebookProvider,
+        @inject(IDataScienceFileSystem) protected readonly fs: IDataScienceFileSystem
     ) {
         traceInfo(`id is ${this._id}`);
 
@@ -214,14 +215,18 @@ export class NativeEditorProvider implements INotebookEditorProvider, CustomEdit
     public async loadModel(file: Uri, contents?: string, backupId?: string): Promise<INotebookModel>;
     // tslint:disable-next-line: no-any
     public async loadModel(file: Uri, contents?: string, options?: any): Promise<INotebookModel> {
-        // Every time we load a new untitled file, up the counter past the max value for this counter
-        this.untitledCounter = getNextUntitledCounter(file, this.untitledCounter);
+        // Get the model that may match this file
+        let model = [...this.models.values()].find((m) => this.fs.arePathsSame(m.file, file));
+        if (!model) {
+            // Every time we load a new untitled file, up the counter past the max value for this counter
+            this.untitledCounter = getNextUntitledCounter(file, this.untitledCounter);
 
-        // Load our model from our storage object.
-        const model = await this.storage.getOrCreateModel(file, contents, options);
+            // Load our model from our storage object.
+            model = await this.storage.getOrCreateModel(file, contents, options);
 
-        // Make sure to listen to events on the model
-        this.trackModel(model);
+            // Make sure to listen to events on the model
+            this.trackModel(model);
+        }
         return model;
     }
 

src/client/datascience/notebookStorage/nativeEditorStorage.ts

@@ -26,7 +26,7 @@ import {
 import detectIndent = require('detect-indent');
 import { VSCodeNotebookModel } from './vscNotebookModel';
 
-const KeyPrefix = 'notebook-storage-';
+export const KeyPrefix = 'notebook-storage-';
 const NotebookTransferKey = 'notebook-transfered';
 
 export function getNextUntitledCounter(file: Uri | undefined, currentValue: number): number {
@@ -178,21 +178,29 @@ export class NativeEditorStorage implements INotebookStorage {
         // Keep track of the time when this data was saved.
         // This way when we retrieve the data we can compare it against last modified date of the file.
         const specialContents = contents ? JSON.stringify({ contents, lastModifiedTimeMs: Date.now() }) : undefined;
-        return this.writeToStorage(filePath, specialContents, cancelToken);
+        return this.writeToStorage(model.file, filePath, specialContents, cancelToken);
     }
 
     private async clearHotExit(file: Uri, backupId?: string): Promise<void> {
         const key = backupId || this.getStaticStorageKey(file);
         const filePath = this.getHashedFileName(key);
-        await this.writeToStorage(filePath);
+        await this.writeToStorage(undefined, filePath);
     }
 
-    private async writeToStorage(filePath: string, contents?: string, cancelToken?: CancellationToken): Promise<void> {
+    private async writeToStorage(
+        owningFile: Uri | undefined,
+        filePath: string,
+        contents?: string,
+        cancelToken?: CancellationToken
+    ): Promise<void> {
         try {
             if (!cancelToken?.isCancellationRequested) {
                 if (contents) {
                     await this.fs.createLocalDirectory(path.dirname(filePath));
                     if (!cancelToken?.isCancellationRequested) {
+                        if (owningFile) {
+                            this.trustService.trustNotebook(owningFile, contents).ignoreErrors();
+                        }
                         await this.fs.writeLocalFile(filePath, contents);
                     }
                 } else {
@@ -374,6 +382,8 @@ export class NativeEditorStorage implements INotebookStorage {
             if (isNotebookTrusted) {
                 model.trust();
             }
+        } else {
+            model.trust();
         }
 
         return model;
@@ -407,9 +417,10 @@ export class NativeEditorStorage implements INotebookStorage {
     }
 
     private async getStoredContentsFromFile(file: Uri, key: string): Promise<string | undefined> {
+        const filePath = this.getHashedFileName(key);
         try {
             // Use this to read from the extension global location
-            const contents = await this.fs.readLocalFile(file.fsPath);
+            const contents = await this.fs.readLocalFile(filePath);
             const data = JSON.parse(contents);
             // Check whether the file has been modified since the last time the contents were saved.
             if (data && data.lastModifiedTimeMs && file.scheme === 'file') {

src/datascience-ui/interactive-common/redux/store.ts

@@ -31,6 +31,7 @@ import { forceLoad } from '../transforms';
 import { isAllowedAction, isAllowedMessage, postActionToExtension } from './helpers';
 import { generatePostOfficeSendReducer } from './postOffice';
 import { generateMonacoReducer, IMonacoState } from './reducers/monaco';
+import { CommonActionType } from './reducers/types';
 import { generateVariableReducer, IVariableState } from './reducers/variables';
 
 function generateDefaultState(
@@ -109,19 +110,21 @@ function createSendInfoMiddleware(): Redux.Middleware<{}, IStore> {
         }
 
         // If cell vm count changed or selected cell changed, send the message
-        const currentSelection = getSelectedAndFocusedInfo(afterState.main);
-        if (
-            prevState.main.cellVMs.length !== afterState.main.cellVMs.length ||
-            getSelectedAndFocusedInfo(prevState.main).selectedCellId !== currentSelection.selectedCellId ||
-            prevState.main.undoStack.length !== afterState.main.undoStack.length ||
-            prevState.main.redoStack.length !== afterState.main.redoStack.length
-        ) {
-            postActionToExtension({ queueAction: store.dispatch }, InteractiveWindowMessages.SendInfo, {
-                cellCount: afterState.main.cellVMs.length,
-                undoCount: afterState.main.undoStack.length,
-                redoCount: afterState.main.redoStack.length,
-                selectedCell: currentSelection.selectedCellId
-            });
+        if (!action.type || action.type !== CommonActionType.UNMOUNT) {
+            const currentSelection = getSelectedAndFocusedInfo(afterState.main);
+            if (
+                prevState.main.cellVMs.length !== afterState.main.cellVMs.length ||
+                getSelectedAndFocusedInfo(prevState.main).selectedCellId !== currentSelection.selectedCellId ||
+                prevState.main.undoStack.length !== afterState.main.undoStack.length ||
+                prevState.main.redoStack.length !== afterState.main.redoStack.length
+            ) {
+                postActionToExtension({ queueAction: store.dispatch }, InteractiveWindowMessages.SendInfo, {
+                    cellCount: afterState.main.cellVMs.length,
+                    undoCount: afterState.main.undoStack.length,
+                    redoCount: afterState.main.redoStack.length,
+                    selectedCell: currentSelection.selectedCellId
+                });
+            }
         }
         return res;
     };
@@ -159,21 +162,26 @@ function createTestMiddleware(): Redux.Middleware<{}, IStore> {
             });
         };
 
-        // Special case for focusing a cell
-        const previousSelection = getSelectedAndFocusedInfo(prevState.main);
-        const currentSelection = getSelectedAndFocusedInfo(afterState.main);
-        if (previousSelection.focusedCellId !== currentSelection.focusedCellId && currentSelection.focusedCellId) {
-            // Send async so happens after render state changes (so our enzyme wrapper is up to date)
-            sendMessage(InteractiveWindowMessages.FocusedCellEditor, { cellId: action.payload.cellId });
-        }
-        if (previousSelection.selectedCellId !== currentSelection.selectedCellId && currentSelection.selectedCellId) {
-            // Send async so happens after render state changes (so our enzyme wrapper is up to date)
-            sendMessage(InteractiveWindowMessages.SelectedCell, { cellId: action.payload.cellId });
-        }
-        // Special case for unfocusing a cell
-        if (previousSelection.focusedCellId !== currentSelection.focusedCellId && !currentSelection.focusedCellId) {
-            // Send async so happens after render state changes (so our enzyme wrapper is up to date)
-            sendMessage(InteractiveWindowMessages.UnfocusedCellEditor);
+        if (!action.type || action.type !== CommonActionType.UNMOUNT) {
+            // Special case for focusing a cell
+            const previousSelection = getSelectedAndFocusedInfo(prevState.main);
+            const currentSelection = getSelectedAndFocusedInfo(afterState.main);
+            if (previousSelection.focusedCellId !== currentSelection.focusedCellId && currentSelection.focusedCellId) {
+                // Send async so happens after render state changes (so our enzyme wrapper is up to date)
+                sendMessage(InteractiveWindowMessages.FocusedCellEditor, { cellId: action.payload.cellId });
+            }
+            if (
+                previousSelection.selectedCellId !== currentSelection.selectedCellId &&
+                currentSelection.selectedCellId
+            ) {
+                // Send async so happens after render state changes (so our enzyme wrapper is up to date)
+                sendMessage(InteractiveWindowMessages.SelectedCell, { cellId: action.payload.cellId });
+            }
+            // Special case for unfocusing a cell
+            if (previousSelection.focusedCellId !== currentSelection.focusedCellId && !currentSelection.focusedCellId) {
+                // Send async so happens after render state changes (so our enzyme wrapper is up to date)
+                sendMessage(InteractiveWindowMessages.UnfocusedCellEditor);
+            }
         }
 
         // Indicate settings updates
@@ -218,7 +226,10 @@ function createTestMiddleware(): Redux.Middleware<{}, IStore> {
             sendMessage(InteractiveWindowMessages.ExecutionRendered);
         }
 
-        if (!action.type || action.type !== InteractiveWindowMessages.FinishCell) {
+        if (
+            !action.type ||
+            (action.type !== InteractiveWindowMessages.FinishCell && action.type !== CommonActionType.UNMOUNT)
+        ) {
             // Might be a non finish but still update cells (like an undo or a delete)
             const prevFinished = prevState.main.cellVMs
                 .filter((c) => c.cell.state === CellState.finished || c.cell.state === CellState.error)

src/test/datascience/interactive-ipynb/nativeEditorStorage.unit.test.ts

@@ -283,7 +283,9 @@ suite('DataScience - Native Editor Storage', () => {
         when(executionProvider.serverStarted).thenReturn(serverStartedEvent.event);
 
         when(trustService.isNotebookTrusted(anything(), anything())).thenReturn(Promise.resolve(true));
-        when(trustService.trustNotebook(anything(), anything())).thenReturn(Promise.resolve());
+        when(trustService.trustNotebook(anything(), anything())).thenCall(() => {
+            return Promise.resolve();
+        });
 
         testIndex += 1;
         when(crypto.createHash(anything(), 'string')).thenReturn(`${testIndex}`);
@@ -351,7 +353,7 @@ suite('DataScience - Native Editor Storage', () => {
             context.object,
             globalMemento,
             localMemento,
-            trustService,
+            instance(trustService),
             new NotebookModelFactory(false)
         );
 

src/test/datascience/mountedWebView.ts

@@ -7,7 +7,7 @@ import {
     IWebPanelOptions,
     WebPanelMessage
 } from '../../client/common/application/types';
-import { traceInfo } from '../../client/common/logger';
+import { traceError, traceInfo } from '../../client/common/logger';
 import { IDisposable } from '../../client/common/types';
 import { createDeferred } from '../../client/common/utils/async';
 import { InteractiveWindowMessages } from '../../client/datascience/interactive-common/interactiveWindowTypes';
@@ -236,6 +236,9 @@ export class MountedWebView implements IMountedWebView, IDisposable {
         }
     }
     private postMessageToWebPanel(msg: any) {
+        if (this.disposed && !msg.type.startsWith(`DISPATCHED`)) {
+            traceError(`Posting to disposed mount.`);
+        }
         if (this.webPanelListener) {
             this.webPanelListener.onMessage(msg.type, msg.payload);
         } else {