[extension] Fix messaging between app and background (#8981)

extension/app/background.ts

@@ -22,15 +22,9 @@ import { generatePKCE } from "./src/lib/utils";
 const log = console.error;
 
 const state: {
-  port: chrome.runtime.Port | undefined;
-  extensionReady: boolean;
-  inputBarReady: boolean;
   refreshingToken: boolean;
   lastHandler: (() => void) | undefined;
 } = {
-  port: undefined,
-  extensionReady: false,
-  inputBarReady: false,
   refreshingToken: false,
   lastHandler: undefined,
 };
@@ -50,6 +44,7 @@ chrome.runtime.onUpdateAvailable.addListener(async (details) => {
  * Listener to open/close the side panel when the user clicks on the extension icon.
  */
 chrome.runtime.onInstalled.addListener(() => {
+  void chrome.storage.local.set({ extensionReady: false });
   void chrome.sidePanel.setPanelBehavior({ openPanelOnActionClick: true });
   chrome.contextMenus.create({
     id: "ask_dust",
@@ -77,14 +72,13 @@ chrome.runtime.onInstalled.addListener(() => {
 chrome.runtime.onConnect.addListener((port) => {
   if (port.name === "sidepanel-connection") {
     console.log("Sidepanel is there");
-    state.port = port;
-    state.extensionReady = true;
-    port.onDisconnect.addListener(() => {
+    void chrome.storage.local.set({ extensionReady: true });
+    port.onDisconnect.addListener(async () => {
       // This fires when sidepanel closes
       console.log("Sidepanel was closed");
-      state.port = undefined;
-      state.extensionReady = false;
-      state.inputBarReady = false;
+      await chrome.storage.local.set({
+        extensionReady: false,
+      });
       state.lastHandler = undefined;
     });
   }
@@ -94,50 +88,42 @@ const getActionHandler = (menuItemId: string | number) => {
   switch (menuItemId) {
     case "ask_dust":
       return () => {
-        if (state.port) {
-          const params = JSON.stringify({
-            includeContent: true,
-            includeCapture: false,
-            text: ":mention[dust]{sId=dust} summarize this page.",
-            configurationId: "dust",
-          });
-          state.port.postMessage({
-            type: "EXT_ROUTE_CHANGE",
-            pathname: "/run",
-            search: `?${params}`,
-          });
-        }
+        const params = JSON.stringify({
+          includeContent: true,
+          includeCapture: false,
+          text: ":mention[dust]{sId=dust} summarize this page.",
+          configurationId: "dust",
+        });
+        void chrome.runtime.sendMessage({
+          type: "EXT_ROUTE_CHANGE",
+          pathname: "/run",
+          search: `?${params}`,
+        });
       };
     case "add_tab_content":
       return () => {
-        if (state.port) {
-          state.port.postMessage({
-            type: "EXT_ATTACH_TAB",
-            includeContent: true,
-            includeCapture: false,
-          });
-        }
+        void chrome.runtime.sendMessage({
+          type: "EXT_ATTACH_TAB",
+          includeContent: true,
+          includeCapture: false,
+        });
       };
     case "add_tab_screenshot":
       return () => {
-        if (state.port) {
-          state.port.postMessage({
-            type: "EXT_ATTACH_TAB",
-            includeContent: false,
-            includeCapture: true,
-          });
-        }
+        void chrome.runtime.sendMessage({
+          type: "EXT_ATTACH_TAB",
+          includeContent: false,
+          includeCapture: true,
+        });
       };
     case "add_selection":
       return () => {
-        if (state.port) {
-          state.port.postMessage({
-            type: "EXT_ATTACH_TAB",
-            includeContent: true,
-            includeCapture: false,
-            includeSelectionOnly: true,
-          });
-        }
+        void chrome.runtime.sendMessage({
+          type: "EXT_ATTACH_TAB",
+          includeContent: true,
+          includeCapture: false,
+          includeSelectionOnly: true,
+        });
       };
   }
 };
@@ -148,15 +134,17 @@ chrome.contextMenus.onClicked.addListener(async (event, tab) => {
     return;
   }
 
-  if (!state.extensionReady && tab) {
-    // Store the handler for later use when the extension is ready.
-    state.lastHandler = handler;
-    void chrome.sidePanel.open({
-      windowId: tab.windowId,
-    });
-  } else {
-    await handler();
-  }
+  chrome.storage.local.get(["extensionReady"], ({ extensionReady }) => {
+    if (!extensionReady && tab) {
+      // Store the handler for later use when the extension is ready.
+      state.lastHandler = handler;
+      void chrome.sidePanel.open({
+        windowId: tab.windowId,
+      });
+    } else {
+      void handler();
+    }
+  });
 });
 
 function capture(sendResponse: (x: CaptureResponse) => void) {
@@ -239,7 +227,7 @@ chrome.runtime.onMessage.addListener(
                     files: ["page.js"],
                   });
                   captures = await new Promise((resolve, reject) => {
-                    if (state?.port && tab?.id) {
+                    if (tab?.id) {
                       const timeout = setTimeout(() => {
                         console.error("Timeout waiting for full page capture");
                         reject(
@@ -255,10 +243,8 @@ chrome.runtime.onMessage.addListener(
                         }
                       );
                     } else {
-                      console.error("No port or tab id");
-                      reject(
-                        new Error("Failed to get content from the current tab.")
-                      );
+                      console.error("No tab id");
+                      reject(new Error("No tab selected."));
                     }
                   });
                 } else {
@@ -323,8 +309,7 @@ chrome.runtime.onMessage.addListener(
 
       case "INPUT_BAR_STATUS":
         // Enable or disable the context menu items based on the input bar status. Actions are only available when the input bar is visible.
-        state.inputBarReady = message.available;
-        if (state.lastHandler && state.inputBarReady) {
+        if (state.lastHandler && message.available) {
           state.lastHandler();
           state.lastHandler = undefined;
         }

extension/app/page.ts

@@ -212,12 +212,23 @@ declare global {
         body.style.overflowY = originalBodyOverflowYStyle;
       }
       window.scrollTo(originalX, originalY);
+      // Ensure the callback is called at least once when cleaning up.
+      // If it was called already with results, this will be ignored by promise.
+      callback([]);
     }
     const screenshots: Screenshot[] = [];
 
     (function processArrangements() {
       const next = arrangements.shift();
       if (!next) {
+        if (callback) {
+          callback(
+            screenshots.map((screenshot) =>
+              screenshot.canvas.toDataURL("image/jpeg", 0.8)
+            )
+          );
+        }
+
         cleanUp();
         return;
       }
@@ -289,17 +300,7 @@ declare global {
                   image.height * zoomFactor
                 );
               });
-              // send back log data for debugging (but keep it truthy to
-              // indicate success)
-              if (!arrangements.length) {
-                if (callback) {
-                  callback(
-                    screenshots.map((screenshot) =>
-                      screenshot.canvas.toDataURL("image/jpeg", 0.8)
-                    )
-                  );
-                }
-              }
+
               // Move on to capture next arrangement.
               processArrangements();
             };

extension/app/src/components/PortContext.tsx

@@ -6,6 +6,7 @@ export const PortProvider = ({ children }: { children: React.ReactNode }) => {
   const [port, setPort] = useState<chrome.runtime.Port | null>(null);
 
   useEffect(() => {
+    console.log("Connecting to sidepanel");
     const port = chrome.runtime.connect({ name: "sidepanel-connection" });
 
     setPort(port);

extension/app/src/components/auth/ProtectedRoute.tsx

@@ -6,12 +6,11 @@ import {
   Spinner,
 } from "@dust-tt/sparkle";
 import { useAuth } from "@extension/components/auth/AuthProvider";
-import { PortContext } from "@extension/components/PortContext";
 import type { RouteChangeMesssage } from "@extension/lib/messages";
 import type { StoredUser } from "@extension/lib/storage";
 import { getPendingUpdate } from "@extension/lib/storage";
 import type { ReactNode } from "react";
-import { useContext, useEffect, useState } from "react";
+import { useEffect, useState } from "react";
 import { useNavigate } from "react-router-dom";
 
 type ProtectedRouteProps = {
@@ -37,22 +36,19 @@ export const ProtectedRoute = ({ children }: ProtectedRouteProps) => {
   const navigate = useNavigate();
   const [isLatestVersion, setIsLatestVersion] = useState(true);
 
-  const port = useContext(PortContext);
   useEffect(() => {
-    if (port) {
-      const listener = (message: RouteChangeMesssage) => {
-        const { type } = message;
-        if (type === "EXT_ROUTE_CHANGE") {
-          navigate({ pathname: message.pathname, search: message.search });
-          return false;
-        }
-      };
-      port.onMessage.addListener(listener);
-      return () => {
-        port.onMessage.removeListener(listener);
-      };
-    }
-  }, [port, navigate]);
+    const listener = (message: RouteChangeMesssage) => {
+      const { type } = message;
+      if (type === "EXT_ROUTE_CHANGE") {
+        navigate({ pathname: message.pathname, search: message.search });
+        return false;
+      }
+    };
+    chrome.runtime.onMessage.addListener(listener);
+    return () => {
+      chrome.runtime.onMessage.removeListener(listener);
+    };
+  }, [navigate]);
 
   useEffect(() => {
     if (!isAuthenticated || !isUserSetup || !user || !workspace) {

extension/app/src/components/input_bar/InputBar.tsx

@@ -12,7 +12,6 @@ import { InputBarCitations } from "@extension/components/input_bar/InputBarCitat
 import type { InputBarContainerProps } from "@extension/components/input_bar/InputBarContainer";
 import { InputBarContainer } from "@extension/components/input_bar/InputBarContainer";
 import { InputBarContext } from "@extension/components/input_bar/InputBarContext";
-import { PortContext } from "@extension/components/PortContext";
 import { useFileUploaderService } from "@extension/hooks/useFileUploaderService";
 import { useDustAPI } from "@extension/lib/dust_api";
 import type { AttachSelectionMessage } from "@extension/lib/messages";
@@ -63,23 +62,21 @@ export function AssistantInputBar({
     getFileBlobs,
     resetUpload,
   } = fileUploaderService;
-  const port = useContext(PortContext);
+
   useEffect(() => {
-    if (port) {
-      void sendInputBarStatus(true);
-      const listener = async (message: AttachSelectionMessage) => {
-        const { type } = message;
-        if (type === "EXT_ATTACH_TAB") {
-          // Handle message
-          void uploadContentTab(message);
-        }
-      };
-      port.onMessage.addListener(listener);
-      return () => {
-        void sendInputBarStatus(false);
-        port.onMessage.removeListener(listener);
-      };
-    }
+    void sendInputBarStatus(true);
+    const listener = async (message: AttachSelectionMessage) => {
+      const { type } = message;
+      if (type === "EXT_ATTACH_TAB") {
+        // Handle message
+        void uploadContentTab(message);
+      }
+    };
+    chrome.runtime.onMessage.addListener(listener);
+    return () => {
+      void sendInputBarStatus(false);
+      chrome.runtime.onMessage.removeListener(listener);
+    };
   }, []);
 
   const { droppedFiles, setDroppedFiles } = useFileDrop();

extension/app/src/lib/messages.ts

@@ -173,7 +173,7 @@ export const sendGetActiveTabMessage = (params: GetActiveTabOptions) => {
 };
 
 export const sendInputBarStatus = (available: boolean) => {
-  return sendMessage<InputBarStatusMessage, void>({
+  void chrome.runtime.sendMessage({
     type: "INPUT_BAR_STATUS",
     available,
   });