Jotai's createJSONStorage code with few tweaks to support custom subscribe functionality for persist storages

jotai-createJSONStorage.ts

import { AsyncStorage, SyncStorage } from 'jotai/vanilla/utils/atomWithStorage';

type Unsubscribe = () => void;

type Subscribe<Value> = (
  key: string,
  callback: (value: Value) => void,
  initialValue: Value
) => Unsubscribe;

type SubscribeHandler<Value> = (
  subscribe: Subscribe<Value>,
  key: string,
  callback: (value: Value) => void,
  initialValue: Value
) => Unsubscribe;

export interface AsyncStringStorage<V = string> {
  getItem: (key: string) => PromiseLike<string | null>;
  setItem: (key: string, newValue: string) => PromiseLike<void>;
  removeItem: (key: string) => PromiseLike<void>;
  subscribe?: Subscribe<V>;
}

export interface SyncStringStorage<V = string> {
  getItem: (key: string) => string | null;
  setItem: (key: string, newValue: string) => void;
  removeItem: (key: string) => void;
  subscribe?: Subscribe<V>;
}

type JsonStorageOptions<V> = {
  reviver?: (key: string, value: unknown) => unknown;
  replacer?: (key: string, value: unknown) => unknown;
  subscribe?: Subscribe<V>;
};

const isPromiseLike = (x: unknown): x is PromiseLike<unknown> =>
  typeof (x as any)?.then === 'function';

export function createJSONStorage__custom<Value>(
  getStringStorage: () => AsyncStringStorage,
  options?: JsonStorageOptions<Value>
): AsyncStorage<Value>;

export function createJSONStorage__custom<Value>(
  getStringStorage: () => SyncStringStorage,
  options?: JsonStorageOptions<Value>
): SyncStorage<Value>;

export function createJSONStorage__custom<Value>(
  getStringStorage: () =>
    | AsyncStringStorage<Value>
    | SyncStringStorage<Value>
    | undefined = () => {
    try {
      return window.localStorage;
    } catch (e) {
      if (process.env.NODE_ENV !== 'production') {
        if (typeof window !== 'undefined') {
          console.warn(e);
        }
      }
      return undefined;
    }
  },
  options?: JsonStorageOptions<Value>
): AsyncStorage<Value> | SyncStorage<Value> {
  let lastStr: string | undefined;
  let lastValue: Value;
  const webStorageSubscribe = ((key, callback) => {
    const storageEventCallback = (e: StorageEvent) => {
      if (e.storageArea === getStringStorage() && e.key === key) {
        callback((e.newValue || '') as Value);
      }
    };
    window.addEventListener('storage', storageEventCallback);
    return () => {
      window.removeEventListener('storage', storageEventCallback);
    };
  }) satisfies Subscribe<Value>;

  const handleSubscribe: SubscribeHandler<Value> = (
    subscriber,
    key,
    callback,
    initialValue
  ) => {
    function cb(v: Value) {
      let newValue: Value;
      try {
        newValue = JSON.parse((v as string) || '');
      } catch {
        newValue = initialValue;
      }

      callback(newValue);
    }

    return subscriber(key, cb, initialValue);
  };

  const storage: AsyncStorage<Value> | SyncStorage<Value> = {
    getItem: (key, initialValue) => {
      const parse = (str: string | null) => {
        str = str || '';
        if (lastStr !== str) {
          try {
            lastValue = JSON.parse(str, options?.reviver);
          } catch {
            return initialValue;
          }
          lastStr = str;
        }
        return lastValue;
      };
      const str = getStringStorage()?.getItem(key) ?? null;
      if (isPromiseLike(str)) {
        return str.then(parse) as never;
      }
      return parse(str) as never;
    },
    setItem: (key, newValue) =>
      getStringStorage()?.setItem(
        key,
        JSON.stringify(newValue, options?.replacer)
      ),
    removeItem: (key) => getStringStorage()?.removeItem(key),
  };

  if (
    typeof window !== 'undefined' &&
    typeof window.addEventListener === 'function'
  ) {
    storage.subscribe = handleSubscribe.bind(
      null,
      getStringStorage()?.subscribe ?? webStorageSubscribe
    );
  }
  return storage;
}

Usage

here's the code to use the custom createJSONStorage function above to implement persist-with-cookies for Jotai atoms:

import {
  createJSONStorage__custom,
  type SyncStringStorage,
} from '@/utils/jotai-createJSONStorage';
import { atom } from 'jotai';
import { atomWithStorage } from 'jotai/utils';
import Cookies, { type CookieAttributes } from 'js-cookie';


const getCookieAtomStorage =
  <T extends string = string>(attributes: CookieAttributes) =>
  () =>
    ({
      getItem: (key) => Cookies.get(key),
      setItem: (key, value) => Cookies.set(key, String(value), attributes),
      removeItem: (key) => Cookies.remove(key, attributes),
      subscribe:
        'cookieStore' in window
          ? (key, callback, initialValue) => {
              const storageEventCallback = (e: any) => {
                const { changed, deleted, type } = e;

                if (type === 'change') {
                  const { value: newTokenCookieValue } =
                    changed.find((i: any) => i.name === key) ?? {};

                  if (newTokenCookieValue) {
                    //@ts-ignore
                    callback(Cookies.converter.read(newTokenCookieValue));
                  }
                }

                if (type === 'delete') {
                  const isTokenCookieDeleted = deleted.find(
                    (i: any) => i.name === key
                  );
                  if (isTokenCookieDeleted) callback(initialValue);
                }
              };

              //@ts-ignore
              window.cookieStore.addEventListener(
                'change',
                storageEventCallback
              );
              return () => {
                //@ts-ignore
                window.cookieStore.removeEventListener(
                  'change',
                  storageEventCallback
                );
              };
            }
          : undefined,
    } as SyncStringStorage<T>);


export const tokenAtom = atomWithStorage<string | null>(
  'token',
  null,
  createJSONStorage__custom(
    getCookieAtomStorage({
      expires: 30,
      domain: process.env.NEXT_PUBLIC_DOMAIN,
    })
  )
);

Note that this code implements a poor way of handling browser feature detection and compatibility.