#7784: skip waiting for managed storage after initial wait

src/__mocks__/browserMock.mjs

@@ -17,6 +17,7 @@
 
 export default {
   runtime: {
+    id: "abcxyz",
     getManifest: () => ({
       homepage_url: "https://www.pixiebrix.com/",
     }),

src/auth/useLinkState.ts

@@ -23,7 +23,7 @@ import {
 import useAsyncExternalStore from "@/hooks/useAsyncExternalStore";
 import type { AsyncState } from "@/types/sliceTypes";
 
-// NOTE: can't share subscribe methods across generators currently for useAsyncExternalStore, because it maintains
+// NOTE: can't share subscribe methods across generators currently for useAsyncExternalStore because it maintains
 // a map of subscriptions to state controllers. See https://github.com/pixiebrix/pixiebrix-extension/issues/7789
 const subscribe = (callback: () => void) => {
   addAuthListener(callback);

src/auth/usePartnerAuthData.ts

@@ -24,7 +24,7 @@ import useAsyncExternalStore from "@/hooks/useAsyncExternalStore";
 import type { AsyncState } from "@/types/sliceTypes";
 import type { PartnerAuthData } from "@/auth/authTypes";
 
-// NOTE: can't share subscribe methods across generators currently for useAsyncExternalStore, because it maintains
+// NOTE: can't share subscribe methods across generators currently for useAsyncExternalStore because it maintains
 // a map of subscriptions to state controllers. See https://github.com/pixiebrix/pixiebrix-extension/issues/7789
 const subscribe = (callback: () => void) => {
   addAuthListener(callback);

src/auth/useRequiredPartnerAuth.ts

@@ -239,21 +239,6 @@ function useRequiredPartnerAuth(): RequiredPartnerState {
     authMethodOverride === "partner-oauth2" ||
     authMethodOverride === "partner-token";
 
-  console.debug("useRequiredPartnerAuth", {
-    partnerAuthState,
-    hasPartner,
-    requiresIntegration,
-    isMeLoading,
-    isLinkedLoading,
-    meError,
-    partnerKey: partner?.theme ?? managedPartnerId,
-    hasConfiguredIntegration: {
-      requiresIntegration,
-      partnerConfiguration,
-      isMissingPartnerJwt,
-    },
-  });
-
   return {
     hasPartner,
     partnerKey: partner?.theme ?? managedPartnerId,

src/background/background.ts

@@ -47,23 +47,13 @@ import { initRuntimeLogging } from "@/development/runtimeLogging";
 import initWalkthroughModalTrigger from "@/background/walkthroughModalTrigger";
 import { initSidePanel } from "./sidePanel";
 import initRestrictUnauthenticatedUrlAccess from "@/background/restrictUnauthenticatedUrlAccess";
-import {
-  initManagedStorage,
-  watchDelayedStorageInitialization,
-} from "@/store/enterprise/managedStorage";
 import { setPlatform } from "@/platform/platformContext";
 import backgroundPlatform from "@/background/backgroundPlatform";
 
 // The background "platform" currently is used to execute API requests from Google Sheets/Automation Anywhere.
 // In the future, it might also run other background tasks from mods (e.g., background intervals)
 setPlatform(backgroundPlatform);
 
-// Try to initialize managed storage as early as possible because it impacts background behavior
-// Call watchDelayedStorageInitialization to handle case where storage is not immediately available within timeout.
-// We might consider putting watchStorageInitialization in initManagedStorage, but having a non-terminating
-// interval complicates testing.
-void initManagedStorage().then(async () => watchDelayedStorageInitialization());
-
 void initLocator();
 void initMessengerLogging();
 void initRuntimeLogging();

src/background/deploymentUpdater.test.ts

@@ -123,7 +123,7 @@ beforeEach(async () => {
     organizationId: "00000000-00000000-00000000-00000000",
   } as any);
 
-  resetManagedStorage();
+  await resetManagedStorage();
 });
 
 describe("updateDeployments", () => {

src/background/installer.test.ts

@@ -18,7 +18,7 @@
 import {
   requirePartnerAuth,
   openInstallPage,
-  handleInstall,
+  showInstallPage,
 } from "@/background/installer";
 import * as auth from "@/auth/authStorage";
 import { locator } from "@/background/locator";
@@ -56,9 +56,9 @@ const getUserData = jest.mocked(auth.getUserData);
 const locateAllForServiceMock = jest.mocked(locator.locateAllForService);
 const browserManagedStorageMock = jest.mocked(browser.storage.managed.get);
 
-afterEach(() => {
+afterEach(async () => {
   jest.clearAllMocks();
-  resetManagedStorage();
+  await resetManagedStorage();
 });
 
 describe("openInstallPage", () => {
@@ -230,7 +230,7 @@ describe("handleInstall", () => {
     // App setup tab isn't open
     queryTabsMock.mockResolvedValue([]);
     isLinkedMock.mockResolvedValue(false);
-    await handleInstall({
+    await showInstallPage({
       reason: "install",
       previousVersion: undefined,
       temporary: false,
@@ -243,7 +243,7 @@ describe("handleInstall", () => {
     // App setup tab isn't open
     queryTabsMock.mockResolvedValue([]);
     isLinkedMock.mockResolvedValue(true);
-    await handleInstall({
+    await showInstallPage({
       reason: "install",
       previousVersion: undefined,
       temporary: false,
@@ -260,7 +260,7 @@ describe("handleInstall", () => {
       });
       queryTabsMock.mockResolvedValue([]);
       isLinkedMock.mockResolvedValue(false);
-      await handleInstall({
+      await showInstallPage({
         reason: "install",
         previousVersion: undefined,
         temporary: false,

src/background/installer.ts

@@ -27,15 +27,17 @@ import { isCommunityControlRoom } from "@/contrib/automationanywhere/aaUtils";
 import { isEmpty } from "lodash";
 import { expectContext } from "@/utils/expectContext";
 import {
-  readManagedStorage,
+  initManagedStorage,
   isInitialized as isManagedStorageInitialized,
+  readManagedStorage,
+  resetInitializationTimestamp as resetManagedStorageInitializationState,
+  watchForDelayedStorageInitialization,
 } from "@/store/enterprise/managedStorage";
 import { Events } from "@/telemetry/events";
-
 import { DEFAULT_SERVICE_URL, UNINSTALL_URL } from "@/urlConstants";
-
 import { CONTROL_ROOM_TOKEN_INTEGRATION_ID } from "@/integrations/constants";
 import { getExtensionConsoleUrl } from "@/utils/extensionUtils";
+import { oncePerSession } from "@/mv3/SessionStorage";
 
 /**
  * The latest version of PixieBrix available in the Chrome Web Store, or null if the version hasn't been fetched.
@@ -213,7 +215,7 @@ export async function requirePartnerAuth(): Promise<void> {
 }
 
 // Exported for testing
-export async function handleInstall({
+export async function showInstallPage({
   reason,
   previousVersion,
 }: Runtime.OnInstalledDetailsType): Promise<void> {
@@ -291,7 +293,9 @@ export async function handleInstall({
   }
 }
 
-function onUpdateAvailable({ version }: Runtime.OnUpdateAvailableDetailsType) {
+function setAvailableVersion({
+  version,
+}: Runtime.OnUpdateAvailableDetailsType): void {
   _availableVersion = version;
 }
 
@@ -320,10 +324,38 @@ async function setUninstallURL(): Promise<void> {
   await browser.runtime.setUninstallURL(url.href);
 }
 
-function initInstaller() {
-  browser.runtime.onUpdateAvailable.addListener(onUpdateAvailable);
-  browser.runtime.onInstalled.addListener(handleInstall);
-  browser.runtime.onStartup.addListener(initTelemetry);
+// Using our own session value vs. webext-events because onExtensionStart has a 100ms delay
+// https://github.com/fregante/webext-events/blob/main/source/on-extension-start.ts#L56
+// eslint-disable-next-line local-rules/persistBackgroundData -- using SessionMp via oncePerSession
+const initManagedStorageOnce = oncePerSession(
+  "initManagedStorage",
+  import.meta.url,
+  async () => {
+    await resetManagedStorageInitializationState();
+    await initManagedStorage();
+    void watchForDelayedStorageInitialization();
+  },
+);
+
+// eslint-disable-next-line local-rules/persistBackgroundData -- using SessionMp via oncePerSession
+const initTelemetryOnce = oncePerSession(
+  "initTelemetry",
+  import.meta.url,
+  async () => {
+    if (await isLinked()) {
+      // Init telemetry is a no-op if not linked. So calling without being linked just delays the throttle
+      await initTelemetry();
+    }
+  },
+);
+
+function initInstaller(): void {
+  void initManagedStorageOnce();
+  void initTelemetryOnce();
+
+  browser.runtime.onInstalled.addListener(showInstallPage);
+  browser.runtime.onUpdateAvailable.addListener(setAvailableVersion);
+
   dntConfig.onChanged(() => {
     void setUninstallURL();
   });

src/background/restrictUnauthenticatedUrlAccess.test.ts

@@ -65,7 +65,7 @@ describe("enforceAuthentication", () => {
 
   afterEach(async () => {
     // eslint-disable-next-line new-cap -- used for testing
-    INTERNAL_reset();
+    await INTERNAL_reset();
     await browser.storage.managed.clear();
     await browser.storage.local.clear();
     axiosMock.reset();

src/background/telemetry.ts

@@ -17,7 +17,6 @@
 
 import { type JsonObject } from "type-fest";
 import { compact, debounce, throttle, uniq } from "lodash";
-import { isLinked } from "@/auth/authStorage";
 import { getModComponentState } from "@/store/extensionsStorage";
 import {
   getLinkedApiClient,
@@ -305,8 +304,9 @@ async function collectUserSummary(): Promise<UserSummary> {
 }
 
 async function init(): Promise<void> {
-  if ((await isLinked()) && (await allowsTrack())) {
-    const client = await getLinkedApiClient();
+  const client = await maybeGetLinkedApiClient();
+
+  if (client && (await allowsTrack())) {
     await client.post("/api/identify/", {
       uid: await getUUID(),
       data: await collectUserSummary(),
@@ -321,7 +321,7 @@ export const initTelemetry = throttle(init, 30 * 60 * 1000, {
 });
 
 /**
- * @deprecated Only allowed in @/background files. Otherwise use: `import reportEvent from "@/telemetry/reportEvent"`
+ * @deprecated Only allowed in @/background files. Otherwise, use: `import reportEvent from "@/telemetry/reportEvent"`
  */
 export async function recordEvent({
   event,

src/components/logViewer/LogTable.stories.tsx

@@ -26,8 +26,6 @@ import { InputValidationError } from "@/bricks/errors";
 import { type Schema } from "@/types/schemaTypes";
 import type { LogEntry } from "@/telemetry/logging";
 
-Object.assign(global, { chrome: { runtime: { id: 42 } } });
-
 export default {
   title: "Editor/LogTable",
   component: LogTable,

src/extensionConsole/Navbar.tsx

@@ -38,7 +38,7 @@ import useLinkState from "@/auth/useLinkState";
 import { DEFAULT_SERVICE_URL } from "@/urlConstants";
 import useAsyncExternalStore from "@/hooks/useAsyncExternalStore";
 
-// NOTE: can't share subscribe methods across generators currently for useAsyncExternalStore, because it maintains
+// NOTE: can't share subscribe methods across generators currently for useAsyncExternalStore because it maintains
 // a map of subscriptions to state controllers. See https://github.com/pixiebrix/pixiebrix-extension/issues/7789
 const subscribe = (callback: () => void) => {
   addAuthListener(callback);

src/extensionConsole/pages/BrowserBanner.test.tsx

@@ -24,7 +24,7 @@ import { INTERNAL_reset } from "@/store/enterprise/managedStorage";
 
 beforeEach(async () => {
   // eslint-disable-next-line new-cap -- test helper method
-  INTERNAL_reset();
+  await INTERNAL_reset();
   await browser.storage.managed.clear();
 });
 

src/extensionConsole/pages/onboarding/SetupPage.test.tsx

@@ -66,7 +66,7 @@ jest.mock("@/data/service/baseService", () => ({
 
 beforeEach(async () => {
   jest.clearAllMocks();
-  resetManagedStorage();
+  await resetManagedStorage();
   await browser.storage.managed.clear();
 });
 

src/mv3/SessionStorage.ts

@@ -26,6 +26,7 @@ import { expectContext } from "@/utils/expectContext";
 import { type OmitIndexSignature, type JsonValue } from "type-fest";
 import { type ManualStorageKey } from "@/utils/storageUtils";
 import { once } from "lodash";
+import pMemoize from "p-memoize";
 
 // Just like chrome.storage.session, this must be "global"
 // eslint-disable-next-line local-rules/persistBackgroundData -- MV2-only
@@ -58,6 +59,19 @@ export class SessionMap<Value extends JsonValue> {
     return `${this.key}::${this.url}::${secondaryKey}` as ManualStorageKey;
   }
 
+  async has(secondaryKey: string): Promise<boolean> {
+    this.validateContext();
+    const rawStorageKey = this.getRawStorageKey(secondaryKey);
+    if (!hasSession) {
+      return storage.has(rawStorageKey);
+    }
+
+    const result = await browser.storage.session.get(rawStorageKey);
+    // OK to use `undefined` because undefined is not a valid storage value
+    // eslint-disable-next-line security/detect-object-injection -- `getRawStorageKey` ensures the format
+    return result[rawStorageKey] !== undefined;
+  }
+
   async get(secondaryKey: string): Promise<Value | undefined> {
     this.validateContext();
     const rawStorageKey = this.getRawStorageKey(secondaryKey);
@@ -113,3 +127,19 @@ export class SessionValue<Value extends OmitIndexSignature<JsonValue>> {
     await this.map.set("#value", value);
   }
 }
+
+/**
+ * Helper to run a method at most once per session from the background page.
+ * @param key the SessionMap key
+ * @param url the ImportMeta.url of the file
+ * @param fn the function to run once per session
+ */
+export const oncePerSession = (
+  key: string,
+  url: string,
+  fn: () => Promise<unknown>,
+) =>
+  pMemoize(fn, {
+    // `fn` has no arguments, so only one value is stored
+    cache: new SessionMap(key, url),
+  });

src/store/enterprise/managedStorage.test.ts

@@ -16,25 +16,43 @@
  */
 
 import {
+  initializationTimestamp,
   INTERNAL_reset,
   readManagedStorage,
   readManagedStorageByKey,
 } from "@/store/enterprise/managedStorage";
+import type { Timestamp } from "@/types/stringTypes";
 
 beforeEach(async () => {
+  jest.clearAllMocks();
   // eslint-disable-next-line new-cap -- test helper method
-  INTERNAL_reset();
+  await INTERNAL_reset();
   await browser.storage.managed.clear();
 });
 
 describe("readManagedStorage", () => {
+  it("reads immediately if managed storage is already initialized", async () => {
+    await initializationTimestamp.set(new Date().toISOString() as Timestamp);
+    await expect(readManagedStorage()).resolves.toStrictEqual({
+      // `jest-webextension-mock`'s storage is shared across sources, the call ends up with the managed storage
+      // and the local storage mixed together. See https://github.com/clarkbw/jest-webextension-mock/issues/183
+      managedStorageInitTimestamp: expect.any(String),
+    });
+
+    // Should only be called once vs. polling
+    expect(browser.storage.managed.get).toHaveBeenCalledOnce();
+  });
+
   it("reads uninitialized managed storage", async () => {
     await expect(readManagedStorage()).resolves.toStrictEqual({});
   });
 
   it("reads managed storage", async () => {
     await browser.storage.managed.set({ partnerId: "taco-bell" });
     await expect(readManagedStorage()).resolves.toStrictEqual({
+      // `jest-webextension-mock`'s storage is shared across sources, the call ends up with the managed storage
+      // and the local storage mixed together. See https://github.com/clarkbw/jest-webextension-mock/issues/183
+      managedStorageInitTimestamp: expect.any(String),
       partnerId: "taco-bell",
     });
   });