fix: implement consent-aware default user provider and remote flag storage

packages/experiment-browser/src/experimentClient.ts

@@ -176,7 +176,7 @@ export class ExperimentClient implements Client {
       ? `${this.config.instanceName}-${internalInstanceName}`
       : this.config.instanceName;
     if (this.isWebExperiment) {
-      storage = new SessionStorage();
+      storage = config?.['consentAwareStorage']?.['sessionStorage'];
     } else {
       storage = new LocalStorage();
     }

packages/experiment-browser/src/factory.ts

@@ -71,7 +71,12 @@ const newExperimentClient = (
 ): ExperimentClient => {
   return new ExperimentClient(apiKey, {
     ...config,
-    userProvider: new DefaultUserProvider(config?.userProvider, apiKey),
+    userProvider: new DefaultUserProvider(
+      config?.userProvider,
+      apiKey,
+      config?.['consentAwareStorage']?.localStorage,
+      config?.['consentAwareStorage']?.sessionStorage,
+    ),
   });
 };
 

packages/experiment-browser/src/providers/default.ts

@@ -4,6 +4,7 @@ import { UAParser } from '@amplitude/ua-parser-js';
 import { LocalStorage } from '../storage/local-storage';
 import { SessionStorage } from '../storage/session-storage';
 import { ExperimentUserProvider } from '../types/provider';
+import { Storage } from '../types/storage';
 import { ExperimentUser } from '../types/user';
 
 export class DefaultUserProvider implements ExperimentUserProvider {
@@ -13,17 +14,24 @@ export class DefaultUserProvider implements ExperimentUserProvider {
       ? this.globalScope?.navigator.userAgent
       : undefined;
   private readonly ua = new UAParser(this.userAgent).getResult();
-  private readonly localStorage = new LocalStorage();
-  private readonly sessionStorage = new SessionStorage();
+  private readonly localStorage: Storage;
+  private readonly sessionStorage: Storage;
   private readonly storageKey: string;
 
   public readonly userProvider: ExperimentUserProvider | undefined;
   private readonly apiKey?: string;
 
-  constructor(userProvider?: ExperimentUserProvider, apiKey?: string) {
+  constructor(
+    userProvider?: ExperimentUserProvider,
+    apiKey?: string,
+    customLocalStorage?: Storage,
+    customSessionStorage?: Storage,
+  ) {
     this.userProvider = userProvider;
     this.apiKey = apiKey;
     this.storageKey = `EXP_${this.apiKey?.slice(0, 10)}_DEFAULT_USER_PROVIDER`;
+    this.localStorage = customLocalStorage || new LocalStorage();
+    this.sessionStorage = customSessionStorage || new SessionStorage();
   }
 
   getUser(): ExperimentUser {

packages/experiment-tag/src/experiment.ts

@@ -18,7 +18,15 @@ import mutate, { MutationController } from 'dom-mutator';
 import { ConsentAwareExposureHandler } from './exposure/consent-aware-exposure-handler';
 import { MessageBus } from './message-bus';
 import { showPreviewModeModal } from './preview/preview';
-import { ConsentAwareStorage } from './storage/consent-aware-storage';
+import {
+  ConsentAwareLocalStorage,
+  ConsentAwareSessionStorage,
+  ConsentAwareStorage,
+} from './storage/consent-aware-storage';
+import {
+  getAndParseStorageItem,
+  setAndStringifyStorageItem,
+} from './storage/storage';
 import { PageChangeEvent, SubscriptionManager } from './subscriptions';
 import {
   ConsentOptions,
@@ -161,6 +169,10 @@ export class DefaultWebExperimentClient implements WebExperimentClient {
       // eslint-disable-next-line @typescript-eslint/ban-ts-comment
       // @ts-ignore
       internalInstanceNameSuffix: 'web',
+      consentAwareStorage: {
+        localStorage: new ConsentAwareLocalStorage(this.storage),
+        sessionStorage: new ConsentAwareSessionStorage(this.storage),
+      },
       initialFlags: initialFlagsString,
       // timeout for fetching remote flags
       fetchTimeoutMillis: 1000,
@@ -878,9 +890,13 @@ export class DefaultWebExperimentClient implements WebExperimentClient {
         }
       });
 
-      this.storage.setItem('sessionStorage', PREVIEW_MODE_SESSION_KEY, {
-        previewFlags: this.previewFlags,
-      });
+      setAndStringifyStorageItem<PreviewState>(
+        'sessionStorage',
+        PREVIEW_MODE_SESSION_KEY,
+        {
+          previewFlags: this.previewFlags,
+        },
+      );
       const previewParamsToRemove = [
         ...Object.keys(this.previewFlags),
         PREVIEW_MODE_PARAM,
@@ -896,7 +912,7 @@ export class DefaultWebExperimentClient implements WebExperimentClient {
       // if in preview mode, listen for ForceVariant messages
       WindowMessenger.setup();
     } else {
-      const previewState: PreviewState | null = this.storage.getItem(
+      const previewState = getAndParseStorageItem<PreviewState>(
         'sessionStorage',
         PREVIEW_MODE_SESSION_KEY,
       );

packages/experiment-tag/src/storage/consent-aware-storage.ts

@@ -4,10 +4,11 @@ import type { Campaign } from '@amplitude/analytics-core';
 import { ConsentStatus } from '../types';
 
 import {
-  getStorage,
-  getStorageItem,
+  getAndParseStorageItem,
+  getRawStorageItem,
   removeStorageItem,
-  setStorageItem,
+  setAndStringifyStorageItem,
+  setRawStorageItem,
   StorageType,
 } from './storage';
 
@@ -16,6 +17,7 @@ import {
  */
 export class ConsentAwareStorage {
   private inMemoryStorage: Map<StorageType, Map<string, unknown>> = new Map();
+  private inMemoryRawStorage: Map<StorageType, Map<string, string>> = new Map();
   private inMemoryMarketingCookies: Map<string, Campaign> = new Map();
   private consentStatus: ConsentStatus;
 
@@ -32,15 +34,19 @@ export class ConsentAwareStorage {
     if (consentStatus === ConsentStatus.GRANTED) {
       for (const [storageType, storageMap] of this.inMemoryStorage.entries()) {
         for (const [key, value] of storageMap.entries()) {
-          try {
-            const jsonString = JSON.stringify(value);
-            getStorage(storageType)?.setItem(key, jsonString);
-          } catch (error) {
-            console.warn(`Failed to persist data for key ${key}:`, error);
-          }
+          setAndStringifyStorageItem(storageType, key, value);
+        }
+      }
+      for (const [
+        storageType,
+        storageMap,
+      ] of this.inMemoryRawStorage.entries()) {
+        for (const [key, value] of storageMap.entries()) {
+          setRawStorageItem(storageType, key, value);
         }
       }
       this.inMemoryStorage.clear();
+      this.inMemoryRawStorage.clear();
       this.persistMarketingCookies().catch();
     }
   }
@@ -73,7 +79,8 @@ export class ConsentAwareStorage {
    */
   public getItem<T>(storageType: StorageType, key: string): T | null {
     if (this.consentStatus === ConsentStatus.GRANTED) {
-      return getStorageItem(storageType, key);
+      const value = getAndParseStorageItem(storageType, key);
+      return value as T;
     }
 
     const storageMap = this.inMemoryStorage.get(storageType);
@@ -89,7 +96,7 @@ export class ConsentAwareStorage {
    */
   public setItem(storageType: StorageType, key: string, value: unknown): void {
     if (this.consentStatus === ConsentStatus.GRANTED) {
-      setStorageItem(storageType, key, value);
+      setAndStringifyStorageItem(storageType, key, value);
     } else {
       if (!this.inMemoryStorage.has(storageType)) {
         this.inMemoryStorage.set(storageType, new Map());
@@ -115,6 +122,42 @@ export class ConsentAwareStorage {
     }
   }
 
+  /**
+   * Get a raw string value from storage with consent awareness
+   * This is used by Storage interface implementations that expect raw strings
+   */
+  public getRawItem(storageType: StorageType, key: string): string {
+    if (this.consentStatus === ConsentStatus.GRANTED) {
+      return getRawStorageItem(storageType, key);
+    }
+
+    const storageMap = this.inMemoryRawStorage.get(storageType);
+    if (storageMap && storageMap.has(key)) {
+      return storageMap.get(key) || '';
+    }
+
+    return '';
+  }
+
+  /**
+   * Set a raw string value in storage with consent awareness
+   * This is used by Storage interface implementations that work with raw strings
+   */
+  public setRawItem(
+    storageType: StorageType,
+    key: string,
+    value: string,
+  ): void {
+    if (this.consentStatus === ConsentStatus.GRANTED) {
+      setRawStorageItem(storageType, key, value);
+    } else {
+      if (!this.inMemoryRawStorage.has(storageType)) {
+        this.inMemoryRawStorage.set(storageType, new Map());
+      }
+      this.inMemoryRawStorage.get(storageType)?.set(key, value);
+    }
+  }
+
   /**
    * Set marketing cookie with consent awareness
    * Parses current campaign data from URL and referrer, then stores it in the marketing cookie
@@ -138,3 +181,35 @@ export class ConsentAwareStorage {
     }
   }
 }
+
+export class ConsentAwareLocalStorage {
+  constructor(private consentAwareStorage: ConsentAwareStorage) {}
+
+  get(key: string): string {
+    return this.consentAwareStorage.getRawItem('localStorage', key);
+  }
+
+  put(key: string, value: string): void {
+    this.consentAwareStorage.setRawItem('localStorage', key, value);
+  }
+
+  delete(key: string): void {
+    this.consentAwareStorage.removeItem('localStorage', key);
+  }
+}
+
+export class ConsentAwareSessionStorage {
+  constructor(private consentAwareStorage: ConsentAwareStorage) {}
+
+  get(key: string): string {
+    return this.consentAwareStorage.getRawItem('sessionStorage', key);
+  }
+
+  put(key: string, value: string): void {
+    this.consentAwareStorage.setRawItem('sessionStorage', key, value);
+  }
+
+  delete(key: string): void {
+    this.consentAwareStorage.removeItem('sessionStorage', key);
+  }
+}

packages/experiment-tag/src/storage/storage.ts

@@ -3,43 +3,54 @@ import { getGlobalScope } from '@amplitude/experiment-core';
 export type StorageType = 'localStorage' | 'sessionStorage';
 
 /**
- * Get a JSON value from storage and parse it
+ * Get a JSON string value from storage
  * @param storageType - The type of storage to use ('localStorage' or 'sessionStorage')
- * @param key - The key to retrieve
- * @returns The parsed JSON value or null if not found or invalid JSON
+ * @param key - The key to retrieve the value for
+ * @returns The JSON string value, or null if not found
  */
-export const getStorageItem = <T>(
+export const getRawStorageItem = <T>(
+  storageType: StorageType,
+  key: string,
+): string => {
+  return getStorage(storageType)?.getItem(key) || '';
+};
+
+/**
+ * Set a JSON string value in storage
+ * @param storageType - The type of storage to use ('localStorage' or 'sessionStorage')
+ * @param key - The key to set the value for
+ * @param value - The JSON string value to set
+ */
+export const setRawStorageItem = (
+  storageType: StorageType,
+  key: string,
+  value: string,
+): void => {
+  getStorage(storageType)?.setItem(key, value);
+};
+
+export const getAndParseStorageItem = <T>(
   storageType: StorageType,
   key: string,
 ): T | null => {
+  const value = getRawStorageItem(storageType, key);
   try {
-    const value = getStorage(storageType)?.getItem(key);
-    if (!value) {
-      return null;
-    }
-    return JSON.parse(value) as T;
-  } catch (error) {
-    console.warn(`Failed to get and parse JSON from ${storageType}:`, error);
+    return JSON.parse(value);
+  } catch {
     return null;
   }
 };
 
-/**
- * Set a JSON value in storage by stringifying it
- * @param storageType - The type of storage to use ('localStorage' or 'sessionStorage')
- * @param key - The key to store the value under
- * @param value - The value to stringify and store
- */
-export const setStorageItem = (
+export const setAndStringifyStorageItem = <T>(
   storageType: StorageType,
   key: string,
-  value: unknown,
+  value: T,
 ): void => {
   try {
-    const jsonString = JSON.stringify(value);
-    getStorage(storageType)?.setItem(key, jsonString);
+    const stringValue = JSON.stringify(value);
+    setRawStorageItem(storageType, key, stringValue);
   } catch (error) {
-    console.warn(`Failed to stringify and set JSON in ${storageType}:`, error);
+    console.warn(`Failed to persist data for key ${key}:`, error);
   }
 };
 
@@ -59,7 +70,7 @@ export const removeStorageItem = (
   }
 };
 
-export const getStorage = (storageType: StorageType): Storage | null => {
+const getStorage = (storageType: StorageType): Storage | null => {
   const globalScope = getGlobalScope();
   if (!globalScope) {
     return null;

packages/experiment-tag/src/util/messenger.ts

@@ -1,6 +1,6 @@
 import { getGlobalScope } from '@amplitude/experiment-core';
 
-import { getStorageItem } from '../storage/storage';
+import { getAndParseStorageItem } from '../storage/storage';
 
 interface VisualEditorSession {
   injectSrc: string;
@@ -73,10 +73,11 @@ export class WindowMessenger {
    * Retrieve stored session data (read-only)
    */
   private static getStoredSession(): VisualEditorSession | null {
-    const sessionData = getStorageItem<VisualEditorSession>(
+    const sessionData = getAndParseStorageItem<VisualEditorSession>(
       'sessionStorage',
       VISUAL_EDITOR_SESSION_KEY,
-    );
+    ) as VisualEditorSession;
+
     if (!sessionData) {
       return null;
     }