The Zustand middleware's set cannot safely be nested inside batch

[jamesbvaughan]

Describe the bug

The zustand store is modified twice when mixing the store's set function with Liveblocks' batch function.

To Reproduce

This code results in a bug:

type TestStore = {
  testArray: string[];
  setTestArray: (array: string[]) => void;
  getTestArray: () => string[];
};

const useTestStore = create<WithLiveblocks<TestStore>>(
  liveblocks(
    (set, get) => ({
      testArray: [],
      setTestArray: (array: string[]) => {
        set({ testArray: array });
      },
      getTestArray: () => get().testArray,
    }),
    {
      client: liveblocksClient,
      storageMapping: {
        testArray: true,
      },
    },
  ),
);

const { setTestArray, getTestArray, liveblocks } = useTestStore.getState();
liveblocks.enterRoom("test-room-id");

liveblocks.room.batch(() => {
  // This line modifies the zustand store, then modifies the current batch of Liveblocks changes.
  setTestArray(["element-test-id-2"]);

  // When this callback finishes, `batch` executes all the batched changes in the Liveblocks storage. Because the Zustand `set` function has finished, `isPatching` is no longer `true` and the Liveblocks zustand middleware will apply these changes to the Zustand store, duplicating them.
  // This isn't a problem for things like `LiveObject.set`, but it is a problem for `LiveList.push`.
});

console.log(
  "testArray after - this will have two elements when it should have just one",
  getTestArray(),
);

Expected behavior

I expected that the zustand middleware and the batch function would work nicely together and that the batched changes would not be duplicated in the zustand store.

Additional context

This is admittedly a pretty unusual edge case. We'd typically not hit this issue because we'd be using a single zustand set call for all the grouped changes, and the zustand middleware already uses batch effectively internally, but there are some cases where we have reason to split the set calls.

Some of our current workarounds are:

• Find ways to use a single set call. (Not always ideal or possible with the structure of our code.)
• Use pause and resume to at least get part of the benefit of batching changes. (Only partially helpful.)
• Directly manipulate the Live data types. (Requires more traversal of the Live-ified version of our document than we'd like and flips the order of operations so that the Liveblocks update happens before the zustand update.)