Zustand (^4.1.5) persist middleware with deferred hydration.

deferedPersist.ts

import { persist, PersistOptions } from "zustand/middleware"
import { StateCreator, StoreApi, StoreMutatorIdentifier } from "zustand/vanilla"

type DeferredPersistOptions<S extends object> = PersistOptions<S> & {
  hydrateOnResolve?: Promise<void>
}

const DEFAULT_GET_STORAGE = () => localStorage

export default function deferredPersist<
  S extends object,
  CustomSetState extends [StoreMutatorIdentifier, unknown][] = [
    StoreMutatorIdentifier,
    unknown
  ][],
  CustomGetState extends [StoreMutatorIdentifier, unknown][] = [
    StoreMutatorIdentifier,
    unknown
  ][]
>(
  config: StateCreator<
    S,
    [...CustomSetState, ["zustand/persist", unknown]],
    CustomGetState,
    S
  >,
  deferredOptions: DeferredPersistOptions<S>
) {
  const hydrateOnResolve = deferredOptions.hydrateOnResolve || Promise.resolve()
  const getStorage = deferredOptions.getStorage || DEFAULT_GET_STORAGE

  const options: PersistOptions<S> = {
    ...deferredOptions,
    getStorage: () => ({
      setItem: async (...args) =>
        hydrateOnResolve.then(() => getStorage().setItem(...args)),
      getItem: async (...args) =>
        hydrateOnResolve.then(() => getStorage().getItem(...args)),
      removeItem: async (...args) =>
        hydrateOnResolve.then(() => getStorage().removeItem?.(...args)),
    }),
  }

  return persist<S, CustomSetState, CustomGetState>(config, options)
}

hey @leosh1d, thank you so much! It works perfectly 🙇