From 62f723e50cf2ea29753a06b2f7799dec9dbe5539 Mon Sep 17 00:00:00 2001
From: Manav Rathi <manav@mrmr.io>
Date: Wed, 31 Jul 2024 11:33:18 +0530
Subject: [PATCH] Adapt

---
 desktop/src/main/services/ml-worker.ts     |  7 +--
 desktop/src/main/utils/comlink-endpoint.ts | 50 ++++++----------------
 2 files changed, 15 insertions(+), 42 deletions(-)

diff --git a/desktop/src/main/services/ml-worker.ts b/desktop/src/main/services/ml-worker.ts
index af5d2add62..e028dd01d5 100644
--- a/desktop/src/main/services/ml-worker.ts
+++ b/desktop/src/main/services/ml-worker.ts
@@ -15,11 +15,10 @@ import { existsSync } from "fs";
 import fs from "node:fs/promises";
 import path from "node:path";
 import * as ort from "onnxruntime-node";
+import { messagePortMainEndpoint } from "../utils/comlink-endpoint";
 import { ensure, wait } from "../utils/common";
 import { writeStream } from "../utils/stream";
 
-const nodeEndpoint = require("comlink/dist/umd/node-adapter");
-
 /**
  * We cannot do
  *
@@ -63,7 +62,6 @@ process.parentPort.once("message", (e) => {
     parseInitData(e.data);
 
     const port = ensure(e.ports[0]);
-    port.on("message", (me: Electron.MessageEvent) => {});
     expose(
         {
             computeCLIPImageEmbedding,
@@ -71,8 +69,7 @@ process.parentPort.once("message", (e) => {
             detectFaces,
             computeFaceEmbeddings,
         },
-        // eslint-disable-next-line @typescript-eslint/no-unsafe-argument, @typescript-eslint/no-explicit-any
-        nodeEndpoint(port as unknown as any),
+        messagePortMainEndpoint(port),
     );
     // port.on("message", (request) => {
     //     void handleMessageFromRenderer(request.data).then((response) =>
diff --git a/desktop/src/main/utils/comlink-endpoint.ts b/desktop/src/main/utils/comlink-endpoint.ts
index ee87d4176e..4572837306 100644
--- a/desktop/src/main/utils/comlink-endpoint.ts
+++ b/desktop/src/main/utils/comlink-endpoint.ts
@@ -5,48 +5,24 @@ import type { MessagePortMain } from "electron";
  * An adaptation of the `nodeEndpoint` function from comlink suitable for use in
  * TypeScript with an Electron utility process.
  *
- * This is an adaption of the 
+ * This is an adaption of the following function from comlink:
+ * https://github.com/GoogleChromeLabs/comlink/blob/main/src/node-adapter.ts
  *
- * Comlink provides a `nodeEndpoint` [function][1] to allow a Node worker_thread
- * to be treated as an {@link Endpoint} and be used with comlink.
- *
- * The first issue we run into when using it is that this the function is not
- * exported as part of the normal comlink.d.ts. Accessing it via this
- * [workaround][2] doesn't work for us either since we cannot currently change
- * our package type to "module".
- *
- * We could skirt around that by doing
- *
- *     const nodeEndpoint = require("comlink/dist/umd/node-adapter");
- *
- * and silencing tsc and eslint. However, we then run into a different issue:
- * the comlink implementation of the adapter adds an extra layer of nesting.
- * This line:
- *
- *       eh({ data } as MessageEvent);
- *
- * Should be
- *
- *       eh(data)
- *
- * I don't currently know if it is because of an impedance mismatch between
- * Node's worker_threads and Electron's UtilityProcesses, or if it is something
- * else that I'm doing wrong somewhere else causing this to happen.
- *
- * To solve both these issues, we create this variant. This also removes the
- * need for us to type cast when passing MessagePortMain.
- *
- * References:
- * 1. https://github.com/GoogleChromeLabs/comlink/blob/main/src/node-adapter.ts
- * 2. https://github.com/GoogleChromeLabs/comlink/pull/542
- * 3. https://github.com/GoogleChromeLabs/comlink/issues/129
+ * It has been modified (somewhat hackily) to be useful with an Electron
+ * MessagePortMain instead of a Node.js worker_thread. Only things that we
+ * currently need have been made to work as you can see by the abuntant type
+ * casts. Caveat emptor.
  */
 export const messagePortMainEndpoint = (mp: MessagePortMain): Endpoint => {
-    const listeners = new WeakMap();
+    type NL = EventListenerOrEventListenerObject;
+    type EL = (data: Electron.MessageEvent) => void;
+    const listeners = new WeakMap<NL, EL>();
     return {
-        postMessage: mp.postMessage.bind(mp),
+        postMessage: (message, transfer) => {
+            mp.postMessage(message, transfer as unknown as MessagePortMain[]);
+        },
         addEventListener: (_, eh) => {
-            const l = (data: Electron.MessageEvent) =>
+            const l: EL = (data) =>
                 "handleEvent" in eh
                     ? eh.handleEvent({ data } as MessageEvent)
                     : eh(data as unknown as MessageEvent);