fix: setAtom should use the correct pending on unmount

pmndrs · Nov 14, 2024 · 0f0e1a2 · 0f0e1a2
1 parent d9091a1
commit 0f0e1a2
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 4 deletions.
diff --git a/src/vanilla/store.ts b/src/vanilla/store.ts
@@ -77,7 +77,7 @@ type Mounted = {
   /** Set of mounted atoms that depends on the atom. */
   readonly t: Set<AnyAtom>
   /** Function to run when the atom is unmounted. */
-  u?: OnUnmount
+  u?: (pending: Pending) => void
 }
 
 /**
@@ -602,18 +602,20 @@ const buildStore = (
         const mounted = atomState.m
         addPendingFunction(pending, () => {
           let isSync = true
+          let currPending = pending
           try {
             const onUnmount = atomOnMount(atom, (...args) => {
               try {
-                return writeAtomState(pending, atom, ...args)
+                return writeAtomState(currPending, atom, ...args)
               } finally {
                 if (!isSync) {
                   flushPending(pending)
                 }
               }
             })
             if (onUnmount) {
-              mounted.u = () => {
+              mounted.u = (pending) => {
+                currPending = pending
                 isSync = true
                 try {
                   onUnmount()
@@ -644,7 +646,7 @@ const buildStore = (
       // unmount self
       const onUnmount = atomState.m.u
       if (onUnmount) {
-        addPendingFunction(pending, onUnmount)
+        addPendingFunction(pending, () => onUnmount(pending))
       }
       delete atomState.m
       if (import.meta.env?.MODE !== 'production') {

diff --git a/tests/vanilla/store.test.tsx b/tests/vanilla/store.test.tsx
@@ -876,3 +876,17 @@ describe('should mount and trigger listeners even when an error is thrown', () =
     expect(listener).toHaveBeenCalledOnce()
   })
 })
+
+it('should use the correct pending on unmount', () => {
+  const store = createStore()
+  const a = atom(0)
+  const b = atom(0, (_, set, update: number) => set(a, update))
+  b.onMount = (setAtom) => () => setAtom(1)
+  const aListener = vi.fn()
+  store.sub(a, aListener)
+  const unsub = store.sub(b, () => {})
+  aListener.mockClear()
+  unsub()
+  expect(store.get(a)).toBe(1)
+  expect(aListener).toHaveBeenCalledTimes(1)
+})