[2.11.1+] Derived atoms are not re-evaluated after referencing them immediately after updating an atom that changes dependencies.

[yuneco]

Bug Description

Starting from jotai v2.11.1, there is an issue where derived atoms do not re-evaluate under certain conditions.
The following code reproduces this problem.

import { atom, useAtomValue, useSetAtom } from "jotai";

const countAtom = atom(0);
const isActiveAtom = atom(false);

const activeCountAtom = atom((get) => {
  return get(isActiveAtom) ? get(countAtom) : undefined;
});

const activateAction = atom(null, (get, set) => {
  set(isActiveAtom, true);
  get(activeCountAtom); // 🌟 This "get" seems to cause the problem.
});

const App = () => {
  const count = useAtomValue(activeCountAtom);
  const setCount = useSetAtom(countAtom);
  const activate = useSetAtom(activateAction);

  return (
    <>
      <div>count: {count ?? "inactive"}</div>
      <button onClick={() => activate()}>1. activate counter</button>
      <button onClick={() => setCount((c) => c + 1)}>2. count up</button>
    </>
  );
};
export default App;

This code implements a simple counter. The displayed count is provided by a derived atom that only return value when isActiveAtom is true.

Running this code and executing activateAction and updating countAtom,although the counter is expected to increment,
the counter value does not change from 0.

Suspected Issue:
The problem seems to be caused by referencing a derived atom immediately after updating a dependent atom, which prevents further re-evaluations. Removing the seemingly innocuous get(activeCountAtom) inside activateAction makes the code function as expected.

When set(isActiveAtom, true) changes the dependencies of activeCountAtom, calling get(activeCountAtom) before the dependency recalculation caused by this change leads to the issue.

Reproduction Link

https://codesandbox.io/p/sandbox/fmk5xf

[dmaskasky]

Thank you for reporting.

Note to self, just thinking about this on my way in to work: when writeAtomState on activateAction is called, setter is called on isActiveAtom which updates the value to true. isActiveAtom is marked as dirty along with all its dependents; since isActiveAtom is currently mounted, this includes activeCountAtom. Calling getter inside writeAtomState should cause an eager readAtomState on activeCountAtom. Since activeCountAtom is dirty, this should recursively reevaluate all dependencies. In this case, countAtom is conditional and is not yet a dependency of activeCountAtom. So the issue lies with not doing a full readAtomState of activeCountAtom in which its read function is called again calling recomputeDependents prior to calling readAtomState on activeCountAtom. Because the readAtomState was called on activeCountAtom, that atom and its dependencies are removed from dirty. Therefore the final recomputeDependents call from flushPending does not process those atoms.

I wonder if this was fixed in 2.11.2. There was a similar issue reported that was caused by early returning from recursively marking dependents as dirty when the current atom was already dirty. That change was reverted in 2.11.2, so the full dependents tree is marked dirty whenever an atom is changed. The issue is present in 2.11.2.

@dai-shi #2965

[dai-shi]

Again, thanks for reporting. It's a bit shame. I'll look into it.