Fix save on close (#13567)

* Pass model through command instead of URI in order to use directly * Add a test to verify we don't regress * Finish fixing tests * Add news entry * Fix unit tests

Author: rchiodo

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 {