Expose comm and custom messages on the widget model (#30)

Author: arjxn-py

src/model.ts

@@ -1,10 +1,11 @@
 import { MapChange } from '@jupyter/ydoc';
+import { Kernel } from '@jupyterlab/services';
 import { JSONExt, JSONObject, PromiseDelegate } from '@lumino/coreutils';
 import { ISignal, Signal } from '@lumino/signaling';
 import * as Y from 'yjs';
 
 import type { Awareness } from 'y-protocols/awareness';
-import { IJupyterYDoc, IJupyterYModel } from './types';
+import { IJupyterYDoc, IJupyterYModel, IYCommProvider } from './types';
 
 export class JupyterYModel implements IJupyterYModel {
   constructor(commMetadata: { [key: string]: any }) {
@@ -57,6 +58,41 @@ export class JupyterYModel implements IJupyterYModel {
     return undefined;
   }
 
+  /**
+   * The raw comm channel backing this model, available once the widget
+   * manager has wired up the comm provider.
+   */
+  get comm(): Kernel.IComm | undefined {
+    return this._commProvider?.comm;
+  }
+
+  /**
+   * Fires when a custom (non Y-protocol) message is received over the comm.
+   */
+  get messageReceived(): ISignal<IJupyterYModel, JSONObject> {
+    return this._messageReceived;
+  }
+
+  /**
+   * Send a custom message to the kernel over the comm.
+   */
+  sendMessage(data: JSONObject): void {
+    this._commProvider?.sendMessage(data);
+  }
+
+  /**
+   * Wire the model to its comm provider. Called by the widget manager once
+   * the comm is open; relays custom comm messages through `messageReceived`.
+   */
+  setCommProvider(provider: IYCommProvider): void {
+    this._commProvider = provider;
+    provider.messageReceived.connect(this._onCommMessage, this);
+  }
+
+  private _onCommMessage = (_: IYCommProvider, data: JSONObject): void => {
+    this._messageReceived.emit(data);
+  };
+
   protected async initialize(commMetadata: { [key: string]: any }) {
     this.ydoc = this.ydocFactory(commMetadata);
     this.sharedModel = new JupyterYDoc(commMetadata, this._ydoc);
@@ -71,6 +107,8 @@ export class JupyterYModel implements IJupyterYModel {
       return;
     }
     this._isDisposed = true;
+    this._commProvider?.messageReceived.disconnect(this._onCommMessage, this);
+    this._commProvider?.dispose();
     this._sharedModel.dispose();
     this._disposed.emit();
     Signal.clearData(this);
@@ -89,6 +127,9 @@ export class JupyterYModel implements IJupyterYModel {
   private _yModelName: string;
   private _sharedModel: IJupyterYDoc;
 
+  private _commProvider?: IYCommProvider;
+  private _messageReceived = new Signal<IJupyterYModel, JSONObject>(this);
+
   private _isDisposed = false;
 
   private _ready: PromiseDelegate<void> = new PromiseDelegate<void>();

src/notebookrenderer/widgetManager.ts

@@ -71,11 +71,12 @@ export class WidgetModelRegistry implements IJupyterYWidgetModelRegistry {
 
     await yModel.ready;
 
-    new YCommProvider({
+    const commProvider = new YCommProvider({
       comm,
       ydoc: yModel.sharedModel.ydoc,
       awareness: yModel.awareness
     });
+    yModel.setCommProvider(commProvider);
     this._yModels.set(comm.commId, yModel);
   };
 

src/notebookrenderer/yCommProvider.ts

@@ -8,7 +8,10 @@ import {
 } from 'y-protocols/awareness';
 import * as syncProtocol from 'y-protocols/sync';
 import * as Y from 'yjs';
-import { IDisposable } from '@lumino/disposable';
+import { JSONObject } from '@lumino/coreutils';
+import { ISignal, Signal } from '@lumino/signaling';
+
+import { IYCommProvider } from '../types';
 
 export enum YMessageType {
   SYNC = 0,
@@ -26,7 +29,7 @@ export interface IYCommProviderOptions {
   awareness?: Awareness;
 }
 
-export class YCommProvider implements IDisposable {
+export class YCommProvider implements IYCommProvider {
   constructor(options: IYCommProviderOptions) {
     this._comm = options.comm;
     this._ydoc = options.ydoc;
@@ -49,10 +52,21 @@ export class YCommProvider implements IDisposable {
     return this._ydoc;
   }
 
+  get comm(): Kernel.IComm {
+    return this._comm;
+  }
+
   get awareness(): Awareness {
     return this._awareness;
   }
 
+  /**
+   * Fires when a custom (non Y-protocol) message is received over the comm.
+   */
+  get messageReceived(): ISignal<IYCommProvider, JSONObject> {
+    return this._messageReceived;
+  }
+
   get synced(): boolean {
     return this._synced;
   }
@@ -78,10 +92,19 @@ export class YCommProvider implements IDisposable {
     }
     this._comm.close();
     this._isDisposed = true;
+    Signal.clearData(this);
+  }
+
+  /**
+   * Send a custom message to the kernel over the comm.
+   */
+  sendMessage(data: JSONObject): void {
+    this._comm.send(data);
   }
 
   private _onMsg = (msg: KernelMessage.ICommMsgMsg<'iopub' | 'shell'>) => {
-    if (msg.buffers) {
+    if (msg.buffers && msg.buffers.length > 0) {
+      // Binary payload: a Y-protocol frame (document sync or awareness).
       const buffer = msg.buffers[0] as ArrayBuffer;
       const buffer_uint8 = new Uint8Array(
         ArrayBuffer.isView(buffer) ? buffer.buffer : buffer
@@ -90,6 +113,9 @@ export class YCommProvider implements IDisposable {
       if (encoding.length(encoder) > 1) {
         this._sendOverComm(encoding.toUint8Array(encoder));
       }
+    } else {
+      // No buffers: a custom application message
+      this._messageReceived.emit(msg.content.data as JSONObject);
     }
   };
 
@@ -140,6 +166,7 @@ export class YCommProvider implements IDisposable {
   private _ownsAwareness: boolean;
   private _synced: boolean;
   private _isDisposed = false;
+  private _messageReceived = new Signal<IYCommProvider, JSONObject>(this);
 }
 
 namespace Private {

src/types.ts

@@ -1,4 +1,5 @@
 import { MapChange, StateChange } from '@jupyter/ydoc';
+import { Kernel } from '@jupyterlab/services';
 import * as Y from 'yjs';
 import { ISignal } from '@lumino/signaling';
 import { JSONObject } from '@lumino/coreutils';
@@ -23,6 +24,31 @@ export interface IJupyterYDoc extends IDisposable {
   disposed: ISignal<any, void>;
 }
 
+/**
+ * A provider bridging a Y.Doc to a Jupyter Comm channel.
+ *
+ * Besides the Y-protocol (sync/awareness) traffic, it exposes the raw comm
+ * and relays custom, non-document messages so consumers can implement
+ * stateless "one shot" actions (e.g. a "zoom-to" request).
+ */
+export interface IYCommProvider extends IDisposable {
+  /**
+   * The underlying comm channel, exposed so consumers can do anything they
+   * need with it beyond the built-in message handling.
+   */
+  comm: Kernel.IComm;
+
+  /**
+   * Fires when a custom (non Y-protocol) message is received over the comm.
+   */
+  messageReceived: ISignal<IYCommProvider, JSONObject>;
+
+  /**
+   * Send a custom message to the kernel over the comm.
+   */
+  sendMessage(data: JSONObject): void;
+}
+
 export interface IJupyterYModel extends IDisposable {
   yModelName: string;
   isDisposed: boolean;
@@ -39,4 +65,27 @@ export interface IJupyterYModel extends IDisposable {
    * instead of creating a second instance on the same Y.Doc
    */
   awareness?: Awareness;
+
+  /**
+   * The raw comm channel backing this model, exposed so consumers can send
+   * and receive custom (non-document) messages. Available once the comm has
+   * been opened and wired up by the widget manager.
+   */
+  comm?: Kernel.IComm;
+
+  /**
+   * Fires when a custom (non Y-protocol) message is received over the comm.
+   */
+  messageReceived: ISignal<IJupyterYModel, JSONObject>;
+
+  /**
+   * Send a custom message to the kernel over the comm.
+   */
+  sendMessage(data: JSONObject): void;
+
+  /**
+   * Wire the model to its comm provider. Called by the widget manager once
+   * the comm is open; enables `comm`, `sendMessage` and `messageReceived`.
+   */
+  setCommProvider(provider: IYCommProvider): void;
 }