fix: subscribing withAtomEffect gives stale value (#80)

* fix: subscribing withAtomEffect gives stale value

* update version to 2.1.5

Author: dmaskasky

package.json

@@ -1,7 +1,7 @@
 {
   "name": "jotai-effect",
   "description": "👻🔁",
-  "version": "2.1.4",
+  "version": "2.1.5",
   "type": "module",
   "author": "David Maskasky",
   "contributors": [

src/withAtomEffect.ts

@@ -43,6 +43,7 @@ export function withAtomEffect<T extends Atom<unknown>>(
     const ensureAtomState = buildingBlocks[11]
     const flushCallbacks = buildingBlocks[12]
     const readAtomState = buildingBlocks[14]
+    const invalidateDependents = buildingBlocks[15]
     const mountDependencies = buildingBlocks[17]
     const mountAtom = buildingBlocks[18]
     const unmountAtom = buildingBlocks[19]
@@ -95,8 +96,16 @@ export function withAtomEffect<T extends Atom<unknown>>(
       }
     })
     storeHooks.m.add(targetWithEffect, function mountEffect() {
+      const atomState = ensureAtomState(store, targetWithEffect)
+      const { n } = atomState
       mountAtom(store, effectAtom)
       flushCallbacks(store)
+      if (n !== atomState.n) {
+        const unsub = storeHooks.f.add(() => {
+          invalidateDependents(store, targetWithEffect)
+          unsub()
+        })
+      }
     })
     storeHooks.u.add(targetWithEffect, function unmountEffect() {
       unmountAtom(store, effectAtom)

tests/withAtomEffect.test.ts

@@ -394,4 +394,33 @@ describe('withAtomEffect', () => {
     u1()
     u2()
   })
+
+  it('subscribing to an enhanced atom after it has been written to does not cause a stale read', () => {
+    const atomA = atom<boolean>()
+    atomA.debugLabel = 'atomA'
+    const atomB = withAtomEffect(atom(false), (get, set) => {
+      if (get(atomA) === true) {
+        set(atomB, true)
+      }
+    })
+    atomB.debugLabel = 'atomB'
+    const atomC = atom((get) => get(atomB))
+    atomC.debugLabel = 'atomC'
+
+    const store = createDebugStore()
+    expect(store.get(atomA)).toBe(undefined)
+    expect(store.get(atomB)).toBe(false)
+    expect(store.get(atomC)).toBe(false)
+
+    store.set(atomA, true)
+    expect(store.get(atomA)).toBe(true)
+    expect(store.get(atomB)).toBe(false)
+    expect(store.get(atomC)).toBe(false)
+
+    store.sub(atomC, () => {})
+    expect(store.get(atomA)).toBe(true)
+    expect(store.get(atomB)).toBe(true)
+
+    expect(store.get(atomC)).toBe(true)
+  })
 })