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)
}

Based on kinoadr gist

hey @leosh1d, thanks for sharing!
qq: how are you passing the hydration event as a promise to hydrateOnResolve?

@tomtobac

in store.ts create promise

let resolveHydratePromise

const HydratePromise = new Promise<void>(function (resolve) {
  resolveHydratePromise = resolve
})

pass HydratePromise to store

{
  hydrateOnResolve: HydratePromise,
  name: "yourStorageName",
}

export promise

  export const ResolveIsHydrated = resolveHydratePromise || function () {}

In _app.tsx resolve promise with useEffect

  useEffect(() => {
    ResolveIsHydrated()
  }, [])

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