RFC(events): lock ownership as a capability for delta flushes

news/+rfc-locked-root.misc.md

@@ -0,0 +1 @@
+RFC: add `LockedRoot`/`mint_locked_root` to the state manager as lock-ownership proof required by delta flushes.

packages/reflex-base/news/+rfc-locked-root.misc.md

@@ -0,0 +1 @@
+RFC: delta flushes take a `LockedRoot` capability minted only by `mint_locked_root`, making a flush without token-lock ownership a `TypeError` at the call site instead of a silent lost update.

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

@@ -13,6 +13,7 @@
 from typing import TYPE_CHECKING, Any
 
 from reflex.istate.data import RouterData
+from reflex.istate.manager import LockedRoot, mint_locked_root
 from reflex.istate.manager.token import BaseStateToken
 from reflex.istate.proxy import StateProxy
 from reflex.utils import types
@@ -196,7 +197,7 @@ async def _route_events(ctx: EventContext, events: Sequence[Event]) -> None:
 async def chain_updates(
     events: EventSpec | list[EventSpec] | None,
     handler_name: str,
-    root_state: BaseState | None = None,
+    root_state: LockedRoot | None = None,
 ) -> None:
     """Chain yielded events and emit a delta to the frontend.
 
@@ -206,8 +207,19 @@ async def chain_updates(
     Args:
         events: The events to queue with the update.
         handler_name: The name of the handler that yielded the events, used for error messages.
-        root_state: The root state of the app, no delta emitted if omitted.
+        root_state: Lock-ownership proof for the root to flush, no delta
+            emitted if omitted. A bare BaseState is refused: flushing without
+            the token lock discards concurrent writers' dirty vars.
+
+    Raises:
+        TypeError: If root_state is a bare BaseState instead of a LockedRoot.
     """
+    if root_state is not None and not isinstance(root_state, LockedRoot):
+        msg = (
+            "chain_updates requires a LockedRoot (see mint_locked_root); "
+            "flushing an unlocked root races concurrent events."
+        )
+        raise TypeError(msg)
     from reflex.event import Event
 
     ctx = EventContext.get()
