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?

[HerrZatacke]

Hi folks!
Do you have a timeframe/estimate on when async migrations will be possible?
The only way I could get this to work for now is to "manually" manipulate localstorage before calling create which feels a bit weird...

await migrateLegacy();

const useItemsStore = create<ItemsState>()(
  persist(
    (set, get) => ({
// ... and so on...

[KiwiKilian]

I'm also a bit worried, as this will result in data loss for user, who are using async migrate and upgraded to the affected versions.

[dai-shi]

Could anyone dig into it if @double-thinker is busy?

[dbritto-dev]

@dai-shi @KiwiKilian here you go -> #2877

[KiwiKilian]

Works in my reproduction. It briefly flashes the initial state during HMR/Migration, is this intended behavior during the migration/rehydration I guess?

[dai-shi]

Yeah, the store itself is created sync, so that behavior is somewhat expected. Might not be ideal though.

[dbritto-dev]

Yeah, that's expected