handle OOB edits in notebooks

Author: dlqqq

src/docprovider/custom_ydocs.ts

@@ -1,6 +1,7 @@
 import {
-  YFile as DefaultYFile
-  // YNotebook as DefaultYNotebook
+  YFile as DefaultYFile,
+  YNotebook as DefaultYNotebook,
+  ISharedNotebook
 } from '@jupyter/ydoc';
 import * as Y from 'yjs';
 import { Awareness } from 'y-protocols/awareness';
@@ -34,14 +35,14 @@ export class YFile extends DefaultYFile {
     (this as any)._ydoc = new Y.Doc();
 
     // Reset all properties derived from `this._ydoc`
-    (this as any).ysource = (this as any)._ydoc.getText('source');
-    (this as any)._ystate = (this as any)._ydoc.getMap('state');
+    (this as any).ysource = this.ydoc.getText('source');
+    (this as any)._ystate = this.ydoc.getMap('state');
     (this as any)._undoManager = new Y.UndoManager([], {
       trackedOrigins: new Set([this]),
       doc: (this as any)._ydoc
     });
     (this as any)._undoManager.addToScope(this.ysource);
-    (this as any)._awareness = new Awareness((this as any)._ydoc);
+    (this as any)._awareness = new Awareness(this.ydoc);
 
     // Emit to `this.resetSignal` to inform consumers immediately
     this._resetSignal.emit(null);
@@ -64,3 +65,52 @@ export class YFile extends DefaultYFile {
 
   _resetSignal: Signal<this, null>;
 }
+
+export class YNotebook extends DefaultYNotebook {
+  constructor(options?: Omit<ISharedNotebook.IOptions, 'data'>) {
+    super(options);
+    this._resetSignal = new Signal(this);
+  }
+
+  /**
+   * See `YFile.reset()`.
+   */
+  reset() {
+    // Remove default observers
+    this._ycells.unobserve((this as any)._onYCellsChanged);
+    this.ymeta.unobserveDeep((this as any)._onMetaChanged);
+    (this as any)._ystate.unobserve(this.onStateChanged);
+
+    // Reset `this._ydoc` to an empty state
+    (this as any)._ydoc = new Y.Doc();
+
+    // Reset all properties derived from `this._ydoc`
+    (this as any)._ystate = this.ydoc.getMap('state');
+    (this as any)._ycells = this.ydoc.getArray('cells');
+    (this as any).cells = [];
+    (this as any).ymeta = this.ydoc.getMap('meta');
+    (this as any)._undoManager = new Y.UndoManager([], {
+      trackedOrigins: new Set([this]),
+      doc: (this as any)._ydoc
+    });
+    (this as any)._undoManager.addToScope(this._ycells);
+    (this as any)._awareness = new Awareness(this.ydoc);
+
+    // Emit to `this.resetSignal` to inform consumers immediately
+    this._resetSignal.emit(null);
+
+    // Add back default observers
+    this._ycells.observe((this as any)._onYCellsChanged);
+    this.ymeta.observeDeep((this as any)._onMetaChanged);
+    (this as any)._ystate.observe(this.onStateChanged);
+  }
+
+  /**
+   * See `YFile.resetSignal`.
+   */
+  get resetSignal(): ISignal<this, null> {
+    return this._resetSignal;
+  }
+
+  _resetSignal: Signal<this, null>;
+}

src/docprovider/filebrowser.ts

@@ -25,8 +25,7 @@ import {
 import { ISettingRegistry } from '@jupyterlab/settingregistry';
 import { ITranslator, nullTranslator } from '@jupyterlab/translation';
 
-import { YNotebook } from '@jupyter/ydoc';
-import { YFile } from './custom_ydocs';
+import { YFile, YNotebook } from './custom_ydocs';
 
 import {
   ICollaborativeContentProvider,

src/docprovider/yprovider.ts

@@ -16,7 +16,7 @@ import { DocumentChange, YDocument } from '@jupyter/ydoc';
 import { Awareness } from 'y-protocols/awareness';
 import { WebsocketProvider as YWebsocketProvider } from 'y-websocket';
 import { requestAPI } from './requests';
-import { YFile } from './custom_ydocs';
+import { YFile, YNotebook } from './custom_ydocs';
 
 /**
  * A class to provide Yjs synchronization over WebSocket.
@@ -146,7 +146,7 @@ export class WebSocketProvider implements IDocumentProvider {
     const close_code = event.code;
 
     // 4000 := server close code on out-of-band change
-    if (close_code === 4000 && this._sharedModel instanceof YFile) {
+    if (close_code === 4000) {
       this._handleOobChange();
       return;
     }
@@ -172,9 +172,8 @@ export class WebSocketProvider implements IDocumentProvider {
    */
   private _handleOobChange() {
     // Reset YDoc
-    // TODO: handle YNotebooks.
     // TODO: is it safe to assume that we only need YFile & YNotebook?
-    const sharedModel = this._sharedModel as YFile;
+    const sharedModel = this._sharedModel as YFile | YNotebook;
     sharedModel.reset();
 
     // Re-connect and display a notification to the user

src/index.ts

@@ -25,9 +25,7 @@ import { KeyboardEvent } from 'react';
 import { IToolbarWidgetRegistry } from '@jupyterlab/apputils';
 import { AwarenessExecutionIndicator } from './executionindicator';
 
-import { IEditorServices } from '@jupyterlab/codeeditor';
 import { requestAPI } from './handler';
-import { RtcNotebookContentFactory } from './notebook';
 
 import { rtcContentProvider, yfile, ynotebook, logger } from './docprovider';
 
@@ -43,6 +41,7 @@ import { URLExt } from '@jupyterlab/coreutils';
 import { AwarenessKernelStatus } from './kernelstatus';
 
 import { codemirrorYjsPlugin } from './codemirror-binding/plugin';
+import { notebookFactoryPlugin } from './notebook-factory';
 
 /**
  * Initialization data for the @jupyter/server-documents extension.
@@ -62,7 +61,10 @@ export const plugin: JupyterFrontEndPlugin<void> = {
       settingRegistry
         .load(plugin.id)
         .then(settings => {
-          console.log('@jupyter/server-documents settings loaded:', settings.composite);
+          console.log(
+            '@jupyter/server-documents settings loaded:',
+            settings.composite
+          );
         })
         .catch(reason => {
           console.error(
@@ -279,21 +281,6 @@ export const kernelStatus: JupyterFrontEndPlugin<IKernelStatusModel> = {
   }
 };
 
-/**
- * The notebook cell factory provider.
- */
-const factory: JupyterFrontEndPlugin<NotebookPanel.IContentFactory> = {
-  id: '@jupyter/server-documents/notebook-extension:factory',
-  description: 'Provides the notebook cell factory.',
-  provides: NotebookPanel.IContentFactory,
-  requires: [IEditorServices],
-  autoStart: true,
-  activate: (app: JupyterFrontEnd, editorServices: IEditorServices) => {
-    const editorFactory = editorServices.factoryService.newInlineEditor;
-    return new RtcNotebookContentFactory({ editorFactory });
-  }
-};
-
 const plugins: JupyterFrontEndPlugin<unknown>[] = [
   rtcContentProvider,
   yfile,
@@ -303,7 +290,7 @@ const plugins: JupyterFrontEndPlugin<unknown>[] = [
   plugin,
   executionIndicator,
   kernelStatus,
-  factory,
+  notebookFactoryPlugin,
   codemirrorYjsPlugin
 ];
 

src/notebook-factory/index.ts

@@ -0,0 +1,2 @@
+export * from './notebook-factory';
+export * from './plugin';

src/notebook.ts renamed to src/notebook-factory/notebook-factory.ts

@@ -1,9 +1,15 @@
-import { CodeCell, CodeCellModel, ICellModel, ICodeCellModel } from '@jupyterlab/cells';
-import { NotebookPanel } from '@jupyterlab/notebook';
+import {
+  CodeCell,
+  CodeCellModel,
+  ICellModel,
+  ICodeCellModel
+} from '@jupyterlab/cells';
+import { IChangedArgs } from '@jupyterlab/coreutils';
+import { Notebook, NotebookPanel } from '@jupyterlab/notebook';
 import { CellChange, createMutex, ISharedCodeCell } from '@jupyter/ydoc';
 import { IOutputAreaModel, OutputAreaModel } from '@jupyterlab/outputarea';
-import { IChangedArgs } from '@jupyterlab/coreutils';
-import { requestAPI } from './handler';
+import { requestAPI } from '../handler';
+import { ResettableNotebook } from './notebook';
 
 const globalModelDBMutex = createMutex();
 
@@ -12,8 +18,6 @@ const globalModelDBMutex = createMutex();
  */
 const DIRTY_CLASS = 'jp-mod-dirty';
 
-
-
 (CodeCellModel.prototype as any)._onSharedModelChanged = function (
   slot: ISharedCodeCell,
   change: CellChange
@@ -142,26 +146,25 @@ class RtcOutputAreaModel extends OutputAreaModel implements IOutputAreaModel {
   }
 }
 
-/** 
- * NOTE: We should upstream this fix. This is a bug in JupyterLab. 
- * 
+/**
+ * NOTE: We should upstream this fix. This is a bug in JupyterLab.
+ *
  * The execution count comes back from the kernel immediately
  * when the execute request is made by the client, even thought
  * cell might still be running. JupyterLab holds this value in
- * memory with a Promise to set it later, once the execution 
- * state goes back to Idle. 
- * 
+ * memory with a Promise to set it later, once the execution
+ * state goes back to Idle.
+ *
  * In CRDT world, we don't need to do this gymnastics, holding
- * the state in a Promise. Instead, we can just watch the 
+ * the state in a Promise. Instead, we can just watch the
  * executionState and executionCount in the CRDT being maintained
  * by the server-side model.
- * 
+ *
  * This is a big win! It means user can close and re-open a
- * notebook while a list of executed cells are queued. 
+ * notebook while a list of executed cells are queued.
  */
 (CodeCell.prototype as any).onStateChanged = function (
-
-  model: ICellModel, 
+  model: ICellModel,
   args: IChangedArgs<any>
 ): void {
   switch (args.name) {
@@ -188,7 +191,7 @@ class RtcOutputAreaModel extends OutputAreaModel implements IOutputAreaModel {
     default:
       break;
   }
-}
+};
 
 CodeCellModel.ContentFactory.prototype.createOutputArea = function (
   options: IOutputAreaModel.IOptions
@@ -203,4 +206,8 @@ export class RtcNotebookContentFactory
   createCodeCell(options: CodeCell.IOptions): CodeCell {
     return new CodeCell(options).initializeState();
   }
+
+  createNotebook(options: Notebook.IOptions): Notebook {
+    return new ResettableNotebook(options);
+  }
 }

src/notebook-factory/notebook.ts

@@ -0,0 +1,65 @@
+import { INotebookModel, Notebook, NotebookModel } from '@jupyterlab/notebook';
+import { YNotebook } from '../docprovider/custom_ydocs';
+
+/**
+ * A custom implementation of `Notebook` that resets the notebook to an empty
+ * state when `YNotebook.resetSignal` is emitted to.
+ *
+ * This requires the custom `YNotebook` class defined by this labextension.
+ */
+export class ResettableNotebook extends Notebook {
+  constructor(options: Notebook.IOptions) {
+    super(options);
+    this._resetSignalSlot = () => this._onReset();
+  }
+
+  get model(): INotebookModel | null {
+    return super.model;
+  }
+
+  set model(newValue: INotebookModel | null) {
+    // if current model exists, remove the `resetSignal` observer
+    if (this.model) {
+      const ynotebook = this.model.sharedModel as YNotebook;
+      ynotebook.resetSignal.disconnect(this._resetSignalSlot);
+    }
+
+    // call parent property setter
+    super.model = newValue;
+
+    // return early if `newValue === null`
+    if (!newValue) {
+      return;
+    }
+
+    // otherwise, listen to `YNotebook.resetSignal`.
+    const ynotebook = newValue.sharedModel as YNotebook;
+    ynotebook.resetSignal.connect(this._resetSignalSlot);
+  }
+
+  /**
+   * Function called when the YDoc has been reset. This recreates the notebook
+   * model using this model's options.
+   *
+   * TODO (?): we may want to use NotebookModelFactory, but that factory only
+   * seems to set some configuration options. The NotebookModel constructor
+   * does not require any arguments so this is OK for now.
+   */
+  _onReset() {
+    if (!this.model) {
+      console.warn(
+        'The notebook was reset without a model. This should never happen.'
+      );
+      return;
+    }
+
+    this.model = new NotebookModel({
+      collaborationEnabled: this.model.collaborative,
+      sharedModel: this.model.sharedModel
+      // other options in `NotebookModel.IOptions` are either unused or
+      // forwarded to `YNotebook`, which is preserved here
+    });
+  }
+
+  _resetSignalSlot: () => void;
+}

src/notebook-factory/plugin.ts

@@ -0,0 +1,26 @@
+import type {
+  JupyterFrontEnd,
+  JupyterFrontEndPlugin
+} from '@jupyterlab/application';
+import { NotebookPanel } from '@jupyterlab/notebook';
+import { IEditorServices } from '@jupyterlab/codeeditor';
+
+import { RtcNotebookContentFactory } from './notebook-factory';
+
+type NotebookFactoryPlugin =
+  JupyterFrontEndPlugin<NotebookPanel.IContentFactory>;
+
+/**
+ * Custom `Notebook` factory plugin.
+ */
+export const notebookFactoryPlugin: NotebookFactoryPlugin = {
+  id: '@jupyter/server-documents/notebook-extension:factory',
+  description: 'Provides the notebook cell factory.',
+  provides: NotebookPanel.IContentFactory,
+  requires: [IEditorServices],
+  autoStart: true,
+  activate: (app: JupyterFrontEnd, editorServices: IEditorServices) => {
+    const editorFactory = editorServices.factoryService.newInlineEditor;
+    return new RtcNotebookContentFactory({ editorFactory });
+  }
+};