Fix stale substates after background task state expiry

news/+7478-refresh-expired-background-state.bugfix.md

@@ -0,0 +1 @@
+Mark recreated state trees dirty so background-task completion deltas refresh expired substate values in the client.

reflex/istate/manager/__init__.py

@@ -5,7 +5,8 @@
 import logging
 import sys
 from abc import ABC, abstractmethod
-from collections.abc import AsyncIterator
+from collections.abc import AsyncIterator, Iterator
+from contextvars import ContextVar
 from types import TracebackType
 from typing import TYPE_CHECKING, Generic, TypedDict, overload
 
@@ -19,6 +20,9 @@
 from reflex.utils import console, prerequisites
 
 logger = logging.getLogger(__name__)
+_STATE_RECREATION_TRACKER: ContextVar[list["BaseState"] | None] = ContextVar(
+    "state_recreation_tracker", default=None
+)
 
 if TYPE_CHECKING:
     from reflex.state import BaseState
@@ -333,6 +337,56 @@ async def __aexit__(
         return await self._state_context.__aexit__(exc_type, exc, tb) or suppressed
 
 
+@contextlib.contextmanager
+def _track_state_recreation() -> Iterator[list["BaseState"]]:
+    """Track state instances recreated inside the current async context.
+
+    Yields:
+        The state instances reported as recreated before the context exits.
+    """
+    recreated_states: list[BaseState] = []
+    tracker_token = _STATE_RECREATION_TRACKER.set(recreated_states)
+    try:
+        yield recreated_states
+    finally:
+        _STATE_RECREATION_TRACKER.reset(tracker_token)
+
+
+@contextlib.contextmanager
+def _no_state_recreation_tracking() -> Iterator[None]:
+    """Temporarily prevent child tasks from inheriting state tracking."""
+    tracker_token = _STATE_RECREATION_TRACKER.set(None)
+    try:
+        yield
+    finally:
+        _STATE_RECREATION_TRACKER.reset(tracker_token)
+
+
+def _report_state_recreation(state: "BaseState") -> None:
+    """Report a state recreated during a tracked modification context.
+
+    Args:
+        state: The newly recreated state.
+    """
+    if (recreated_states := _STATE_RECREATION_TRACKER.get()) is not None:
+        recreated_states.append(state)
+
+
+def _report_state_tree_recreation(state: "BaseState") -> None:
+    """Report every state in a newly recreated tree.
+
+    Args:
+        state: The root of the recreated state tree.
+    """
+    if (recreated_states := _STATE_RECREATION_TRACKER.get()) is None:
+        return
+    pending = [state]
+    while pending:
+        current = pending.pop()
+        recreated_states.append(current)
+        pending.extend(current.substates.values())
+
+
 def _release_state_tree(state: "BaseState"):
     """Clear the substate links of a purged state tree.
 

reflex/istate/manager/disk.py

@@ -19,6 +19,7 @@
     StateModificationContext,
     _default_token_expiration,
     _release_state_tree,
+    _report_state_recreation,
 )
 from reflex.istate.manager.token import TOKEN_TYPE, BaseStateToken, StateToken
 from reflex.state import BaseState
@@ -186,6 +187,8 @@ async def populate_substates(
             instance.parent_state = state
 
             await self.populate_substates(token, instance, root_state)
+            if instance is fresh_instance:
+                _report_state_recreation(instance)
 
     @override
     async def get_state(
@@ -216,6 +219,7 @@ async def get_state(
             fresh_root_state = root_state_cls(_reflex_internal_init=True)
             if root_state is None:
                 root_state = fresh_root_state
+                _report_state_recreation(root_state)
             elif not isinstance(root_state, BaseState):
                 msg = "Deserialized state is not an instance of BaseState, cannot populate substates."
                 raise TypeError(msg)

reflex/istate/manager/memory.py

@@ -12,7 +12,9 @@
     StateManager,
     StateModificationContext,
     _default_token_expiration,
+    _no_state_recreation_tracking,
     _release_state_tree,
+    _report_state_tree_recreation,
 )
 from reflex.istate.manager.token import TOKEN_TYPE, BaseStateToken, StateToken
 
@@ -56,9 +58,9 @@ def _get_or_create_state(self, token: StateToken[TOKEN_TYPE]) -> TOKEN_TYPE:
         key = token.cache_key
         if key not in self.states:
             if isinstance(token, BaseStateToken):
-                self.states[key] = token.cls.get_root_state()(
-                    _reflex_internal_init=True
-                )
+                state = token.cls.get_root_state()(_reflex_internal_init=True)
+                _report_state_tree_recreation(state)
+                self.states[key] = state
             else:
                 self.states[key] = token.cls()
         return cast(TOKEN_TYPE, self.states[key])
@@ -139,10 +141,11 @@ def _ensure_expiration_task(self):
         """Ensure the expiration background task is running."""
         if self._expiration_task is None or self._expiration_task.done():
             asyncio.get_running_loop()  # Ensure we're in an event loop.
-            self._expiration_task = asyncio.create_task(
-                self._expire_states(),
-                name="StateManagerMemory|Expiration",
-            )
+            with _no_state_recreation_tracking():
+                self._expiration_task = asyncio.create_task(
+                    self._expire_states(),
+                    name="StateManagerMemory|Expiration",
+                )
 
     @override
     async def get_state(self, token: StateToken[TOKEN_TYPE]) -> TOKEN_TYPE:

reflex/istate/manager/redis.py

@@ -28,6 +28,8 @@
     StateManager,
     StateModificationContext,
     _default_token_expiration,
+    _no_state_recreation_tracking,
+    _report_state_recreation,
 )
 from reflex.istate.manager.token import TOKEN_TYPE, BaseStateToken, StateToken
 from reflex.state import BaseState
@@ -127,16 +129,19 @@ class OplockFound(Exception):  # noqa: N818
     """Indicates that an opportunistic lock was found."""
 
 
-# KEYS: the lock key, then the state keys. ARGV: the lock id, the state
-# expiration in seconds, then the serialized states in KEYS order. Writes only
-# while the lock is held by ARGV[1], and returns the lock's remaining PTTL, or
-# nil when the lock changed hands or expired and nothing was written.
+# KEYS: the lock key, the state manifest key, then the state keys.
+# ARGV: the lock id, the state expiration in seconds, then serialized states
+# in KEYS order. The manifest tracks which per-state keys have been persisted.
 _FENCED_SAVE_SCRIPT = """
 if redis.call('GET', KEYS[1]) ~= ARGV[1] then
     return nil
 end
-for i = 2, #KEYS do
-    redis.call('SET', KEYS[i], ARGV[i + 1], 'EX', ARGV[2])
+for i = 3, #KEYS do
+    redis.call('SET', KEYS[i], ARGV[i], 'EX', ARGV[2])
+    redis.call('SADD', KEYS[2], KEYS[i])
+end
+if #KEYS > 2 then
+    redis.call('PEXPIRE', KEYS[2], tonumber(ARGV[2]) * 1000)
 end
 return redis.call('PTTL', KEYS[1])
 """
@@ -357,18 +362,39 @@ async def get_state(
             key=lambda x: x.get_full_name(),
         )
 
-        # Read the tree atomically with one command instead of a transaction
-        # containing a GET for each state. An already populated tree needs no IO.
-        redis_states = (
-            await self.redis.mget([
-                token._state_key(state_cls) for state_cls in required_state_classes
-            ])
-            if required_state_classes
-            else []
-        )
+        # Read state values and their persisted-key manifest in one round trip.
+        state_keys = [
+            token._state_key(state_cls) for state_cls in required_state_classes
+        ]
+        if required_state_classes:
+            pipeline = self.redis.pipeline(transaction=False)
+            pipeline.mget(state_keys)
+            pipeline.smembers(token._state_manifest_key())
+            redis_states, manifest_keys = await pipeline.execute()
+            persisted_state_keys = {
+                key.decode() if isinstance(key, bytes) else key for key in manifest_keys
+            }
+            untracked_keys = [
+                key
+                for key, redis_state in zip(state_keys, redis_states, strict=True)
+                if redis_state is not None and key not in persisted_state_keys
+            ]
+            if untracked_keys:
+                # Learn existing Redis keys on first access after introducing the manifest.
+                manifest_pipeline = self.redis.pipeline(transaction=False)
+                for key in untracked_keys:
+                    manifest_pipeline.sadd(token._state_manifest_key(), key)
+                manifest_pipeline.pexpire(
+                    token._state_manifest_key(), self.token_expiration * 1000
+                )
+                await manifest_pipeline.execute()
+        else:
+            redis_states = []
+            persisted_state_keys = set()
 
-        for state_cls, redis_state in zip(
+        for state_cls, state_key, redis_state in zip(
             required_state_classes,
+            state_keys,
             redis_states,
             strict=True,
         ):
@@ -379,11 +405,13 @@ async def get_state(
                 with contextlib.suppress(StateSchemaMismatchError):
                     state = BaseState._deserialize(data=redis_state)
             if state is None:
-                # Key didn't exist or schema mismatch so create a new instance for this token.
+                # A missing never-persisted state has no stale client value to reset.
                 state = state_cls(
                     init_substates=False,
                     _reflex_internal_init=True,
                 )
+                if redis_state is not None or state_key in persisted_state_keys:
+                    _report_state_recreation(state)
             state_full_name = state.get_full_name()
             flat_state_tree[state_full_name] = state
             if state.get_parent_state() is not None:
@@ -438,11 +466,18 @@ async def set_state(
             pickle_state = token.serialize(state)
             writes = [(str(token), pickle_state)] if pickle_state else []
 
+        manifest_key = (
+            token._state_manifest_key()
+            if isinstance(token, BaseStateToken)
+            else f"{token.lock_key}_reflex_state_manifest"
+        )
         if lock_id is None:
             if writes:
                 pipeline = self.redis.pipeline(transaction=False)
                 for key, pickle_state in writes:
                     pipeline.set(key, pickle_state, ex=self.token_expiration)
+                    pipeline.sadd(manifest_key, key)
+                pipeline.pexpire(manifest_key, self.token_expiration * 1000)
                 await pipeline.execute()
             return
 
@@ -460,8 +495,9 @@ async def set_state(
         fenced_save = cast("Callable[..., Awaitable[int | None]]", self.redis.eval)
         pttl = await fenced_save(
             _FENCED_SAVE_SCRIPT,
-            1 + len(writes),
+            2 + len(writes),
             lock_key,
+            manifest_key,
             *(key for key, _ in writes),
             lock_id,
             self.token_expiration,
@@ -798,10 +834,11 @@ async def lease_breaker():
                 lock_key not in self._local_leases
                 and await self._n_lock_contenders(self._lock_key(token)) == 0
             ):
-                self._local_leases[lock_key] = task = asyncio.create_task(
-                    lease_breaker(),
-                    name=f"reflex_lease_breaker|{lock_key}|{lock_id.decode()}",
-                )
+                with _no_state_recreation_tracking():
+                    self._local_leases[lock_key] = task = asyncio.create_task(
+                        lease_breaker(),
+                        name=f"reflex_lease_breaker|{lock_key}|{lock_id.decode()}",
+                    )
                 # Fetch the requested state into the cache.
                 self._cached_states[lock_key] = await self.get_state(token)
                 return task
@@ -908,12 +945,13 @@ async def _subscribe_lock_updates(self):
 
     def _ensure_lock_task(self) -> None:
         """Ensure the lock updates subscriber task is running."""
-        ensure_task(
-            owner=self,
-            task_attribute="_lock_task",
-            coro_function=self._subscribe_lock_updates,
-            suppress_exceptions=[Exception],
-        )
+        with _no_state_recreation_tracking():
+            ensure_task(
+                owner=self,
+                task_attribute="_lock_task",
+                coro_function=self._subscribe_lock_updates,
+                suppress_exceptions=[Exception],
+            )
 
     async def _ensure_lock_task_subscribed(self, timeout: float | None = None) -> None:
         """Ensure the lock updates subscriber task is running and subscribed to avoid missing notifications.

reflex/istate/manager/token.py

@@ -162,6 +162,14 @@ def _state_key(self, cls: type[BaseState]) -> str:
         """
         return f"{self.ident}_{cls.get_full_name()}"
 
+    def _state_manifest_key(self) -> str:
+        """Get the Redis key recording state keys previously persisted for this token.
+
+        Returns:
+            The Redis key for this token's persisted-state manifest.
+        """
+        return f"{self.ident}_reflex_state_manifest"
+
     def __str__(self) -> str:
         """The key used in the underlying StateManager store.
 

reflex/istate/proxy.py

@@ -211,6 +211,7 @@ def __init__(
         self._self_actx_lock = asyncio.Lock()
         self._self_actx_lock_holder = None
         self._self_parent_state_proxy = parent_state_proxy
+        self._self_recreated_states: list[BaseState] = []
         # Whether `async with self` was ever entered; a background handler that
         # never did emitted no delta, so the processor flushes once for it.
         self._self_entered_context = False
@@ -266,12 +267,29 @@ async def __aenter__(self) -> Self:
             self._self_actx = ctx.state_manager.modify_state_with_links(
                 token=self._self_substate_token, event=self._self_event
             )
-            mutable_state = await self._self_actx.__aenter__()
+            from reflex.istate.manager import _track_state_recreation
+
+            with _track_state_recreation() as recreated_states:
+                mutable_state = await self._self_actx.__aenter__()
+            self._self_recreated_states = recreated_states
+            from reflex.state import State
+
+            for state in recreated_states:
+                state_cls = type(state)
+                if state_cls is not State:
+                    state.dirty_vars.update(
+                        name
+                        for name, field in state_cls.get_fields().items()
+                        if field._owner is state_cls
+                    )
+                state.dirty_vars.update(state.computed_vars)
+                for computed_var in state.computed_vars.values():
+                    computed_var.mark_dirty(instance=state)
+                state._mark_dirty()
+            refreshed_state = mutable_state.get_substate(self._self_substate_path)
             self._self_mutable = True
             self._self_entered_context = True
-            super().__setattr__(
-                "__wrapped__", mutable_state.get_substate(self._self_substate_path)
-            )
+            super().__setattr__("__wrapped__", refreshed_state)
         except (Exception, asyncio.CancelledError):
             # Restore the proxy to a consistent state since __aexit__ will not be called when __aenter__ raises.
             await self.__aexit__(*sys.exc_info())
@@ -292,7 +310,15 @@ async def __aexit__(self, *exc_info: Any) -> None:
         try:
             if self._self_mutable and self._self_actx is not None:
                 root_state = self.__wrapped__._get_root_state()
-                delta = await root_state._get_resolved_delta()
+                touched_states = [
+                    (state, state._was_touched) for state in self._self_recreated_states
+                ]
+                try:
+                    delta = await root_state._get_resolved_delta()
+                finally:
+                    for state, was_touched in touched_states:
+                        state._was_touched = was_touched
+                    self._self_recreated_states.clear()
                 root_state._clean()
                 # When the frontend vars are modified emit the delta to the frontend.
                 if delta: