async migrate will result in Data Loss

[KiwiKilian]

Bug Description

Using an async migrate function with persist does not set the result as the store value and resets to the initial state. See this store with async migrate:

import { create, StateCreator } from 'zustand';
import { persist } from 'zustand/middleware';

type StoreState = {
  example: string;
};

const store: StateCreator<StoreState> = () => ({
  example: 'initial',
});

export const useAsyncStore = create<StoreState>()(
  persist(store, {
    name: 'ASYNC_STORE',
    version: 1, // must be raised to 2 after once persisted
    migrate: async (persistedState, version) => {
      const migratedState = persistedState as StoreState;

      if (version < 2) {
        migratedState.example = await new Promise((resolve) => {
          setTimeout(() => {
            resolve('migrated ✅');
          }, 100);
        });
      }

      console.log('async migrate result', JSON.stringify(migratedState));

      return migratedState;
    },
  })
);

And compare to this non-async migrate:

import { create, StateCreator } from 'zustand';
import { persist } from 'zustand/middleware';

type StoreState = {
  example: string;
};

const store: StateCreator<StoreState> = () => ({
  example: 'initial',
});

export const useSyncStore = create<StoreState>()(
  persist(store, {
    name: 'SYNC_STORE',
    version: 1, // must be raised to 2 after once persisted
    migrate: (persistedState, version) => {
      const migratedState = persistedState as StoreState;

      if (version < 2) {
        migratedState.example = 'migrated ✅';
      }

      console.log('sync migrate result', JSON.stringify(migratedState));

      return migratedState;
    },
  })
);

When there is a persisted value in the storage and raising the version to 2 the console will show the correctly migrated value async migrate result {"example":"migrated ✅"} but the store will not accept it and reset to the initial value. This can be produced by following the steps in the linked example:

0. Initially displayed

Result:

Async
initial

Sync
initial

1. Click Button to initially persist stores

Result:

Async
persisted ✅

Sync
persisted ✅

2. Manually change version in both stores to 2

Result:

Async
initial

Sync
migrated ✅

→ Both stores should have changed the value to migrated ✅ but the async will fallback to initial

This is broken from >= 4.5.5 upwards (I've reproduced with 4.5.5, 5.0.0 and 5.0.1), similarly to the behavior change #2763 I suspect #2678, to have broken this functionality.

If I'm not doing something wrong this should be quite prominently announced. People are facing data loss, if they have an async migrate and it's executed using any of the affected versions. It might even be a case for a security advisory.

Reproduction Link

https://stackblitz.com/edit/vitejs-vite-jd3xf5?file=src%2FuseAsyncStore.ts

[dai-shi]

Thanks for reporting. Nice catch.
@double-thinker Can you look into it?

[KiwiKilian]

@double-thinker I've added tests for this case in #2833.