fix(store): prevent batch race condition causing silent state loss in nested batches

packages/store/src/store.test.ts

@@ -290,6 +290,136 @@ describe('batch', () => {
         expect(spy).toHaveBeenCalledOnce()
         expect(spy.mock.calls[0][0]).toEqual({ a: 1, b: 2 })
     })
+
+    it('nested batch does not lose intermediate state from outer batch', async () => {
+        const useStore = createStore((set) => ({
+            count: 0,
+            name: '',
+        }))
+        const spy = vi.fn()
+        useStore.subscribe(spy)
+
+        batch(() => {
+            useStore.setState({ count: 1 })
+            batch(() => {
+                useStore.setState({ name: 'test' })
+            })
+        })
+
+        await new Promise(resolve => queueMicrotask(resolve))
+
+        expect(spy).toHaveBeenCalledOnce()
+        expect(useStore.getState()).toEqual({ count: 1, name: 'test' })
+    })
+
+    it('consecutive batches do not interfere with each other', async () => {
+        const useStore = createStore((set) => ({
+            a: 0,
+            b: 0,
+        }))
+        const spy = vi.fn()
+        useStore.subscribe(spy)
+
+        batch(() => {
+            useStore.setState({ a: 1 })
+        })
+
+        batch(() => {
+            useStore.setState({ b: 2 })
+        })
+
+        await new Promise(resolve => queueMicrotask(resolve))
+
+        expect(useStore.getState()).toEqual({ a: 1, b: 2 })
+    })
+
+    it('stale microtask from old batch does not dispatch when new batch starts', async () => {
+        const useStore = createStore((set) => ({
+            x: 0,
+        }))
+        const spy = vi.fn()
+        useStore.subscribe(spy)
+
+        // Start first batch but manually prevent its microtask from firing
+        batch(() => {
+            useStore.setState({ x: 1 })
+        })
+
+        // Immediately start a second batch before the first microtask fires
+        batch(() => {
+            useStore.setState({ x: 2 })
+        })
+
+        await new Promise(resolve => queueMicrotask(resolve))
+
+        // Listener should only see the final value, not be called twice
+        expect(spy).toHaveBeenCalledOnce()
+        expect(useStore.getState().x).toBe(2)
+    })
+
+    it('mutate inside batch merges correctly with setState', async () => {
+        const useStore = createStore((set) => ({
+            count: 0,
+            label: '',
+        }))
+        const spy = vi.fn()
+        useStore.subscribe(spy)
+
+        batch(() => {
+            useStore.setState({ count: 5 })
+            useStore.mutate((draft) => {
+                draft.label = 'mutated'
+            })
+        })
+
+        await new Promise(resolve => queueMicrotask(resolve))
+
+        expect(spy).toHaveBeenCalledOnce()
+        expect(useStore.getState()).toEqual({ count: 5, label: 'mutated' })
+    })
+
+    it('multiple mutate calls inside batch coalesce into one notification', async () => {
+        const useStore = createStore((set) => ({
+            a: 0,
+            b: 0,
+            c: 0,
+        }))
+        const spy = vi.fn()
+        useStore.subscribe(spy)
+
+        batch(() => {
+            useStore.mutate((draft) => { draft.a = 1 })
+            useStore.mutate((draft) => { draft.b = 2 })
+            useStore.mutate((draft) => { draft.c = 3 })
+        })
+
+        await new Promise(resolve => queueMicrotask(resolve))
+
+        expect(spy).toHaveBeenCalledOnce()
+        expect(useStore.getState()).toEqual({ a: 1, b: 2, c: 3 })
+    })
+
+    it('batch rollback restores state before any updates including mutate', () => {
+        const useStore = createStore((set) => ({
+            count: 0,
+            label: '',
+        }))
+        const spy = vi.fn()
+        useStore.subscribe(spy)
+
+        try {
+            batch(() => {
+                useStore.setState({ count: 5 })
+                useStore.mutate((draft) => {
+                    draft.label = 'mutated'
+                })
+                throw new Error('abort')
+            })
+        } catch {}
+
+        expect(useStore.getState()).toEqual({ count: 0, label: '' })
+        expect(spy).not.toHaveBeenCalled()
+    })
 })
 
 describe('middleware', () => {

packages/store/src/store.ts

@@ -37,6 +37,7 @@ interface BatchEntry<T> {
 }
 
 let _batchDepth = 0;
+let _batchEpoch = 0;
 // Map store instance to batch entry. Using any for listener set type because
 // the batch mechanism operates on the raw Set<Listener<T>> without knowing T at this level.
 const _batchStores = new Map<Set<any>, BatchEntry<any>>();
@@ -61,7 +62,9 @@ const _batchStores = new Map<Set<any>, BatchEntry<any>>();
  * ```
  */
 export function batch<T>(fn: () => T): T {
+    const isOutermost = _batchDepth === 0;
     _batchDepth++;
+    if (isOutermost) _batchEpoch++;
     let threw = false;
     let res: any;
     try {
@@ -99,16 +102,21 @@ export function batch<T>(fn: () => T): T {
 
 function flushBatch(threw: boolean) {
     if (threw) {
-        for (const [, { prevState, rollback }] of _batchStores) {
+        for (const [, { rollback }] of _batchStores) {
             rollback();
         }
-        _batchStores.clear(); // Don't notify listeners with partial state
+        _batchStores.clear();
     } else {
         if (_batchStores.size === 0) return;
+        // Snapshot the current epoch so the microtask can bail out if a new
+        // batch has started before it runs.
+        const epochAtFlush = _batchEpoch;
         queueMicrotask(() => {
+            // A new batch started between the flush and this microtask — skip.
+            if (_batchEpoch !== epochAtFlush) return;
             const stores = Array.from(_batchStores.entries());
             _batchStores.clear();
-            for (const [listeners, { prevState, nextState ,commit }] of stores) {
+            for (const [listeners, { prevState, nextState, commit }] of stores) {
                 commit();
                 for (const listener of listeners) {
                     listener(nextState, prevState);
@@ -398,6 +406,8 @@ export function createStore<T extends object>(
                 });
             } else {
                 existing.nextState = nextState;
+                existing.commit = () => { state = nextState; persistState(); };
+                existing.rollback = () => { state = existing.prevState; };
             }
         } else {
             for (const listener of listeners) {