#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/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

@@ -29,19 +29,26 @@ import { expectContext } from "@/utils/expectContext";
 import {
   readManagedStorage,
   isInitialized as isManagedStorageInitialized,
+  resetInitializationTimestamp as resetManagedStorageInitializationState,
+  initManagedStorage,
 } 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 { SessionValue } from "@/mv3/SessionStorage";
+import { allSettled } from "@/utils/promiseUtils";
 
 /**
  * The latest version of PixieBrix available in the Chrome Web Store, or null if the version hasn't been fetched.
  */
 let _availableVersion: string | null = null;
 
+const initialized = new SessionValue<boolean>(
+  "installerInitialized",
+  import.meta.url,
+);
+
 /**
  * Returns true if this appears to be a Chrome Web Store install and/or the user has an app URL where they're
  * authenticated so the extension can be linked.
@@ -213,7 +220,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 +298,9 @@ export async function handleInstall({
   }
 }
 
-function onUpdateAvailable({ version }: Runtime.OnUpdateAvailableDetailsType) {
+function setAvailableVersion({
+  version,
+}: Runtime.OnUpdateAvailableDetailsType): void {
   _availableVersion = version;
 }
 
@@ -320,10 +329,45 @@ 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);
+/**
+ * Run initialization that should occur once per extension session.
+ */
+async function initSessionOnce(): Promise<void> {
+  // 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
+
+  // Can't use onStartup because it doesn't mean "on extension start", but "on browser profile start".
+  // Events registered onStartup won't run when the user first installs or when they re-enable the extension.
+  // See: https://developer.chrome.com/docs/extensions/reference/api/runtime#event-onStartup
+
+  if (await initialized.get()) {
+    // Session already initialized
+    return;
+  }
+
+  console.debug("Running extension session initialization");
+
+  const storagePromise = (async () => {
+    await resetManagedStorageInitializationState();
+    await initManagedStorage();
+  })();
+
+  const telemetryPromise = (async () => {
+    if (await isLinked()) {
+      // Init telemetry is a no-op if not linked. So calling without being linked just delays the throttle
+      await initTelemetry();
+    }
+  })();
+
+  await allSettled([storagePromise, telemetryPromise], { catch: "ignore" });
+}
+
+function initInstaller(): void {
+  void initSessionOnce();
+
+  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/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",
     });
   });