cogsbox-state 0.5.476 → 0.5.479

package/src/CogsState.tsx

@@ -1,6 +1,10 @@
 'use client';
 import { pluginStore } from './pluginStore';
-import { FormWrapperParams, type CogsPlugin } from './plugins';
+import {
+  FormWrapperParams,
+  type ChainMethodCallables,
+  type CogsPlugin,
+} from './plugins';
 import {
   createElement,
   startTransition,
@@ -79,7 +83,8 @@ type CutFunctionType<T> = (
   options?: { waitForSync?: boolean }
 ) => StateObject<T>;
 
-export type InferArrayElement<T> = T extends (infer U)[] ? U : never;
+export type InferArrayElement<T> =
+  NonNullable<T> extends (infer U)[] ? U : never;
 
 export type FormControl<T> = (obj: FormElementParams<T>) => JSX.Element;
 
@@ -95,7 +100,15 @@ export type InsertTypeObj<T> = (payload: InsertParams<T>) => void;
 type EffectFunction<T, R> = (state: T, deps: any[]) => R;
 export type PerPathFormOptsType<
   TState,
-  TPlugins extends readonly CogsPlugin<any, any, any, any, any>[] = [],
+  TPlugins extends readonly CogsPlugin<
+    any,
+    any,
+    any,
+    any,
+    any,
+    any,
+    any
+  >[] = [],
 > = Omit<FormOptsType, 'formElements'> & {
   formElements?: FormsElementsType<TState, TPlugins>;
 };
@@ -123,7 +136,15 @@ export type ArrayElementExtras<TParentArray = unknown> = {
 // EndType - NO $cutThis here, it's purely contextual
 export type EndType<
   T,
-  TPlugins extends readonly CogsPlugin<any, any, any, any, any>[] = [],
+  TPlugins extends readonly CogsPlugin<
+    any,
+    any,
+    any,
+    any,
+    any,
+    any,
+    any
+  >[] = [],
 > = {
   $getPluginMetaData: (pluginName: string) => Record<string, any>;
   $addPluginMetaData: (key: string, data: Record<string, any>) => void;
@@ -175,6 +196,7 @@ export type EndType<
     selectText: () => void;
   };
   $removeStorage: () => void;
+  $setRaw: (value: T) => void;
   $sync: () => void;
   $validationWrapper: ({
     children,
@@ -190,12 +212,20 @@ export type EndType<
 export type ArrayElementState<
   TElement,
   TParentArray,
-  TPlugins extends readonly CogsPlugin<any, any, any, any, any>[] = [],
+  TPlugins extends readonly CogsPlugin<
+    any,
+    any,
+    any,
+    any,
+    any,
+    any,
+    any
+  >[] = [],
 > = StateObject<TElement, TPlugins> & ArrayElementExtras<TParentArray>;
 
 export type ArrayEndType<
   TShape extends unknown,
-  TPlugins extends readonly CogsPlugin<any, any, any, any, any>[],
+  TPlugins extends readonly CogsPlugin<any, any, any, any, any, any, any>[],
 > = {
   (): TShape;
   (newValue: TShape | ((prev: TShape) => TShape)): void;
@@ -299,24 +329,33 @@ export type ArrayEndType<
 
 export type StateObject<
   T,
-  TPlugins extends readonly CogsPlugin<any, any, any, any, any>[] = [],
+  TPlugins extends readonly CogsPlugin<
+    any,
+    any,
+    any,
+    any,
+    any,
+    any,
+    any
+  >[] = [],
 > = {
   (): T;
   (newValue: T | ((prev: T) => T)): void;
-} & (T extends any[]
+} & ([NonNullable<T>] extends [any[]] // <-- Wrap in brackets
   ? ArrayEndType<T, TPlugins>
-  : T extends Record<string, unknown> | object
-    ? { [K in keyof T]-?: StateObject<T[K], TPlugins> }
-    : EndType<T, TPlugins>) & // primitives just get EndType
+  : [NonNullable<T>] extends [Record<string, unknown> | object] // <-- Wrap in brackets
+    ? {
+        [K in keyof NonNullable<T>]-?: StateObject<NonNullable<T>[K], TPlugins>;
+      }
+    : {}) & // Fallback to {} since we intersect EndType below anyway
   EndType<T, TPlugins> & {
-    // NO third param = no $cutThis
-    $toggle: T extends boolean ? () => void : never;
+    $toggle: NonNullable<T> extends boolean ? () => void : never;
     $validate: () => { success: boolean; data?: T; error?: any };
     $_componentId: string | null;
     $getComponents: () => ComponentsType;
     $_initialState: T;
     $updateInitialState: (newState: T | null) => {
-      fetchId: (field: keyof T) => string | number;
+      fetchId: (field: keyof NonNullable<T>) => string | number;
     };
     $initializeAndMergeShadowState: (newState: any | null) => void;
     $_isLoading: boolean;
@@ -332,10 +371,23 @@ export type StateObject<
       }) => void
     ) => void;
     $getLocalStorage: (key: string) => LocalStorageData<T> | null;
-  };
-
+  } & PluginChainMethodCallables<TPlugins>;
 export type CogsUpdate<T extends unknown> = UpdateType<T>;
 
