From 9d5cfe107c2da9050e8262296dc0bc40847c3496 Mon Sep 17 00:00:00 2001
From: Ian Thomas <ianthomas23@gmail.com>
Date: Mon, 15 Jul 2024 12:51:36 +0100
Subject: [PATCH] Replace postMessage from webworker with comlink callback

---
 src/terminal.ts | 24 +++++++++---------------
 src/tokens.ts   | 10 +++++++++-
 src/worker.ts   | 17 +++++++++++------
 3 files changed, 29 insertions(+), 22 deletions(-)

diff --git a/src/terminal.ts b/src/terminal.ts
index 57873c9..e6c592e 100644
--- a/src/terminal.ts
+++ b/src/terminal.ts
@@ -3,13 +3,13 @@
 
 import { JSONPrimitive } from '@lumino/coreutils';
 
+import { proxy, wrap } from 'comlink';
+
 import {
   Server as WebSocketServer,
   Client as WebSocketClient
 } from 'mock-socket';
 
-import { wrap } from 'comlink';
-
 import { ITerminal, IRemoteWorkerTerminal } from './tokens';
 
 export class Terminal implements ITerminal {
@@ -28,17 +28,13 @@ export class Terminal implements ITerminal {
     this._remote = wrap(this._worker);
     const { baseUrl } = this.options;
     await this._remote.initialize({ baseUrl });
+    this._remote.registerCallbacks(proxy(this._outputCallback.bind(this)));
   }
 
-  /**
-   * Process a message coming from the JavaScript web worker.
-   *
-   * @param msg The worker message to process.
-   */
-  private _processWorkerMessage(msg: any, socket: WebSocketClient): void {
-    if (msg.type === 'output') {
-      const ret = JSON.stringify(['stdout', msg.text]);
-      socket.send(ret);
+  private async _outputCallback(text: string): Promise<void> {
+    if (this._socket) {
+      const ret = JSON.stringify(['stdout', text]);
+      this._socket.send(ret);
     }
   }
 
@@ -56,10 +52,7 @@ export class Terminal implements ITerminal {
 
     server.on('connection', async (socket: WebSocketClient) => {
       console.log('==> server connection', this, socket);
-
-      this._worker!.onmessage = e => {
-        this._processWorkerMessage(e.data, socket);
-      };
+      this._socket = socket;
 
       socket.on('message', async (message: any) => {
         const data = JSON.parse(message) as JSONPrimitive[];
@@ -95,4 +88,5 @@ export class Terminal implements ITerminal {
 
   private _worker?: Worker;
   private _remote?: IRemoteWorkerTerminal;
+  private _socket?: WebSocketClient;
 }
diff --git a/src/tokens.ts b/src/tokens.ts
index e03fbe4..5070672 100644
--- a/src/tokens.ts
+++ b/src/tokens.ts
@@ -3,9 +3,11 @@
 
 import { TerminalAPI } from '@jupyterlab/services';
 
+import { IOutputCallback } from '@jupyterlite/cockle';
+
 import { Token } from '@lumino/coreutils';
 
-import { Remote } from 'comlink';
+import { ProxyMarked, Remote } from 'comlink';
 
 /**
  * The token for the Terminals service.
@@ -71,11 +73,17 @@ export namespace IWorkerTerminal {
   }
 }
 
+export namespace IRemote {
+  export type OutputCallback = IOutputCallback & ProxyMarked;
+}
+
 export interface IRemote extends IWorkerTerminal {
   /**
    * Handle any lazy initialization activities.
    */
   initialize(options: IWorkerTerminal.IOptions): Promise<void>;
+
+  registerCallbacks(outputCallback: IRemote.OutputCallback): void;
 }
 
 /**
diff --git a/src/worker.ts b/src/worker.ts
index 1883a2e..4f90692 100644
--- a/src/worker.ts
+++ b/src/worker.ts
@@ -3,7 +3,7 @@ import { DriveFS } from '@jupyterlite/contents';
 
 import { expose } from 'comlink';
 
-import { IWorkerTerminal } from './tokens';
+import { IRemote, IWorkerTerminal } from './tokens';
 
 class WorkerTerminal implements IWorkerTerminal {
   async initialize(options: IWorkerTerminal.IOptions): Promise<void> {
@@ -15,12 +15,16 @@ class WorkerTerminal implements IWorkerTerminal {
     await this._shell!.input(text);
   }
 
+  registerCallbacks(outputCallback: IRemote.OutputCallback) {
+    this._outputCallback = outputCallback;
+  }
+
   async setSize(rows: number, columns: number): Promise<void> {
     await this._shell!.setSize(rows, columns);
   }
 
   async start(): Promise<void> {
-    this._shell = new Shell(this.output, this._mountpoint);
+    this._shell = new Shell(this.output.bind(this), this._mountpoint);
     const { FS, PATH, ERRNO_CODES } = await this._shell.initFilesystem();
 
     if (this._wantDriveFS) {
@@ -45,10 +49,9 @@ class WorkerTerminal implements IWorkerTerminal {
   }
 
   private async output(text: string): Promise<void> {
-    postMessage({
-      type: 'output',
-      text: text
-    });
+    if (this._outputCallback) {
+      await this._outputCallback(text);
+    }
   }
 
   private _options: IWorkerTerminal.IOptions | null = null;
@@ -56,6 +59,8 @@ class WorkerTerminal implements IWorkerTerminal {
   private _mountpoint: string = '/drive';
   private _wantDriveFS: boolean = true;
   private _driveFS?: DriveFS;
+
+  private _outputCallback?: IRemote.OutputCallback;
 }
 
 const obj = new WorkerTerminal();