@@ -219,12 +231,13 @@ async def chain_updates(
         # 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).
+        root = root_state.root
         try:
-            delta = await root_state._get_resolved_delta()
+            delta = await root._get_resolved_delta()
             if delta:
                 await ctx.emit_delta(delta)
         finally:
-            root_state._clean()
+            root._clean()
 
     # Convert valid EventHandler and EventSpec into Event
     if fixed_events := Event.from_event_type(
@@ -234,50 +247,52 @@ async def chain_updates(
 
 
 def ensure_locked(
-    state: BaseState | StateProxy, root_state: BaseState | None
-) -> BaseState | None:
+    state: BaseState | StateProxy, root_state: LockedRoot | None
+) -> LockedRoot | None:
     """The root to flush deltas from, only while the state lock is held.
 
-    Foreground handlers pass the locked root in. A background handler
+    Foreground handlers pass lock-ownership proof in. A background handler
     yielding from inside ``async with self`` is suspended while its proxy
-    still holds the lock, so flushing through the proxy's root keeps the
-    documented ordering: deltas reach the frontend before the yielded event.
-    Outside the proxy context there is no lock, and flushing there is the
-    unlocked snapshot/clean that discards concurrent writes. Evaluated per
-    yield: a generator can move between inside and outside the context.
+    still holds the lock, so minting for the proxy's root here is the audited
+    claim of that precondition, and flushing it keeps the documented ordering:
+    deltas reach the frontend before the yielded event. Outside the proxy
+    context there is no lock, and flushing there is the unlocked
+    snapshot/clean that discards concurrent writes. Evaluated per yield: a
+    generator can move between inside and outside the context.
 
     Args:
         state: The state the handler runs against, possibly a StateProxy.
-        root_state: The locked root passed by foreground callers, if any.
+        root_state: Lock-ownership proof passed by foreground callers, if any.
 
     Returns:
-        The root state to flush, or None when the lock is not held.
+        Lock-ownership proof for the root to flush, or None when the lock is
+        not held.
     """
     if root_state is not None:
         return root_state
     if isinstance(state, StateProxy) and state._is_mutable():
-        return state.__wrapped__._get_root_state()
+        return mint_locked_root(state.__wrapped__._get_root_state())
     return None
 
 
 async def process_event(
     handler: EventHandler,
     payload: dict,
     state: BaseState | StateProxy,
-    root_state: BaseState | None,
+    root_state: LockedRoot | None,
 ):
     """Process event.
 
     Args:
         handler: EventHandler to process.
         payload: The event payload.
         state: State to process the handler.
-        root_state: The root state of the app, used for emitting deltas. Pass
-            None when the caller does not hold the state lock (background
-            tasks): computing and cleaning a delta on an unlocked root races
-            concurrent events on a shared state tree, and background state
-            changes are emitted by the ``async with self`` context exits
-            instead.
+        root_state: Lock-ownership proof for the root to flush deltas from.
+            Pass None when the caller does not hold the state lock
+            (background tasks): computing and cleaning a delta on an unlocked
+            root races concurrent events on a shared state tree, and
+            background state changes are emitted by the ``async with self``
+            context exits instead.
 
     Raises:
         ValueError: If a string value is received for an int or float type and cannot be converted.
@@ -368,17 +383,18 @@ async def _rehydrate(self, root_state: BaseState):
         ):
             return
 
+        locked = mint_locked_root(root_state)
         await process_event(
             handler=State.event_handlers["hydrate"],
             payload={},
             state=root_state,
-            root_state=root_state,
+            root_state=locked,
         )
         await process_event(
             handler=OnLoadInternalState.event_handlers["on_load_internal"],
             payload={},
             state=await root_state.get_state(OnLoadInternalState),
-            root_state=root_state,
+            root_state=locked,
         )
 
     async def _execute_event(
@@ -445,7 +461,7 @@ async def _execute_event(
                     handler=registered_handler.handler,
                     payload=event.payload,
                     state=substate,
-                    root_state=root_state,
+                    root_state=mint_locked_root(root_state),
                 )
                 return
         # Otherwise drop the state lock and start processing the background task
@@ -477,7 +493,7 @@ async def _execute_event(
             ) as flush_state:
                 await chain_updates(
                     None,
-                    root_state=flush_state._get_root_state(),
+                    root_state=mint_locked_root(flush_state._get_root_state()),
                     handler_name=registered_handler.handler.fn.__qualname__,
                 )
 

reflex/istate/manager/__init__.py

@@ -266,6 +266,56 @@ def reset_disk_state_manager():
             path.unlink()
 
 
+_LOCKED_ROOT_MINT = object()
+
+
+@dataclasses.dataclass(frozen=True)
+class LockedRoot:
+    """Proof that the holder acquired the token state lock for this root.
+
+    Delta work (snapshot, resolve, emit, clean) on a shared state tree is only
+    safe while the token lock is held; a bare ``BaseState`` carries no record
+    of that. Functions that flush deltas accept this wrapper instead, so an
+    unlocked flush is unrepresentable rather than a code-review catch.
+
+    Only ``mint_locked_root`` creates instances. Do not construct directly.
+    """
+
+    root: "BaseState"
+    _mint: object = None
+
+    def __post_init__(self):
+        """Refuse construction outside mint_locked_root.
+
+        Raises:
+            TypeError: If constructed without the module-private mint token.
+        """
+        if self._mint is not _LOCKED_ROOT_MINT:
+            msg = (
+                "LockedRoot must be minted via mint_locked_root, from a caller "
+                "that holds the token state lock."
+            )
+            raise TypeError(msg)
+
+
+def mint_locked_root(root: "BaseState") -> LockedRoot:
+    """Assert lock ownership of a root state, making it flushable.
+
+    Call only while the token state lock for this root is held: inside a
+    ``state_manager.modify_state``/``modify_state_with_links`` block, or on a
+    ``StateProxy``'s freshly-fetched root while the proxy is mutable (its
+    ``async with self`` holds the same lock). Every call site is an auditable
+    claim of that precondition.
+
+    Args:
+        root: The root state acquired under the token lock.
+
+    Returns:
+        The lock-ownership wrapper accepted by delta-flushing functions.
+    """
+    return LockedRoot(root=root, _mint=_LOCKED_ROOT_MINT)
+
+
 def get_state_manager() -> StateManager:
     """Get the state manager for the app that is currently running.
 

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

@@ -543,10 +543,11 @@ async def test_ensure_locked_returns_a_root_only_while_the_lock_is_held(
     real_base_state_processor: BaseStateEventProcessor,
     token: str,
 ):
-    """ensure_locked passes a locked root through and refuses everything else.
+    """ensure_locked passes lock-ownership proof through and refuses the rest.
 
-    Foreground callers hand in the root they locked; a plain substate or an
-    un-entered proxy holds no lock, so there is nothing safe to flush.
+    Foreground callers hand in the proof they minted under the lock; a plain
+    substate or an un-entered proxy holds no lock, so there is nothing safe
+    to flush.
 
     Args:
         wired_app: The App wired to the processor's state manager.
@@ -555,6 +556,7 @@ async def test_ensure_locked_returns_a_root_only_while_the_lock_is_held(
     """
     from reflex_base.event.processor.base_state_processor import ensure_locked
 
+    from reflex.istate.manager import mint_locked_root
     from reflex.istate.proxy import StateProxy
 
     root_ctx = real_base_state_processor._root_context
@@ -565,11 +567,48 @@ async def test_ensure_locked_returns_a_root_only_while_the_lock_is_held(
     )
     substate = await root.get_state(OnLoadInternalState)
 
-    assert ensure_locked(substate, root) is root
+    minted = mint_locked_root(root)
+    assert ensure_locked(substate, minted) is minted
     assert ensure_locked(substate, None) is None
     assert ensure_locked(StateProxy(substate), None) is None
 
 
+async def test_chain_updates_refuses_a_bare_root_state(
+    wired_app: App,
+    real_base_state_processor: BaseStateEventProcessor,
+    token: str,
+):
+    """Flushing requires lock-ownership proof, not a bare state.
+
+    An unlocked flush discards concurrent writers' dirty vars; requiring a
+    ``LockedRoot`` makes that mistake a TypeError at the call site instead of
+    a silent lost update.
+
+    Args:
+        wired_app: The App wired to the processor's state manager.
+        real_base_state_processor: The unmocked BaseStateEventProcessor.
+        token: The client token.
+    """
+    from reflex_base.event.processor.base_state_processor import chain_updates
+
+    from reflex.istate.manager import LockedRoot, mint_locked_root
+
+    root_ctx = real_base_state_processor._root_context
+    assert root_ctx is not None
+    EventContext.set(root_ctx.fork(token=token))
+    root = await root_ctx.state_manager.get_state(
+        BaseStateToken(ident=token, cls=State)
+    )
+
+    with pytest.raises(TypeError, match="LockedRoot"):
+        await chain_updates(None, root_state=root, handler_name="bare")  # pyright: ignore[reportArgumentType]
+
+    with pytest.raises(TypeError, match="mint_locked_root"):
+        LockedRoot(root=root)
+
+    await chain_updates(None, root_state=mint_locked_root(root), handler_name="minted")
+
+
 async def test_failed_context_enter_does_not_mark_the_proxy_entered(
     wired_app: App,
     real_base_state_processor: BaseStateEventProcessor,