diff --git a/extension/app/background.ts b/extension/app/background.ts index 42329dd91bfb..26ed8b1f6232 100644 --- a/extension/app/background.ts +++ b/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; } diff --git a/extension/app/page.ts b/extension/app/page.ts index 26a2cecae77c..e797087b9d94 100644 --- a/extension/app/page.ts +++ b/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(); }; diff --git a/extension/app/src/components/PortContext.tsx b/extension/app/src/components/PortContext.tsx index df4e6930685c..04e69e5be595 100644 --- a/extension/app/src/components/PortContext.tsx +++ b/extension/app/src/components/PortContext.tsx @@ -6,6 +6,7 @@ export const PortProvider = ({ children }: { children: React.ReactNode }) => { const [port, setPort] = useState(null); useEffect(() => { + console.log("Connecting to sidepanel"); const port = chrome.runtime.connect({ name: "sidepanel-connection" }); setPort(port); diff --git a/extension/app/src/components/auth/ProtectedRoute.tsx b/extension/app/src/components/auth/ProtectedRoute.tsx index 6547dbfc3ab4..f2101faae5b4 100644 --- a/extension/app/src/components/auth/ProtectedRoute.tsx +++ b/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) { diff --git a/extension/app/src/components/input_bar/InputBar.tsx b/extension/app/src/components/input_bar/InputBar.tsx index 9fa6036a81ef..96066b543387 100644 --- a/extension/app/src/components/input_bar/InputBar.tsx +++ b/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(); diff --git a/extension/app/src/lib/messages.ts b/extension/app/src/lib/messages.ts index e786beb41ad0..c95dca4d53ef 100644 --- a/extension/app/src/lib/messages.ts +++ b/extension/app/src/lib/messages.ts @@ -173,7 +173,7 @@ export const sendGetActiveTabMessage = (params: GetActiveTabOptions) => { }; export const sendInputBarStatus = (available: boolean) => { - return sendMessage({ + void chrome.runtime.sendMessage({ type: "INPUT_BAR_STATUS", available, });