Handle YDoc reset & out-of-band changes in notebooks

jupyter_server_documents/rooms/yroom.py

@@ -405,18 +405,23 @@ def _on_ydoc_update(self, event: TransactionEvent) -> None:
         self._broadcast_message(message, message_type="SyncUpdate")
 
 
-    def _on_jupyter_ydoc_update(self, updated_key: str, *_) -> None:
+    def _on_jupyter_ydoc_update(self, updated_key: str, event: Any) -> None:
         """
         This method is an observer on `self._jupyter_ydoc` which saves the file
-        whenever the YDoc changes, unless `updated_key == "state"`.
+        whenever the YDoc changes.
 
-        The `state` key is used by `jupyter_ydoc` to store temporary data like
-        whether a file is 'dirty' (has unsaved changes). This data is not
-        persisted, so changes to the `state` key should be ignored. Otherwise,
-        an infinite loop of saves will result, as saving sets `dirty = False`.
+        - This observer ignores updates to the 'state' dictionary if they have
+        no effect. See `should_ignore_state_update()` documentation for info.
 
-        This observer is separate because `pycrdt.Doc.observe()` does not pass
-        `updated_key` to `self._on_ydoc_update()`.
+        - This observer is separate from the `pycrdt` observer because we must
+        check if the update should be ignored. This requires the `updated_key`
+        and `event` arguments exclusive to `jupyter_ydoc` observers, not
+        available to `pycrdt` observers.
+
+        - The type of the `event` argument depends on the shared type that
+        `updated_key` references. If it references a `pycrdt.Map`, then event
+        will always be of type `pycrdt.MapEvent`. Same applies for other shared
+        types, like `pycrdt.Array` or `pycrdt.Text`.
         """
         # Do nothing if there is no file API for this room (e.g. global awareness)
         if self.file_api is None:
@@ -429,10 +434,16 @@ def _on_jupyter_ydoc_update(self, updated_key: str, *_) -> None:
         if content_loading:
             return
 
-        # Save only when the content of the YDoc is updated.
-        # See this method's docstring for more context.
-        if updated_key != "state":
-            self.file_api.schedule_save()
+        # Do nothing if the event updates the 'state' dictionary with no effect
+        if updated_key == "state":
+            # The 'state' key always refers to a `pycrdt.Map` shared type, so
+            # event always has type `pycrdt.MapEvent`.
+            map_event = cast(pycrdt.MapEvent, event)
+            if should_ignore_state_update(map_event):
+                return
+
+        # Otherwise, save the file
+        self.file_api.schedule_save()
 
 
     def handle_awareness_update(self, client_id: str, message: bytes) -> None:
@@ -614,3 +625,40 @@ async def stop(self) -> None:
         # JupyterYDoc in the process.
         if self.file_api:
             await self.file_api.stop_then_save()
+
+def should_ignore_state_update(event: pycrdt.MapEvent) -> bool:
+    """
+    Returns whether an update to the `state` dictionary should be ignored by
+    `_on_jupyter_ydoc_update()`. Every Jupyter YDoc includes this dictionary in
+    its YDoc.
+
+    This function returns `False` if the update has no effect, i.e. the event
+    consists of updating each key to the same value it had originally.
+
+    Motivation: `pycrdt` emits update events on the 'state' key even when they have no
+    effect. Without ignoring those updates, saving the file triggers an
+    infinite loop of saves, as setting `jupyter_ydoc.dirty = False` always
+    emits an update, even if that attribute was already `False`. See PR #50 for
+    more info.
+    """
+    # Iterate through the keys added/updated/deleted by this event. Return
+    # `False` immediately if:
+    # - a key was updated to a value different from the previous value
+    # - a key was added with a value different from the previous value
+    for key in event.keys.keys():
+        update_info = event.keys[key]
+        action = update_info.get('action', None)
+        if action == 'update':
+            old_value = update_info.get('oldValue', None)
+            new_value = update_info.get('newValue', None)
+            if old_value != new_value:
+                return False
+        elif action == "add":
+            old_value = event.target.get(key, None)
+            new_value = update_info.get('newValue', None)
+            if old_value != new_value:
+                return False
+
+    # Otherwise, return `True`.
+    return True
+

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
@@ -191,9 +190,6 @@ export class WebSocketProvider implements IDocumentProvider {
     if (isSynced) {
       if (this._yWebsocketProvider) {
         this._yWebsocketProvider.off('sync', this._onSync);
-
-        const state = this._sharedModel.ydoc.getMap('state');
-        state.set('document_id', this._yWebsocketProvider.roomname);
       }
       this._ready.resolve();
     }

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 → 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);
+  }
 }