fix(browser): handle port disconnection b/w client and worker

src/core/client.mjs

@@ -116,11 +116,12 @@ const initMenu = async (db) => {
 };
 
 export default async (api, db) => {
+  const { sendMessage } = globalThis.__enhancerApi;
   await Promise.all([
     overrideThemes(db),
     insertCustomStyles(db),
     initMenu(db),
     sendTelemetryPing(),
   ]);
-  api.sendMessage("notion-enhancer", "load-complete");
+  sendMessage("notion-enhancer", "load-complete");
 };

src/shared/system.js

@@ -33,17 +33,24 @@ const platform = IS_ELECTRON
     IS_ELECTRON && !IS_RENDERER ? require(`../../../${target}`) : undefined;
 
 let __port;
-const onMessage = (channel, listener) => {
+const connectToPort = () => {
+    if (__port) return;
+    __port = chrome.runtime.connect();
+    __port.onDisconnect.addListener(() => (__port = null));
+  },
+  onMessage = (channel, listener) => {
     // from worker to client
     if (IS_RENDERER) {
       const { ipcRenderer } = require("electron");
       ipcRenderer.on(channel, listener);
     } else if (!IS_ELECTRON) {
-      __port ??= chrome.runtime.connect();
-      __port.onMessage.addListener((msg) => {
+      const onMessage = (msg) => {
         if (msg?.channel !== channel || msg?.invocation) return;
         listener(msg.message);
-      });
+      };
+      connectToPort();
+      __port.onMessage.addListener(onMessage);
+      chrome.runtime.onMessage.addListener(onMessage);
     }
   },
   sendMessage = (channel, message) => {
@@ -52,7 +59,7 @@ const onMessage = (channel, listener) => {
       const { ipcRenderer } = require("electron");
       ipcRenderer.send(channel, message);
     } else if (!IS_ELECTRON) {
-      __port ??= chrome.runtime.connect();
+      connectToPort();
       __port.postMessage({ channel, message });
     }
   },
@@ -67,7 +74,7 @@ const onMessage = (channel, listener) => {
       // the browser: uses a long-lived ipc connection to
       // pass messages and handle responses asynchronously
       let fulfilled;
-      __port ??= chrome.runtime.connect();
+      connectToPort();
       const id = crypto.randomUUID();
       return new Promise((res, rej) => {
         __port.onMessage.addListener((msg) => {
@@ -118,7 +125,7 @@ const initDatabase = (namespace, fallbacks = {}) => {
     const query = (query, args = {}) =>
       IS_ELECTRON && !IS_RENDERER
         ? globalThis.__enhancerApi.queryDatabase(namespace, query, args)
-        : invokeInWorker("notion-enhancer:db", {
+        : invokeInWorker("notion-enhancer", {
             action: "query-database",
             data: { namespace, query, args },
           });

src/worker.js

@@ -116,7 +116,7 @@ const initDatabase = async () => {
 if (IS_ELECTRON) {
   const { ipcMain } = require("electron"),
     { reloadApp } = globalThis.__enhancerApi;
-  ipcMain.handle("notion-enhancer:db", ({}, message) => {
+  ipcMain.handle("notion-enhancer", ({}, message) => {
     if (message?.action !== "query-database") return;
     const { namespace, query, args } = message.data;
     return queryDatabase(namespace, query, args);
@@ -128,19 +128,19 @@ if (IS_ELECTRON) {
   const notionUrl = "https://www.notion.so/",
     isNotionTab = (tab) => tab?.url?.startsWith(notionUrl);
 
-  const tabQueue = new Set(),
+  const connectedTabs = new Set(),
+    openMenuInTabs = new Set(),
     openMenu = { channel: "notion-enhancer", message: "open-menu" },
     openEnhancerMenu = async (tab) => {
       if (!isNotionTab(tab)) {
-        const windowId = chrome.windows.WINDOW_ID_CURRENT,
-          windowTabs = await chrome.tabs.query({ windowId });
-        tab = windowTabs.find(isNotionTab);
+        const windowId = chrome.windows.WINDOW_ID_CURRENT;
+        tab = (await chrome.tabs.query({ windowId })).find(isNotionTab);
         tab ??= await chrome.tabs.create({ url: notionUrl });
       }
       chrome.tabs.highlight({ tabs: [tab.index] });
-      if (tab.status === "complete") {
+      if (connectedTabs.has(tab.id)) {
         chrome.tabs.sendMessage(tab.id, openMenu);
-      } else tabQueue.add(tab.id);
+      } else openMenuInTabs.add(tab.id);
     },
     reloadNotionTabs = async () => {
       const windowId = chrome.windows.WINDOW_ID_CURRENT;
@@ -149,24 +149,29 @@ if (IS_ELECTRON) {
         .forEach((tab) => chrome.tabs.reload(tab.id));
     };
 
-  // listen for invoke: https://developer.chrome.com/docs/extensions/mv3/messaging/
   chrome.action.onClicked.addListener(openEnhancerMenu);
+  // long-lived connection for rapid two-way messaging
+  // b/w client and worker, primarily used for db wrapper:
+  // https://developer.chrome.com/docs/extensions/mv3/messaging/
   chrome.runtime.onConnect.addListener((port) => {
-    port.onMessage.addListener((msg, sender) => {
+    const tabId = port.sender.tab.id;
+    connectedTabs.add(tabId);
+    port.onMessage.addListener(async (msg) => {
       if (msg?.channel !== "notion-enhancer") return;
       const { message, invocation } = msg;
       if (message.action === "query-database") {
         const { namespace, query, args } = message.data,
-          res = queryDatabase(namespace, query, args);
+          res = await queryDatabase(namespace, query, args);
         if (invocation) port.postMessage({ invocation, message: res });
       }
       if (message === "load-complete") {
-        if (!tabQueue.has(sender?.tab?.id)) return;
-        chrome.tabs.sendMessage(sender.tab.id, openMenu);
-        tabQueue.delete(sender.tab.id);
+        if (!openMenuInTabs.has(tabId)) return;
+        openMenuInTabs.delete(tabId);
+        port.postMessage(openMenu);
       }
       if (message === "reload-app") reloadNotionTabs();
     });
+    port.onDisconnect.addListener(() => connectedTabs.delete(tabId));
   });
 }