Handle YChat document resets when jupyterlab_chat is installed

package.json

@@ -59,6 +59,7 @@
         "watch:labextension": "jupyter labextension watch ."
     },
     "dependencies": {
+        "@jupyter/chat": ">=0.6.0 <1",
         "@jupyter/collaborative-drive": "^4",
         "@jupyter/ydoc": "^2.1.3 || ^3.0.0",
         "@jupyterlab/application": "^4.4.0",
@@ -76,6 +77,7 @@
         "@lumino/coreutils": "^2.2.1",
         "@lumino/disposable": "^2.1.4",
         "@lumino/signaling": "^2.1.4",
+        "jupyterlab-chat": ">=0.6.0 <1",
         "y-protocols": "^1.0.5",
         "y-websocket": "^1.3.15",
         "yjs": "^13.5.40"
@@ -134,7 +136,15 @@
         "disabledExtensions": [
             "@jupyterlab/codemirror-extension:binding",
             "@jupyter/docprovider-extension"
-        ]
+        ],
+        "sharedPackages": {
+            "@jupyter/chat": {
+                "singleton": true
+            },
+            "jupyterlab-chat": {
+                "singleton": true
+            }
+        }
     },
     "eslintIgnore": [
         "node_modules",

src/docprovider/custom_ydocs.ts

@@ -3,6 +3,7 @@ import {
   YNotebook as DefaultYNotebook,
   ISharedNotebook
 } from '@jupyter/ydoc';
+import { YChat as DefaultYChat } from 'jupyterlab-chat';
 import * as Y from 'yjs';
 import { Awareness } from 'y-protocols/awareness';
 import { ISignal, Signal } from '@lumino/signaling';
@@ -114,3 +115,49 @@ export class YNotebook extends DefaultYNotebook {
 
   _resetSignal: Signal<this, null>;
 }
+
+export class YChat extends DefaultYChat {
+  constructor() {
+    super();
+    this._resetSignal = new Signal(this);
+  }
+
+  /**
+   * See `YFile.reset()`.
+   */
+  reset() {
+    // Remove default observers
+    (this as any)._users.unobserve((this as any)._usersObserver);
+    (this as any)._messages.unobserve((this as any)._messagesObserver);
+    (this as any)._attachments.unobserve((this as any)._attachmentsObserver);
+    (this as any)._metadata.unobserve((this as any)._metadataObserver);
+
+    // Reset `this._ydoc` to an empty state
+    (this as any)._ydoc = new Y.Doc();
+
+    // Reset all properties derived from `this._ydoc`
+    (this as any)._users = this.ydoc.getMap('users');
+    (this as any)._messages = this.ydoc.getArray('messages');
+    (this as any)._attachments = this.ydoc.getMap('attachments');
+    (this as any)._metadata = this.ydoc.getMap('metadata');
+    (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 as any)._users.observe((this as any)._usersObserver);
+    (this as any)._messages.observe((this as any)._messagesObserver);
+    (this as any)._attachments.observe((this as any)._attachmentsObserver);
+    (this as any)._metadata.observe((this as any)._metadataObserver);
+  }
+
+  /**
+   * See `YFile.resetSignal`.
+   */
+  get resetSignal(): ISignal<this, null> {
+    return this._resetSignal;
+  }
+
+  _resetSignal: Signal<this, null>;
+}

src/docprovider/filebrowser.ts

@@ -8,7 +8,7 @@ import {
   JupyterFrontEndPlugin
 } from '@jupyterlab/application';
 import { Dialog, showDialog } from '@jupyterlab/apputils';
-import { IDocumentWidget } from '@jupyterlab/docregistry';
+import { DocumentWidget, IDocumentWidget } from '@jupyterlab/docregistry';
 import { ContentsManager } from '@jupyterlab/services';
 
 import {
@@ -34,6 +34,10 @@ import {
 import { RtcContentProvider } from './ydrive';
 import { Awareness } from 'y-protocols/awareness';
 
+import { YChat } from './custom_ydocs';
+import { IChatFactory } from 'jupyterlab-chat';
+import { AbstractChatModel } from '@jupyter/chat';
+
 const TWO_SESSIONS_WARNING =
   'The file %1 has been opened with two different views. ' +
   'This is not supported. Please close this view; otherwise, ' +
@@ -154,6 +158,76 @@ export const ynotebook: JupyterFrontEndPlugin<void> = {
   }
 };
 
+/**
+ * This plugin provides the YChat shared model and handles document resets by
+ * listening to the `YChat.resetSignal` property automatically.
+ *
+ * Whenever a YChat is reset, this plugin will iterate through all of the app's
+ * document widgets and find the one containing the `YChat` shared model which
+ * was reset. It then clears the content.
+ */
+export const ychat: JupyterFrontEndPlugin<void> = {
+  id: '@jupyter/server-documents:ychat',
+  description:
+    'Plugin to register a custom YChat factory and handle document resets.',
+  autoStart: true,
+  requires: [ICollaborativeContentProvider],
+  optional: [IChatFactory],
+  activate: (
+    app: JupyterFrontEnd,
+    contentProvider: ICollaborativeContentProvider,
+    chatFactory?: IChatFactory
+  ): void => {
+    if (!chatFactory) {
+      console.warn(
+        'No existing shared model factory found for chat. Not providing custom chat shared model.'
+      );
+      return;
+    }
+
+    const onYChatReset = (ychat: YChat) => {
+      for (const widget of app.shell.widgets()) {
+        if (!(widget instanceof DocumentWidget)) {
+          continue;
+        }
+        const model = widget.content.model;
+        const sharedModel = model && model._sharedModel;
+        if (
+          !(model instanceof AbstractChatModel && sharedModel instanceof YChat)
+        ) {
+          continue;
+        }
+        if (sharedModel !== ychat) {
+          continue;
+        }
+
+        // If this point is reached, we have identified the correct parent
+        // `model: AbstractChatModel` that maintains the message state for the
+        // `YChat` which was reset. We clear its content directly & emit a
+        // `contentChanged` signal to update the UI.
+        (model as any)._messages = [];
+        (model as any)._messagesUpdated.emit();
+        break;
+      }
+    };
+
+    // Override the existing `YChat` factory to provide a custom `YChat` with a
+    // `resetSignal`, which is automatically subscribed to & refreshes the UI
+    // state upon document reset.
+    const yChatFactory = () => {
+      const ychat = new YChat();
+      ychat.resetSignal.connect(() => {
+        onYChatReset(ychat);
+      });
+      return ychat;
+    };
+    contentProvider.sharedModelFactory.registerDocumentFactory(
+      'chat',
+      yChatFactory as any
+    );
+  }
+};
+
 /**
  * The default collaborative drive provider.
  */

src/docprovider/ydrive.ts

@@ -272,7 +272,10 @@ class SharedModelFactory implements ISharedModelFactory {
     factory: SharedDocumentFactory
   ) {
     if (this.documentFactories.has(type)) {
-      throw new Error(`The content type ${type} already exists`);
+      // allow YChat shared model factory to be overridden
+      if (type !== 'chat') {
+        throw new Error(`The content type ${type} already exists.`);
+      }
     }
     this.documentFactories.set(type, factory);
   }

src/docprovider/yprovider.ts

@@ -21,6 +21,7 @@ import { JupyterFrontEnd } from '@jupyterlab/application';
 import { DocumentWidget } from '@jupyterlab/docregistry';
 import { FileEditor } from '@jupyterlab/fileeditor';
 import { Notebook } from '@jupyterlab/notebook';
+import { ChatWidget } from '@jupyter/chat';
 
 /**
  * A class to provide Yjs synchronization over WebSocket.
@@ -87,9 +88,15 @@ export class WebSocketProvider implements IDocumentProvider {
         continue;
       }
 
-      // Skip widgets that don't contain a YFile / YNotebook
+      // Skip widgets that don't contain a YFile / YNotebook / YChat
       const widget = docWidget.content;
-      if (!(widget instanceof FileEditor || widget instanceof Notebook)) {
+      if (
+        !(
+          widget instanceof FileEditor ||
+          widget instanceof Notebook ||
+          widget instanceof ChatWidget
+        )
+      ) {
         continue;
       }
 

src/index.ts

@@ -28,7 +28,13 @@ import { AwarenessExecutionIndicator } from './executionindicator';
 
 import { requestAPI } from './handler';
 
-import { rtcContentProvider, yfile, ynotebook, logger } from './docprovider';
+import {
+  rtcContentProvider,
+  yfile,
+  ynotebook,
+  ychat,
+  logger
+} from './docprovider';
 
 import { IStateDB, StateDB } from '@jupyterlab/statedb';
 import { IGlobalAwareness } from '@jupyter/collaborative-drive';
@@ -317,7 +323,8 @@ const plugins: JupyterFrontEndPlugin<unknown>[] = [
   notebookFactoryPlugin,
   codemirrorYjsPlugin,
   backupCellExecutorPlugin,
-  disableSavePlugin
+  disableSavePlugin,
+  ychat
 ];
 
 export default plugins;