#7784: skip waiting for managed storage after initial wait (#7790)

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 SessionMap via oncePerSession
+const initManagedStorageOncePerSession = oncePerSession(
+  "initManagedStorage",
+  import.meta.url,
+  async () => {
+    await resetManagedStorageInitializationState();
+    await initManagedStorage();
+    void watchForDelayedStorageInitialization();
+  },
+);
+
+// eslint-disable-next-line local-rules/persistBackgroundData -- using SessionMap via oncePerSession
+const initTelemetryOncePerSession = 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 initManagedStorageOncePerSession();
+  void initTelemetryOncePerSession();
+
+  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.test.ts

@@ -23,21 +23,49 @@ test("SessionMap", async () => {
 
   await map.set("alpha", 1);
   await expect(map.get("alpha")).resolves.toBe(1);
+  await expect(map.has("alpha")).resolves.toBeTrue();
 
   // Other props should be left untouched
   await expect(map.get("beta")).resolves.toBeUndefined();
 
+  await map.delete("alpha");
+  await expect(map.get("alpha")).resolves.toBeUndefined();
+  await expect(map.has("alpha")).resolves.toBeFalse();
+});
+
+test("SessionMap accepts undefined", async () => {
+  const map = new SessionMap("jester", import.meta.url);
+  await expect(map.get("alpha")).resolves.toBeUndefined();
+
+  await map.set("alpha", 1);
+  await expect(map.get("alpha")).resolves.toBe(1);
+  await expect(map.has("alpha")).resolves.toBeTrue();
+
+  // `undefined` is a type-error because it's not a JsonValue. Keeping this test just to document the current behavior.
   await map.set("alpha", undefined);
   await expect(map.get("alpha")).resolves.toBeUndefined();
+  // SessionMap will unset the value
+  await expect(map.has("alpha")).resolves.toBeFalse();
 });
 
 test("SessionValue", async () => {
-  const map = new SessionValue("jester", import.meta.url);
-  await expect(map.get()).resolves.toBeUndefined();
+  const value = new SessionValue("jester", import.meta.url);
+  await expect(value.get()).resolves.toBeUndefined();
+
+  await value.set(1);
+  await expect(value.get()).resolves.toBe(1);
+
+  await value.unset();
+  await expect(value.get()).resolves.toBeUndefined();
+});
+
+test("SessionValue allows setting undefined", async () => {
+  const value = new SessionValue("jester", import.meta.url);
 
-  await map.set(1);
-  await expect(map.get()).resolves.toBe(1);
+  await value.set(1);
+  await expect(value.get()).resolves.toBe(1);
 
-  await map.set(undefined);
-  await expect(map.get()).resolves.toBeUndefined();
+  // `undefined` is a type-error because it's not a JsonValue. Keeping this test just to document the current behavior.
+  await value.set(undefined);
+  await expect(value.get()).resolves.toBeUndefined();
 });