Zustand (^3.7.0) persist middleware with deferred hydration.

deferredPersist.ts

import { persist, PersistOptions } from 'zustand/middleware';
import { GetState, SetState, StoreApi } from 'zustand/vanilla';

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

const DEFAULT_GET_STORAGE = () => localStorage;

export default function deferredPersist<
    S extends object,
    CustomSetState extends SetState<S> = SetState<S>,
    CustomGetState extends GetState<S> = GetState<S>,
    CustomStoreApi extends StoreApi<S> = StoreApi<S>,
>(
    config: (set: CustomSetState, get: CustomGetState, api: CustomStoreApi) => 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, CustomStoreApi>(config, options);
}

Hey! I'm trying to test your implementation in my project, but it throws a lot of type errors.
I'm using Zustand 4.0.0 with TypeScript 4.7.4 and 4.8.0

First of all, it says that GetState and SetState are deprecated and it can be beautifully fixed with StoreApi<S>['getState'] and StoreApi<S>['setState']. Maybe you should consider changing them too.

And then only 2 errors will remain, see the screenshots below.

For CustomSetState error looks like that:

And for options it won't even fit in the window:

As a result, it also makes errors while creating the store and while using it.

I tried different ways, but couldn't fix it, I'm not a pro in TS either. Can you help me with that? Really want to use persist in my NextJS project

@PAXANDDOS This was implemented to work with zustand ^3.7.0. Maybe you should consider downgrading from 4.0.0 to this version ?

Impementation for ^4.1.5