+type UnionToIntersection<T> = (
+  T extends any ? (arg: T) => void : never
+) extends (arg: infer I) => void
+  ? I
+  : never;
+
+type PluginChainMethodCallables<
+  TPlugins extends readonly CogsPlugin<any, any, any, any, any, any, any>[],
+> = UnionToIntersection<
+  TPlugins[number] extends CogsPlugin<any, any, any, any, any, infer TMethods>
+    ? ChainMethodCallables<TMethods>
+    : {}
+>;
+
 type EffectiveSetStateArg<
   T,
   UpdateType extends 'update' | 'insert' | 'cut',
@@ -407,7 +459,15 @@ type SyncOptionsType<TApiParams> = {
 
 export type CreateStateOptionsType<
   T extends unknown = unknown,
-  TPlugins extends readonly CogsPlugin<any, any, any, any, any>[] = [],
+  TPlugins extends readonly CogsPlugin<
+    any,
+    any,
+    any,
+    any,
+    any,
+    any,
+    any
+  >[] = [],
 > = {
   formElements?: FormsElementsType<T, TPlugins>;
   validation?: ValidationOptionsType;
@@ -470,7 +530,15 @@ type ScopedPluginApi<THookReturn, TFieldMetaData> = {
 };
 export type FormsElementsType<
   TState,
-  TPlugins extends readonly CogsPlugin<any, any, any, any, any>[] = [],
+  TPlugins extends readonly CogsPlugin<
+    any,
+    any,
+    any,
+    any,
+    any,
+    any,
+    any
+  >[] = [],
 > = {
   // These optional, built-in wrappers are unchanged.
   validation?: (options: {
@@ -504,16 +572,6 @@ export type FormsElementsType<
     key?: string;
   }) => React.ReactNode;
 };
-export type CogsInitialState<T> =
-  | {
-      initialState: T;
-    }
-  | CreateStateOptionsType<T>;
-
-export type TransformedStateType<T> = {
-  [P in keyof T]: T[P] extends CogsInitialState<infer U> ? U : T[P];
-};
-
 const {
   getInitialOptions,
   updateInitialStateGlobal,
@@ -692,18 +750,8 @@ function setOptions<StateKey, Opt>({
 
   return mergedOptions;
 }
-export function addStateOptions<T>(
-  initialState: T,
-  options: CreateStateOptionsType<T>
-) {
-  return {
-    ...options,
-    initialState,
-    _addStateOptions: true,
-  };
-}
 export type PluginData = {
-  plugin: CogsPlugin<any, any, any, any, any>;
+  plugin: CogsPlugin<any, any, any, any, any, any, any>;
   options: any;
   hookData?: any;
 };
@@ -712,7 +760,15 @@ export type PluginData = {
 
 export const createCogsState = <
   State extends Record<string, unknown>,
-  const TPlugins extends readonly CogsPlugin<string, any, any, any, any>[] = [],
+  const TPlugins extends readonly CogsPlugin<
+    string,
+    any,
+    any,
+    any,
+    any,
+    any,
+    any
+  >[] = [],
 >(
   initialState: State,
   opt?: {
@@ -721,60 +777,92 @@ export const createCogsState = <
     validation?: ValidationOptionsType;
   }
 ) => {
-  type PluginOptions = {
-    [K in TPlugins[number] as K['name']]?: K extends CogsPlugin<
-      string,
-      infer O,
-      any,
-      any,
-      any
-    >
+  type ExtractPluginOptions<T> = T extends {
+    useHook?: (params: { options: infer O }) => any; // infers O here
+  }
+    ? O
+    : T extends CogsPlugin<string, infer O, any, any, any, any, any> // AND here
       ? O
       : never;
+
+  type PluginOptions = {
+    [K in TPlugins[number] as K['name']]?: ExtractPluginOptions<K>;
   };
 
   if (opt?.plugins) {
     pluginStore.getState().setRegisteredPlugins(opt.plugins as any);
   }
+  type UnionToIntersection<U> = (
+    U extends any ? (k: U) => void : never
+  ) extends (k: infer I) => void
+    ? I
+    : never;
+
+  // Extract plugin initial state from the CogsPlugin generic (7th param) first,
+  // falling back to structural inference from initialState() return type.
+  type ExtractPluginState<T> = T extends CogsPlugin<
+    string,
+    any,
+    any,
+    any,
+    any,
+    any,
+    infer S
+  >
+    ? S extends Record<string, unknown>
+      ? S
+      : {}
+    : T extends {
+          initialState?: () => infer S;
+        }
+      ? S extends Record<string, unknown>
+        ? S
+        : {}
+      : {};
+
+  // Extract the merged plugin states into a helper type
+  type PluginStates = UnionToIntersection<ExtractPluginState<TPlugins[number]>>;
+
+  // User keys win over plugin keys — omit overlapping plugin keys before merging.
+  type FullState = Prettify<
+    State & {
+      [K in keyof PluginStates as K extends keyof State
+        ? never
+        : K]: PluginStates[K];
+    }
+  >;
+  const pluginState: Record<string, unknown> = {};
+  if (opt?.plugins) {
+    for (const plugin of opt.plugins) {
+      if (typeof plugin.initialState === 'function') {
+        Object.assign(pluginState, plugin.initialState());
+      }
+    }
+  }
+
+  const mergedInitialState = { ...pluginState, ...initialState } as FullState;
 
   const [statePart, initialOptionsPart] =
-    transformStateFunc<State>(initialState);
+    transformStateFunc<FullState>(mergedInitialState);
 
-  // FIX: Store options INCLUDING validation for each state key
   Object.keys(statePart).forEach((key) => {
-    let existingOptions = initialOptionsPart[key] || {};
-
-    const mergedOptions: any = {
-      ...existingOptions,
-    };
+    let mergedOptions: any = {};
 
     if (opt?.formElements) {
-      mergedOptions.formElements = {
-        ...opt.formElements,
-        ...(existingOptions.formElements || {}),
-      };
+      mergedOptions.formElements = opt.formElements;
     }
 
     mergedOptions.validation = {
       onBlur: 'error',
       ...opt?.validation,
-      ...(existingOptions.validation || {}),
     };
 
-    if (opt?.validation?.key && !existingOptions.validation?.key) {
-      mergedOptions.validation.key = `${opt.validation.key}.${key}`;
-    }
-
     const existingGlobalOptions = getInitialOptions(key);
 
     const finalOptions = existingGlobalOptions
       ? {
           ...existingGlobalOptions,
-          ...mergedOptions,
-          formElements: {
-            ...existingGlobalOptions.formElements,
-            ...mergedOptions.formElements,
-          },
+          formElements: opt?.formElements,
           validation: {
             ...existingGlobalOptions.validation,
             ...mergedOptions.validation,
@@ -782,35 +870,56 @@ export const createCogsState = <
         }
       : mergedOptions;
 
-    setInitialStateOptions(key, finalOptions);
+    if (Object.keys(finalOptions).length > 0) {
+      setInitialStateOptions(key, finalOptions);
+    }
   });
 
   Object.keys(statePart).forEach((key) => {
     initializeShadowState(key, statePart[key]);
   });
-  type Prettify<T> = {
-    [K in keyof T]: T[K];
-  } & {};
-  type StateKeys = keyof typeof statePart;
+  type StateKeys = keyof FullState;
+
+  // Flattens the 4x spam into a single clean object
+  type CleanIntersection<T> = T extends object ? { [K in keyof T]: T[K] } : T;
+
+  // Helper to find which keys are "keyed" (like queryParams)
+  type KeyedKeys<P> = {
+    [K in keyof P]-?: NonNullable<P[K]> extends { __key: 'keyed'; map: any }
+      ? K
+      : never;
+  }[keyof P];
+
+  type StateSlice<StateKey extends StateKeys> = FullState[StateKey];
 
   const useCogsState = <StateKey extends StateKeys>(
     stateKey: StateKey,
     options?: Prettify<
-      OptionsType<(typeof statePart)[StateKey], never> & {
+      OptionsType<StateSlice<StateKey>, never> & {
         [PName in keyof PluginOptions]?: PluginOptions[PName] extends infer P
           ? P extends Record<string, any>
-            ? {
-                [K in keyof P]: P[K] extends { __key: 'keyed'; map: infer TMap }
-                  ? StateKey extends keyof TMap
-                    ? TMap[StateKey]
-                    : never
-                  : P[K];
-              }
+            ? Prettify<
+                Partial<Pick<P, Exclude<keyof P, KeyedKeys<P>>>> & {
+                  [K in KeyedKeys<P> as StateKey extends keyof NonNullable<
+                    P[K]
+                  >['map']
+                    ? NonNullable<P[K]>['map'][StateKey] extends undefined
+                      ? never
+                      : keyof NonNullable<P[K]>['map'][StateKey] extends never
+                        ? never
+                        : K
+                    : never]: CleanIntersection<
+                    StateKey extends keyof NonNullable<P[K]>['map']
+                      ? NonNullable<P[K]>['map'][StateKey]
+                      : never
+                  >;
+                }
+              >
             : P
           : never;
       }
     >
-  ) => {
+  ): StateObject<StateSlice<StateKey>> => {
     const [componentId] = useState(options?.componentId ?? uuidv4());
 
     const currentOptions = setOptions({
@@ -825,21 +934,18 @@ export const createCogsState = <
     const thiState =
       getShadowValue(stateKey as string, []) || statePart[stateKey as string];
 
-    const updater = useCogsStateFn<(typeof statePart)[StateKey], TPlugins>(
-      thiState,
-      {
-        stateKey: stateKey as string,
-        syncUpdate: options?.syncUpdate,
-        componentId,
-        localStorage: options?.localStorage,
-        middleware: options?.middleware,
-        reactiveType: options?.reactiveType,
-        reactiveDeps: options?.reactiveDeps,
-        defaultState: options?.defaultState as any,
-        dependencies: options?.dependencies,
-        serverState: options?.serverState,
-      }
-    );
+    const updater = useCogsStateFn<StateSlice<StateKey>>(thiState, {
+      stateKey: stateKey as string,
+      syncUpdate: options?.syncUpdate,
+      componentId,
+      localStorage: options?.localStorage,
+      middleware: options?.middleware,
+      reactiveType: options?.reactiveType,
+      reactiveDeps: options?.reactiveDeps,
+      defaultState: options?.defaultState as any,
+      dependencies: options?.dependencies,
+      serverState: options?.serverState,
+    });
 
     useEffect(() => {
       if (options) {
@@ -858,14 +964,14 @@ export const createCogsState = <
       };
     }, [stateKey, updater]);
 
-    return updater;
+    return updater as StateObject<StateSlice<StateKey>>;
   };
 
   function setCogsOptionsByKey<StateKey extends StateKeys>(
     stateKey: StateKey,
-    options: CreateStateOptionsType<(typeof statePart)[StateKey], TPlugins> &
+    options: CreateStateOptionsType<StateSlice<StateKey>, TPlugins> &
       Omit<
-        OptionsType<(typeof statePart)[StateKey]>,
+        OptionsType<StateSlice<StateKey>>,
         keyof CreateStateOptionsType
       >
   ) {
@@ -984,6 +1090,18 @@ const loadFromLocalStorage = (localStorageKey: string) => {
     return null;
   }
 };
+
+const removeFromLocalStorage = (localStorageKey?: string) => {
+  if (!localStorageKey) return;
+
+  try {
+    if (typeof window !== 'undefined' && window.localStorage) {
+      window.localStorage.removeItem(localStorageKey);
+    }
+  } catch (error) {
+    console.error('Error removing from localStorage:', error);
+  }
+};
 const loadAndApplyLocalStorage = (stateKey: string, options: any) => {
   const currentState = getShadowValue(stateKey, []);
   const { sessionId } = useCogsConfig();
@@ -1210,7 +1328,7 @@ function getComponentNotifications(
       const pathMeta = getShadowMetadata(stateKey, currentPath);
 
       if (pathMeta?.pathComponents) {
-        pathMeta.pathComponents.forEach((componentId: string) => {
+        pathMeta.pathComponents.forEach((componentId) => {
           const component = rootMeta.components?.get(componentId);
           if (component) {
             componentsToNotify.add(component);
@@ -1499,7 +1617,15 @@ function createEffectiveSetState<T>(
 
 export function useCogsStateFn<
   TStateObject extends unknown,
-  const TPlugins extends readonly CogsPlugin<any, any, any, any, any>[],
+  const TPlugins extends readonly CogsPlugin<
+    any,
+    any,
+    any,
+    any,
+    any,
+    any,
+    any
+  >[] = [],
 >(
   stateObject: TStateObject,
   {
@@ -1959,9 +2085,72 @@ function getScopedData(stateKey: string, path: string[], meta?: MetaData) {
   };
 }
 
+function pathMatchesPattern(path: string[], pattern?: string[]) {
+  if (!pattern) return true;
+  if (path.length !== pattern.length) return false;
+
+  return pattern.every((segment, index) => {
+    return segment === '*' || segment === path[index];
+  });
+}
+
+function valueMatchesChainTarget(value: any, target: string) {
+  if (target === 'any') return true;
+  if (target === 'array') return Array.isArray(value);
+  if (target === 'boolean') return typeof value === 'boolean';
+  if (target === 'object') {
+    return value !== null && typeof value === 'object' && !Array.isArray(value);
+  }
+  if (target === 'primitive') {
+    return (
+      value === null || (typeof value !== 'object' && !Array.isArray(value))
+    );
+  }
+  return false;
+}
+
+function getFieldRefsForPath(stateKey: string, path: string[]) {
+  const meta = getGlobalStore.getState().getShadowMetadata(stateKey, path);
+  if (!meta?.clientActivityState?.elements) return [];
+  const refs: RefObject<any>[] = [];
+  meta.clientActivityState.elements.forEach((entry: any) => {
+    if (entry.domRef?.current) refs.push(entry.domRef);
+  });
+  return refs;
+}
+
+function getFieldElementsForPath(stateKey: string, path: string[]) {
+  const refs = getFieldRefsForPath(stateKey, path);
+  return refs.map((ref) => ref.current).filter(Boolean);
+}
+
+function setFieldDisabledForPath(
+  stateKey: string,
+  path: string[],
+  disabled: boolean
+) {
+  getFieldElementsForPath(stateKey, path).forEach((el: any) => {
+    if ('disabled' in el) {
+      el.disabled = disabled;
+      return;
+    }
+
+    el.style.pointerEvents = disabled ? 'none' : '';
+    el.setAttribute('aria-disabled', String(disabled));
+  });
+}
+
 function createProxyHandler<
   T,
-  const TPlugins extends readonly CogsPlugin<any, any, any, any, any>[],
+  const TPlugins extends readonly CogsPlugin<
+    any,
+    any,
+    any,
+    any,
+    any,
+    any,
+    any
+  >[],
 >(
   stateKey: string,
   effectiveSetState: EffectiveSetState<T>,
@@ -1979,6 +2168,11 @@ function createProxyHandler<
     componentId: string;
     meta?: MetaData;
   }): any {
+    const rawMeta = getShadowMetadata(stateKey, path);
+    if (rawMeta?.isRaw) {
+      return getShadowValue(stateKey, path);
+    }
+
     const derivationSignature = meta
       ? JSON.stringify(meta.arrayViews || meta.transforms)
       : '';
@@ -2055,6 +2249,91 @@ function createProxyHandler<
         }
 
         if (typeof prop === 'string' && !prop.startsWith('$')) {
+          const { value } = getScopedData(stateKey, path, meta);
+          const hasStateChild =
+            value !== null &&
+            typeof value === 'object' &&
+            !Array.isArray(value) &&
+            Object.prototype.hasOwnProperty.call(value, prop);
+
+          if (hasStateChild) {
+            const nextPath = [...path, prop];
+            return rebuildStateShape({
+              path: nextPath,
+              componentId: componentId!,
+              meta,
+            });
+          }
+
+          const registeredPlugins = pluginStore.getState().registeredPlugins;
+          for (const plugin of registeredPlugins) {
+            const chainMethod = (plugin.chainMethods as any)?.[prop];
+            if (!chainMethod) continue;
+            if (!pathMatchesPattern(path, chainMethod.pathPattern)) continue;
+            if (!valueMatchesChainTarget(value, chainMethod.target)) continue;
+
+            return (...args: any[]) => {
+              const store = pluginStore.getState();
+              const options = store.pluginOptions
+                .get(stateKey)
+                ?.get(plugin.name);
+              const hookData = store.getHookResult(stateKey, plugin.name);
+
+              return chainMethod.handler(
+                {
+                  stateKey,
+                  path,
+                  pluginName: plugin.name,
+                  options,
+                  hookData,
+                  $get: () => getScopedData(stateKey, path, meta).value,
+                  $update: (payload: any) => {
+                    effectiveSetState(payload, path, {
+                      updateType: 'update',
+                    });
+
+                    return {
+                      synced: () => {
+                        const shadowMeta = getGlobalStore
+                          .getState()
+                          .getShadowMetadata(stateKey, path);
+
+                        setShadowMetadata(stateKey, path, {
+                          ...shadowMeta,
+                          isDirty: false,
+                          stateSource: 'server',
+                          lastServerSync: Date.now(),
+                        });
+                      },
+                    };
+                  },
+                  $applyOperation: (
+                    operation: UpdateTypeDetail,
+                    metaData?: Record<string, any>
+                  ) => {
+                    effectiveSetState(operation.newValue, operation.path, {
+                      updateType: operation.updateType,
+                      itemId: operation.itemId,
+                      metaData,
+                    });
+                  },
+                  getFieldMetaData: () =>
+                    getPluginMetaDataMap(stateKey, path)?.get(plugin.name),
+                  setFieldMetaData: (data: Record<string, any>) =>
+                    setPluginMetaData(stateKey, path, plugin.name, data),
+                  removeFieldMetaData: () =>
+                    removePluginMetaData(stateKey, path, plugin.name),
+                  getFieldRefs: () => getFieldRefsForPath(stateKey, path),
+                  getFieldElements: () =>
+                    getFieldElementsForPath(stateKey, path),
+                  setFieldDisabled: (disabled: boolean) =>
+                    setFieldDisabledForPath(stateKey, path, disabled),
+                },
+                ...args
+              );
+            };
+          }
+
           const nextPath = [...path, prop];
           return rebuildStateShape({
             path: nextPath,
@@ -2171,8 +2450,18 @@ function createProxyHandler<
             const localKey = isFunction(initalOptionsGet?.localStorage?.key)
               ? initalOptionsGet.localStorage.key(initialState)
               : initalOptionsGet?.localStorage?.key;
-            const storageKey = `${sessionId}-${stateKey}-${localKey}`;
-            if (storageKey) localStorage.removeItem(storageKey);
+            const storageKey =
+              sessionId && localKey
+                ? `${sessionId}-${stateKey}-${localKey}`
+                : undefined;
+            removeFromLocalStorage(storageKey);
+          };
+        }
+        if (prop === '$setRaw') {
+          return (value: any) => {
+            const currentMeta = getShadowMetadata(stateKey, path) || {};
+            setShadowMetadata(stateKey, path, { ...currentMeta, isRaw: true });
+            effectiveSetState(value, path, { updateType: 'update' });
           };
         }
 
@@ -3323,10 +3612,11 @@ function createProxyHandler<
       const localKey = isFunction(initalOptionsGet?.localStorage?.key)
         ? initalOptionsGet?.localStorage?.key(revertState)
         : initalOptionsGet?.localStorage?.key;
-      const storageKey = `${sessionId}-${stateKey}-${localKey}`;
-      if (storageKey) {
-        localStorage.removeItem(storageKey);
-      }
+      const storageKey =
+        sessionId && localKey
+          ? `${sessionId}-${stateKey}-${localKey}`
+          : undefined;
+      removeFromLocalStorage(storageKey);
 
       notifyComponents(stateKey);
 
@@ -3350,11 +3640,11 @@ function createProxyHandler<
         ? initalOptionsGet?.localStorage?.key(initialState)
         : initalOptionsGet?.localStorage?.key;
 
-      const storageKey = `${sessionId}-${stateKey}-${localKey}`;
-
-      if (localStorage.getItem(storageKey)) {
-        localStorage.removeItem(storageKey);
-      }
+      const storageKey =
+        sessionId && localKey
+          ? `${sessionId}-${stateKey}-${localKey}`
+          : undefined;
+      removeFromLocalStorage(storageKey);
       startTransition(() => {
         updateInitialStateGlobal(stateKey, newState);
         initializeShadowState(stateKey, newState);