diff --git a/packages/store/src/store.test.ts b/packages/store/src/store.test.ts index 33f88f33..d7c7090a 100644 --- a/packages/store/src/store.test.ts +++ b/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', () => { diff --git a/packages/store/src/store.ts b/packages/store/src/store.ts index 6cf14a36..22c99a8f 100644 --- a/packages/store/src/store.ts +++ b/packages/store/src/store.ts @@ -37,6 +37,7 @@ interface BatchEntry { } 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> without knowing T at this level. const _batchStores = new Map, BatchEntry>(); @@ -61,7 +62,9 @@ const _batchStores = new Map, BatchEntry>(); * ``` */ export function batch(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(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( }); } else { existing.nextState = nextState; + existing.commit = () => { state = nextState; persistState(); }; + existing.rollback = () => { state = existing.prevState; }; } } else { for (const listener of listeners) {