Chat subprotocol (#167)

* Creates a subprotocol to broadcast messages

* Creates a WebSocketAwarenessProvider

* Pre commit

Author: hbcarlos

jupyter_collaboration/handlers.py

@@ -16,11 +16,16 @@
 from tornado.websocket import WebSocketHandler
 from ypy_websocket.websocket_server import YRoom
 from ypy_websocket.ystore import BaseYStore
-from ypy_websocket.yutils import YMessageType
+from ypy_websocket.yutils import YMessageType, write_var_uint
 
 from .loaders import FileLoaderMapping
 from .rooms import DocumentRoom, TransientRoom
-from .utils import JUPYTER_COLLABORATION_EVENTS_URI, LogLevel, decode_file_path
+from .utils import (
+    JUPYTER_COLLABORATION_EVENTS_URI,
+    LogLevel,
+    MessageType,
+    decode_file_path,
+)
 from .websocketserver import JupyterWebsocketServer
 
 YFILE = YDOCS["file"]
@@ -197,6 +202,8 @@ def on_message(self, message):
         On message receive.
         """
         message_type = message[0]
+        print("message type:", message_type)
+
         if message_type == YMessageType.AWARENESS:
             # awareness
             skip = False
@@ -222,6 +229,18 @@ def on_message(self, message):
                 )
                 return skip
 
+        if message_type == MessageType.CHAT:
+            msg = message[2:].decode("utf-8")
+            user = self.get_current_user()
+            data = json.dumps({"username": user.username, "msg": msg}).encode("utf8")
+            for client in self.room.clients:
+                if client != self:
+                    task = asyncio.create_task(
+                        client.send(bytes([MessageType.CHAT]) + write_var_uint(len(data)) + data)
+                    )
+                    self._websocket_server.background_tasks.add(task)
+                    task.add_done_callback(self._websocket_server.background_tasks.discard)
+
         self._message_queue.put_nowait(message)
         self._websocket_server.ypatch_nb += 1
 

jupyter_collaboration/utils.py

@@ -2,13 +2,19 @@
 # Distributed under the terms of the Modified BSD License.
 
 import pathlib
-from enum import Enum
+from enum import Enum, IntEnum
 from typing import Tuple
 
 JUPYTER_COLLABORATION_EVENTS_URI = "https://schema.jupyter.org/jupyter_collaboration/session/v1"
 EVENTS_SCHEMA_PATH = pathlib.Path(__file__).parent / "events" / "session.yaml"
 
 
+class MessageType(IntEnum):
+    SYNC = 0
+    AWARENESS = 1
+    CHAT = 125
+
+
 class LogLevel(Enum):
     INFO = "INFO"
     DEBUG = "DEBUG"

packages/collaboration-extension/src/collaboration.ts

@@ -16,14 +16,14 @@ import {
 } from '@jupyterlab/codemirror';
 import {
   CollaboratorsPanel,
-  IAwareness,
   IGlobalAwareness,
   IUserMenu,
   remoteUserCursors,
   RendererUserMenu,
   UserInfoPanel,
   UserMenu
 } from '@jupyter/collaboration';
+import { IAwareness, WebSocketAwarenessProvider } from '@jupyter/docprovider';
 import { SidePanel, usersIcon } from '@jupyterlab/ui-components';
 import { Menu, MenuBar } from '@lumino/widgets';
 import { URLExt } from '@jupyterlab/coreutils';
@@ -33,7 +33,6 @@ import { ITranslator, nullTranslator } from '@jupyterlab/translation';
 
 import * as Y from 'yjs';
 import { Awareness } from 'y-protocols/awareness';
-import { WebsocketProvider } from 'y-websocket';
 
 /**
  * Jupyter plugin providing the IUserMenu.
@@ -90,26 +89,20 @@ export const rtcGlobalAwarenessPlugin: JupyterFrontEndPlugin<IAwareness> = {
   provides: IGlobalAwareness,
   activate: (app: JupyterFrontEnd, state: StateDB): IAwareness => {
     const { user } = app.serviceManager;
-    const ydoc = new Y.Doc();
 
+    const ydoc = new Y.Doc();
     const awareness = new Awareness(ydoc);
 
     const server = ServerConnection.makeSettings();
     const url = URLExt.join(server.wsUrl, 'api/collaboration/room');
 
-    new WebsocketProvider(url, 'JupyterLab:globalAwareness', ydoc, {
-      awareness: awareness
+    new WebSocketAwarenessProvider({
+      url: url,
+      roomID: 'JupyterLab:globalAwareness',
+      awareness: awareness,
+      user: user
     });
 
-    const userChanged = () => {
-      awareness.setLocalStateField('user', user.identity);
-    };
-    if (user.isReady) {
-      userChanged();
-    }
-    user.ready.then(userChanged).catch(e => console.error(e));
-    user.userChanged.connect(userChanged);
-
     state.changed.connect(async () => {
       const data: any = await state.toJSON();
       const current = data['layout-restorer:data']?.main?.current || '';

packages/collaboration/package.json

@@ -41,6 +41,7 @@
   "dependencies": {
     "@codemirror/state": "^6.2.0",
     "@codemirror/view": "^6.7.0",
+    "@jupyter/docprovider": "1.0.0",
     "@jupyterlab/apputils": "^4.0.0",
     "@jupyterlab/coreutils": "^6.0.0",
     "@jupyterlab/services": "^7.0.0",

packages/collaboration/src/tokens.ts

@@ -3,7 +3,7 @@
 
 import type { Menu } from '@lumino/widgets';
 import { Token } from '@lumino/coreutils';
-import type { Awareness } from 'y-protocols/awareness';
+import { IAwareness } from '@jupyter/docprovider';
 import type { User } from '@jupyterlab/services';
 
 /**
@@ -23,11 +23,6 @@ export const IGlobalAwareness = new Token<IAwareness>(
   '@jupyter/collaboration:IGlobalAwareness'
 );
 
-/**
- * The awareness interface.
- */
-export type IAwareness = Awareness;
-
 /**
  * An interface describing the user menu.
  */

packages/docprovider/src/awareness.ts

@@ -0,0 +1,141 @@
+/* -----------------------------------------------------------------------------
+| Copyright (c) Jupyter Development Team.
+| Distributed under the terms of the Modified BSD License.
+|----------------------------------------------------------------------------*/
+
+import { User } from '@jupyterlab/services';
+
+import { IDisposable } from '@lumino/disposable';
+import { ISignal, Signal } from '@lumino/signaling';
+
+import * as decoding from 'lib0/decoding';
+import * as encoding from 'lib0/encoding';
+import { WebsocketProvider } from 'y-websocket';
+
+import { IAwareness, IAwarenessProvider } from './tokens';
+
+export enum MessageType {
+  CHAT = 125
+}
+
+export interface IChatMessage {
+  username: string;
+  msg: string;
+}
+
+/**
+ * A class to provide Yjs synchronization over WebSocket.
+ *
+ * We specify custom messages that the server can interpret. For reference please look in yjs_ws_server.
+ *
+ */
+export class WebSocketAwarenessProvider
+  extends WebsocketProvider
+  implements IAwarenessProvider, IDisposable
+{
+  /**
+   * Construct a new WebSocketAwarenessProvider
+   *
+   * @param options The instantiation options for a WebSocketAwarenessProvider
+   */
+  constructor(options: WebSocketAwarenessProvider.IOptions) {
+    super(options.url, options.roomID, options.awareness.doc, {
+      awareness: options.awareness
+    });
+
+    this._awareness = options.awareness;
+
+    const user = options.user;
+    user.ready
+      .then(() => this._onUserChanged(user))
+      .catch(e => console.error(e));
+    user.userChanged.connect(this._onUserChanged, this);
+
+    this._chatMessage = new Signal(this);
+
+    this.messageHandlers[MessageType.CHAT] = (
+      encoder,
+      decoder,
+      provider,
+      emitSynced,
+      messageType
+    ) => {
+      const content = decoding.readVarString(decoder);
+      const data = JSON.parse(content) as IChatMessage;
+      console.debug('Chat:', data);
+      this._chatMessage.emit(data);
+    };
+  }
+
+  get isDisposed(): boolean {
+    return this._isDisposed;
+  }
+
+  /**
+   * A signal to subscribe for incoming messages.
+   */
+  get chatMessage(): ISignal<this, IChatMessage> {
+    return this._chatMessage;
+  }
+
+  dispose(): void {
+    if (this._isDisposed) {
+      return;
+    }
+
+    this.destroy();
+    this._isDisposed = true;
+  }
+
+  /**
+   * Send a message to every collaborator.
+   *
+   * @param msg message
+   */
+  sendMessage(msg: string): void {
+    console.debug('Send message:', msg);
+    const encoder = encoding.createEncoder();
+    encoding.writeVarUint(encoder, MessageType.CHAT);
+    encoding.writeVarString(encoder, msg);
+    this.ws!.send(encoding.toUint8Array(encoder));
+  }
+
+  private _onUserChanged(user: User.IManager): void {
+    this._awareness.setLocalStateField('user', user.identity);
+  }
+
+  private _isDisposed = false;
+  private _awareness: IAwareness;
+
+  private _chatMessage: Signal<this, IChatMessage>;
+}
+
+/**
+ * A namespace for WebSocketAwarenessProvider statics.
+ */
+export namespace WebSocketAwarenessProvider {
+  /**
+   * The instantiation options for a WebSocketAwarenessProvider.
+   */
+  export interface IOptions {
+    /**
+     * The server URL
+     */
+    url: string;
+
+    /**
+     * The room ID
+     */
+    roomID: string;
+
+    /**
+     * The awareness object
+     */
+    awareness: IAwareness;
+
+    /**
+     * The user data
+     */
+    user: User.IManager;
+  }
+}

packages/docprovider/src/index.ts

@@ -7,6 +7,7 @@
  * @module docprovider
  */
 
+export * from './awareness';
 export * from './ydrive';
 export * from './yprovider';
 export * from './tokens';

packages/docprovider/src/tokens.ts

@@ -3,7 +3,13 @@
 
 import { DocumentChange, YDocument } from '@jupyter/ydoc';
 import { Contents } from '@jupyterlab/services';
+
 import { Token } from '@lumino/coreutils';
+import { ISignal } from '@lumino/signaling';
+
+import type { Awareness } from 'y-protocols/awareness';
+
+import { IChatMessage } from './awareness';
 
 /**
  * The collaborative drive.
@@ -45,3 +51,27 @@ export interface ISharedModelFactory extends Contents.ISharedFactory {
     factory: SharedDocumentFactory
   ): void;
 }
+
+/**
+ * The awareness interface.
+ *
+ * TODO: Move to @jupyter/YDoc
+ */
+export type IAwareness = Awareness;
+
+/**
+ * A provider interface for global awareness features.
+ */
+export interface IAwarenessProvider {
+  /**
+   * A signal to subscribe for incoming messages.
+   */
+  get chatMessage(): ISignal<this, IChatMessage>;
+
+  /**
+   * Send a message to every collaborator.
+   *
+   * @param msg message
+   */
+  sendMessage(msg: string): void;
+}

yarn.lock

@@ -2054,6 +2054,7 @@ __metadata:
   dependencies:
     "@codemirror/state": ^6.2.0
     "@codemirror/view": ^6.7.0
+    "@jupyter/docprovider": 1.0.0
     "@jupyterlab/apputils": ^4.0.0
     "@jupyterlab/coreutils": ^6.0.0
     "@jupyterlab/services": ^7.0.0
@@ -2070,7 +2071,7 @@ __metadata:
   languageName: unknown
   linkType: soft
 
-"@jupyter/docprovider@^1.0.0, @jupyter/docprovider@workspace:packages/docprovider":
+"@jupyter/docprovider@1.0.0, @jupyter/docprovider@^1.0.0, @jupyter/docprovider@workspace:packages/docprovider":
   version: 0.0.0-use.local
   resolution: "@jupyter/docprovider@workspace:packages/docprovider"
   dependencies: