RFC(events): selective clean via a resolution-dirt ledger

news/+rfc-resolution-ledger.misc.md

@@ -0,0 +1 @@
+RFC: delta flushes clean selectively via a resolution-dirt ledger, so a write landing during delta resolution survives for the next harvest instead of being discarded by a clean that never snapshotted it.

packages/reflex-base/news/+rfc-resolution-ledger.misc.md

@@ -0,0 +1 @@
+RFC: chain_updates cleans selectively (dirty snapshot plus resolution ledger), keeping the SharedState fan-out capture complete while concurrent writers' dirt survives.

packages/reflex-base/src/reflex_base/event/processor/base_state_processor.py

@@ -209,22 +209,50 @@ async def chain_updates(
         root_state: The root state of the app, no delta emitted if omitted.
     """
     from reflex.event import Event
+    from reflex.state import _recording_resolution_dirt, _resolve_delta
 
     ctx = EventContext.get()
 
     if root_state is not None:
         # Emit deltas first, so any frontend events are processed with the
-        # latest state. The clean deliberately runs after resolution: the
-        # SharedState fan-out captures its dirty vars at clean time, and
-        # resolving the delta is what re-marks linked vars through the patch
-        # machinery, so cleaning earlier would fan out a stale set (see
-        # tests/integration/test_linked_state.py).
-        try:
-            delta = await root_state._get_resolved_delta()
-            if delta:
-                await ctx.emit_delta(delta)
-        finally:
-            root_state._clean()
+        # latest state. The clean clears exactly what this flush published:
+        # the dirty snapshot plus whatever resolution itself dirtied (the
+        # SharedState patch machinery reached through computed vars records
+        # into the ledger), so the fan-out capture stays complete while a
+        # concurrent writer's dirt survives for the next harvest instead of
+        # being discarded by a clean that never snapshotted it.
+        delta = root_state.get_delta()
+        if not delta or not any(
+            inspect.iscoroutine(value)
+            for subdelta in delta.values()
+            for value in subdelta.values()
+        ):
+            # A coroutine-free resolution never yields the loop, so nothing
+            # can interleave with this flush before the emit; the selective
+            # machinery below would record and subtract nothing. Take the
+            # plain path the sync pipeline has always had.
+            try:
+                if delta:
+                    await ctx.emit_delta(delta)
+            finally:
+                root_state._clean()
+        else:
+            # The flush suspends while resolving, which is the window the
+            # selective clean exists for: snapshot what this flush publishes,
+            # record what resolution itself dirties, and clean exactly that
+            # union, so a concurrent writer's dirt survives for the next
+            # harvest instead of being discarded by a clean that never
+            # snapshotted it.
+            flushed = root_state._snapshot_dirty_vars()
+            try:
+                with _recording_resolution_dirt(flushed_by_resolution := {}):
+                    delta = await _resolve_delta(delta)
+                for state_name, var_names in flushed_by_resolution.items():
+                    flushed.setdefault(state_name, set()).update(var_names)
+                if delta:
+                    await ctx.emit_delta(delta)
+            finally:
+                root_state._clean_flushed(flushed)
 
     # Convert valid EventHandler and EventSpec into Event
     if fixed_events := Event.from_event_type(

reflex/istate/proxy.py

@@ -614,7 +614,7 @@ def _mark_dirty(
         Returns:
             The result of the wrapped function.
         """
-        self._self_state.dirty_vars.add(self._self_field_name)
+        self._self_state._record_dirty_var(self._self_field_name)
         self._self_state._mark_dirty()
         if wrapped is not None:
             return wrapped(*args, **(kwargs or {}))

reflex/istate/shared.py

@@ -149,6 +149,27 @@ def _clean(self):
             previous_dirty_vars.update(self.dirty_vars)
         super()._clean()
 
+    @_override_base_method
+    def _clean_flushed(self, flushed: dict[str, set[str]]) -> None:
+        """Selective clean that still captures this flush's fan-out seed.
+
+        The capture mirrors ``_clean``: what this event published for this
+        state, which under a selective clean is exactly its slice of
+        ``flushed``, including resolution-created dirt from the ledger. A
+        concurrent writer's dirt is neither captured nor cleared.
+
+        Args:
+            flushed: Mapping of full state name to the var names the flush
+                snapshotted or created during resolution.
+        """
+        mine = self.dirty_vars & flushed.get(self.get_full_name(), set())
+        if (
+            previous_dirty_vars := getattr(self, "_previous_dirty_vars", None)
+        ) is not None and mine:
+            previous_dirty_vars.clear()
+            previous_dirty_vars.update(mine)
+        super()._clean_flushed(flushed)
+
     @_override_base_method
     def _mark_dirty(self):
         """Override BaseState._mark_dirty to avoid marking certain vars as dirty.

reflex/state.py

@@ -5,6 +5,7 @@
 import asyncio
 import builtins
 import contextlib
+import contextvars
 import copy
 import dataclasses
 import functools
@@ -279,6 +280,43 @@ def get_var_for_field(cls: type[BaseState], name: str, f: Field) -> Var:
 _DROP_FROM_DELTA: Final = object()
 
 
+_RESOLUTION_DIRT: contextvars.ContextVar[dict[str, set[str]] | None] = (
+    contextvars.ContextVar("reflex_resolution_dirt", default=None)
+)
+
+# Count of open recording contexts, so the per-dirty-mark fast path is a plain
+# global check instead of a contextvar lookup while no flush is resolving.
+_RESOLUTION_RECORDING_DEPTH = 0
+
+
+@contextlib.contextmanager
+def _recording_resolution_dirt(ledger: dict[str, set[str]]) -> Iterator[None]:
+    """Record vars dirtied by the current task tree into the ledger.
+
+    Delta resolution is not a pure read: computed vars reach ``get_state``,
+    which drives machinery (e.g. the SharedState patch) that marks vars
+    dirty. Contextvars are task-local and inherited by tasks created during
+    resolution, so writes made by the resolving task tree land in the ledger
+    while concurrent events' writes, made from tasks created elsewhere, do
+    not. That distinction is what lets a flush clean exactly what it
+    published and preserve everything else.
+
+    Args:
+        ledger: Mapping of full state name to var names, filled in place.
+
+    Yields:
+        None.
+    """
+    global _RESOLUTION_RECORDING_DEPTH
+    token = _RESOLUTION_DIRT.set(ledger)
+    _RESOLUTION_RECORDING_DEPTH += 1
+    try:
+        yield
+    finally:
+        _RESOLUTION_RECORDING_DEPTH -= 1
+        _RESOLUTION_DIRT.reset(token)
+
+
 async def _resolve_delta(delta: Delta) -> Delta:
     """Await all coroutines in the delta, dropping keys that resolve to the drop sentinel.
 
@@ -1519,7 +1557,7 @@ def __setattr__(self, name: str, value: Any):
 
         if name in self.backend_vars:
             self._backend_vars.__setitem__(name, value)
-            self.dirty_vars.add(name)
+            self._record_dirty_var(name)
             self._mark_dirty()
             return
 
@@ -1552,12 +1590,12 @@ def __setattr__(self, name: str, value: Any):
 
         # Add the var to the dirty list.
         if name in self.base_vars:
-            self.dirty_vars.add(name)
+            self._record_dirty_var(name)
             self._mark_dirty()
 
         # For now, handle router_data updates as a special case
         if name == constants.ROUTER_DATA:
-            self.dirty_vars.add(name)
+            self._record_dirty_var(name)
             self._mark_dirty()
 
     def reset(self):
@@ -1819,7 +1857,7 @@ def _mark_dirty_computed_vars(self) -> None:
                     defining_state = self._get_root_state().get_substate(
                         tuple(state_name.split("."))
                     )
-                defining_state.dirty_vars.add(cvar)
+                defining_state._record_dirty_var(cvar)
                 actual_var = defining_state.computed_vars.get(cvar)
                 if actual_var is not None:
                     actual_var.mark_dirty(instance=defining_state)
@@ -1888,10 +1926,16 @@ def get_delta(self) -> Delta:
         if len(subdelta) > 0:
             delta[self.get_full_name()] = subdelta
 
-        # Recursively find the substate deltas.
+        # Recursively find the substate deltas. A dirty name without an
+        # attached substate is legitimate under selective cleaning: a prior
+        # flush preserved another writer's dirt, the tree was then fetched
+        # partially (redis fetches the handler's slice), and the dirt lives in
+        # the unfetched substate's own record until an event fetches it. It
+        # cannot contribute to this delta, so skip it, as _clean always has.
         substates = self.substates
         for substate in self.dirty_substates.union(self._always_dirty_substates):
-            delta.update(substates[substate].get_delta())
+            if (substate_instance := substates.get(substate)) is not None:
+                delta.update(substate_instance.get_delta())
 
         # Return the delta.
         return delta
@@ -1904,6 +1948,79 @@ async def _get_resolved_delta(self) -> Delta:
         """
         return await _resolve_delta(self.get_delta())
 
+    def _record_dirty_var(self, name: str) -> None:
+        """Mark one var dirty, recording it when delta resolution is active.
+
+        The single choke point for "a var became dirty", so a flush can tell
+        its own resolution-created dirt apart from concurrent writers. Known
+        boundary: dirt is name-only, so a concurrent rewrite of a var already
+        in a flush's snapshot, landing during that flush's resolution, is
+        still cleaned with it; distinguishing that needs per-var write
+        versions. Unreachable while every flush holds the token lock, which
+        all current callers do.
+
+        Args:
+            name: The var name to mark dirty on this state.
+        """
+        self.dirty_vars.add(name)
+        if _RESOLUTION_RECORDING_DEPTH:
+            self._note_resolution_dirt(name)
+
+    def _note_resolution_dirt(self, name: str) -> None:
+        """Record one var into the active flush's resolution ledger.
+
+        The cold half of ``_record_dirty_var``: reached only while a flush is
+        resolving, so the hot path pays one set-add and one falsy global
+        check.
+
+        Args:
+            name: The var name that became dirty during resolution.
+        """
+        if (ledger := _RESOLUTION_DIRT.get()) is not None:
+            ledger.setdefault(self.get_full_name(), set()).add(name)
+
+    def _snapshot_dirty_vars(self) -> dict[str, set[str]]:
+        """The dirty var names per state, for the flush that just snapshotted.
+
+        Walked over the same states ``get_delta`` visits, and including
+        backend vars, which never reach a delta but are captured for the
+        SharedState fan-out and must be cleaned once flushed.
+
+        Returns:
+            Mapping of full state name to a copy of its dirty var names.
+        """
+        snapshot: dict[str, set[str]] = {}
+        if self.dirty_vars:
+            snapshot[self.get_full_name()] = set(self.dirty_vars)
+        for substate_name in self.dirty_substates.union(self._always_dirty_substates):
+            substate = self.substates.get(substate_name)
+            if substate is not None:
+                snapshot.update(substate._snapshot_dirty_vars())
+        return snapshot
+
+    def _clean_flushed(self, flushed: dict[str, set[str]]) -> None:
+        """Clear exactly the dirty vars a flush published, keeping the rest.
+
+        Unlike ``_clean``, a var dirtied by a concurrent writer after the
+        flush's snapshot survives for the next harvest instead of being
+        discarded by a clean that never saw it.
+
+        Args:
+            flushed: Mapping of full state name to the var names the flush
+                snapshotted or created during resolution.
+        """
+        self._update_was_touched()
+        for substate_name in tuple(
+            self.dirty_substates.union(self._always_dirty_substates)
+        ):
+            substate = self.substates.get(substate_name)
+            if substate is None:
+                continue
+            substate._clean_flushed(flushed)
+            if not substate.dirty_vars and not substate.dirty_substates:
+                self.dirty_substates.discard(substate_name)
+        self.dirty_vars -= flushed.get(self.get_full_name(), set())
+
     def _mark_dirty(self):
         """Mark the substate and all parent states as dirty."""
         state_name = self.get_name()

tests/units/reflex_base/event/processor/test_base_state_processor.py

@@ -615,3 +615,72 @@ def raise_on_modify(*args, **kwargs):
         object.__setattr__(root_ctx.state_manager, "modify_state_with_links", original)
 
     assert proxy._self_entered_context is False
+
+
+async def test_chain_updates_keeps_writes_made_during_delta_resolution(
+    wired_app: App,
+    real_base_state_processor: BaseStateEventProcessor,
+    emitted_deltas: list[tuple[str, Mapping[str, Mapping[str, Any]]]],
+    token: str,
+):
+    """chain_updates must not clean a write it never snapshotted.
+
+    The snapshot and the clean happen in one step, before the resolved delta
+    is awaited or emitted, so a concurrent write landing during resolution
+    stays dirty for the next harvest even if a caller ever runs chain_updates
+    without holding the state lock.
+
+    Args:
+        wired_app: The App wired to the processor's state manager.
+        real_base_state_processor: The unmocked BaseStateEventProcessor.
+        emitted_deltas: List to capture emitted deltas.
+        token: The client token.
+    """
+    from reflex_base.event.processor.base_state_processor import chain_updates
+
+    resolving = asyncio.Event()
+    release = asyncio.Event()
+    hold_resolution = [False]
+
+    class MidResolveState(State):
+        victim: str = ""
+
+        @rx.var(cache=False)
+        async def window(self) -> int:
+            if hold_resolution[0]:
+                resolving.set()
+                await release.wait()
+            return 0
+
+    root_ctx = real_base_state_processor._root_context
+    assert root_ctx is not None
+    EventContext.set(root_ctx.fork(token=token))
+    state_manager = root_ctx.state_manager
+
+    try:
+        root = await state_manager.get_state(BaseStateToken(ident=token, cls=State))
+        substate = await root.get_state(MidResolveState)
+        root._clean()
+
+        hold_resolution[0] = True
+        flush = asyncio.ensure_future(
+            chain_updates(None, root_state=root, handler_name="unlocked_flush")
+        )
+        await resolving.wait()
+        substate.victim = "written"
+        hold_resolution[0] = False
+        release.set()
+        await flush
+
+        assert "victim" in substate.dirty_vars, (
+            "the write made during delta resolution was cleaned away"
+        )
+        await chain_updates(None, root_state=root, handler_name="second_flush")
+    finally:
+        State._always_dirty_substates.discard(MidResolveState.get_name())
+
+    state_name = MidResolveState.get_full_name()
+    victim_key = "victim" + FIELD_MARKER
+    assert any(
+        d.get(state_name, {}).get(victim_key) == "written" for _, d in emitted_deltas
+    ), f"the surviving write never reached a delta: {emitted_deltas}"

tests/units/test_state.py

@@ -5188,3 +5188,18 @@ async def _coro(value):  # noqa: RUF029 - a trivial coroutine value for the delt
     }
     resolved = await _resolve_delta(delta)
     assert resolved == {"s2": {"keep": 1}}
+
+
+def test_get_delta_skips_dirty_names_without_attached_substates() -> None:
+    """A dirty substate name with no attached instance is skipped, not a KeyError.
+
+    Under selective cleaning a prior flush can preserve another writer's dirt,
+    and a later redis fetch of the handler's slice restores the parent's
+    dirty_substates naming a substate that was not fetched. The dirt lives in
+    the unfetched substate's own record; the partial tree's delta walk must
+    pass over it the way _clean always has.
+    """
+    root = State(_reflex_internal_init=True)  # pyright: ignore [reportCallIssue]
+    root._clean()
+    root.dirty_substates.add("not___an___attached____substate")
+    assert root.get_delta() == {}