fix: rehydrate state trees that expire during a background task

news/+background-task-expired-state-rehydrate.bugfix.md

@@ -0,0 +1 @@
+A background task that re-enters an expired session now rehydrates the loaded state tree, so the browser no longer shows stale values for those states.

packages/reflex-base/news/+background-task-expired-state-rehydrate.bugfix.md

@@ -0,0 +1 @@
+A background task that re-enters an expired session now rehydrates the loaded states. The browser drops their stale values and keeps its route and browser storage; its next event runs the current page's on_load.

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

@@ -393,7 +393,8 @@ class BaseStateEventProcessor(EventProcessor):
     frontend.
     """
 
-    async def _rehydrate(self, root_state: BaseState):
+    @staticmethod
+    async def _rehydrate(root_state: BaseState):
         """Rehydrate the state by calling the hydrate event handler.
 
         Args:

reflex/istate/proxy.py

@@ -272,6 +272,16 @@ async def __aenter__(self) -> Self:
             super().__setattr__(
                 "__wrapped__", mutable_state.get_substate(self._self_substate_path)
             )
+            from reflex.state import State
+
+            if type(mutable_state) is State and not mutable_state.router_data:
+                # The tree expired while the lock was released: rehydrate it, so
+                # the client drops stale values. The task's route may be a page
+                # the client left, so the tree stays routeless and the client's
+                # next event loads its current page.
+                from reflex_base.event.processor import BaseStateEventProcessor
+
+                await BaseStateEventProcessor._rehydrate(mutable_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())

tests/units/istate/test_proxy.py

@@ -16,6 +16,7 @@
 import pytest
 from reflex_base.constants.state import FIELD_MARKER
 from reflex_base.event.context import EventContext
+from reflex_base.event.processor import BaseStateEventProcessor
 from reflex_base.utils.exceptions import ImmutableStateError
 
 import reflex as rx
@@ -186,6 +187,30 @@ async def mock_modify_state_context(*args, **kwargs):
         pass
 
 
+@pytest.mark.asyncio
+async def test_state_proxy_releases_lock_when_rehydrate_fails(
+    attached_mock_event_context: EventContext, monkeypatch: pytest.MonkeyPatch
+):
+    """A failed rehydrate of an expired tree must release the state lock."""
+    state_proxy = StateProxy(ProxyTestState())
+
+    async def failing_rehydrate(root_state: BaseState):  # noqa: RUF029
+        msg = "Simulated rehydrate failure"
+        raise RuntimeError(msg)
+
+    monkeypatch.setattr(BaseStateEventProcessor, "_rehydrate", failing_rehydrate)
+    with pytest.raises(RuntimeError, match="Simulated rehydrate failure"):
+        async with state_proxy:
+            pass
+
+    state_token = BaseStateToken(
+        ident=attached_mock_event_context.token, cls=ProxyTestState
+    )
+    async with asyncio.timeout(1):
+        async with attached_mock_event_context.state_manager.modify_state(state_token):
+            pass
+
+
 @dataclasses.dataclass
 class TaggedModel:
     """A dataclass with a mutable attribute for attr-path refresh tests."""

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

@@ -492,6 +492,177 @@ async def fg(self):
     )
 
 
+def _expire_token(state_manager: StateManager, token: str) -> None:
+    """Drop every state stored for a token, as its expiration would.
+
+    Args:
+        state_manager: The memory or redis state manager holding the token.
+        token: The client token to expire.
+    """
+    if isinstance(state_manager, StateManagerMemory):
+        state_manager._purge_token(BaseStateToken(ident=token, cls=State))
+        return
+    assert isinstance(state_manager, StateManagerRedis)
+    keys = state_manager.redis._internals["keys"]  # pyright: ignore[reportAttributeAccessIssue]
+    for key in [key for key in keys if key.startswith(token.encode())]:
+        del keys[key]
+
+
+@pytest.mark.parametrize(
+    "processor_state_manager", ["in_process", "redis"], indirect=True
+)
+async def test_background_task_rehydrates_state_tree_expired_while_unlocked(
+    wired_app: App,
+    real_base_state_processor: BaseStateEventProcessor,
+    emitted_deltas: list[tuple[str, Mapping[str, Mapping[str, Any]]]],
+    token: str,
+):
+    """A background task re-entering an expired tree sends the full recreated state.
+
+    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.
+    """
+    waiting = asyncio.Event()
+    release = asyncio.Event()
+
+    class ExpiryState(State):
+        value: int = 0
+
+        @event(background=True)
+        async def bg(self):
+            waiting.set()
+            await release.wait()
+            async with self:
+                self.value = 100
+
+    class ExpiryChildState(ExpiryState):
+        child: int = 0
+
+        @event
+        def set_child(self):
+            self.child = 13
+
+    router_data = {
+        RouteVar.PATH: "/post/[pid]",
+        RouteVar.ORIGIN: "/post/abc",
+        RouteVar.QUERY: {"pid": "abc"},
+        RouteVar.CLIENT_TOKEN: token,
+    }
+
+    set_child, bg = (
+        Event.from_event_type(spec, router_data=router_data)[0]
+        for spec in (ExpiryChildState.set_child(), ExpiryState.bg())
+    )
+    assert real_base_state_processor._root_context is not None
+    state_manager = real_base_state_processor._root_context.state_manager
+    if isinstance(state_manager, StateManagerRedis):
+        # Oplock keeps the tree cached in the process, so deleting keys does not expire it.
+        state_manager._oplock_enabled = False
+    async with real_base_state_processor as processor:
+        await (await processor.enqueue(token, set_child))
+        await processor.enqueue(token, bg)
+        await asyncio.wait_for(waiting.wait(), timeout=2)
+        _expire_token(state_manager, token)
+        emitted_deltas.clear()
+        release.set()
+        await processor.join(5)
+
+    root_name = State.get_full_name()
+    child_name = ExpiryChildState.get_full_name()
+    is_hydrated_key = CompileVars.IS_HYDRATED + FIELD_MARKER
+    snapshots = [d for _, d in emitted_deltas if child_name in d]
+    assert snapshots, f"The expired child state was never resent: {emitted_deltas}"
+    snapshot = snapshots[0]
+    # Not hydrated, so the browser does not overwrite its client storage.
+    assert snapshot[root_name][is_hydrated_key] is False
+    assert snapshot[child_name]["child" + FIELD_MARKER] == 0
+    # Routeless: the client's next event sets its current page.
+    page = snapshot[root_name][constants.ROUTER_PAGE + FIELD_MARKER]
+    assert page.raw_path == ""
+    assert emitted_deltas[-1][1][ExpiryState.get_full_name()] == {
+        "value" + FIELD_MARKER: 100
+    }
+
+
+@pytest.mark.parametrize(
+    "processor_state_manager", ["in_process", "redis"], indirect=True
+)
+async def test_background_rehydrate_does_not_load_the_route_the_client_left(
+    wired_app: App,
+    real_base_state_processor: BaseStateEventProcessor,
+    token: str,
+):
+    """A tree that expires during a background task does not reload the task's old page.
+
+    Args:
+        wired_app: The App wired to the processor's state manager.
+        real_base_state_processor: The unmocked BaseStateEventProcessor.
+        token: The client token.
+    """
+    waiting = asyncio.Event()
+    release = asyncio.Event()
+
+    class LeftPageState(State):
+        loaded: list[str] = []
+
+        @event
+        def load_a(self):
+            self.loaded = [*self.loaded, "a"]
+
+        @event
+        def load_b(self):
+            self.loaded = [*self.loaded, "b"]
+
+        @event
+        def ping(self):
+            pass
+
+        @event(background=True)
+        async def bg(self):
+            waiting.set()
+            await release.wait()
+            async with self:
+                pass
+
+    wired_app.add_page(
+        lambda: rx.text("a"), route="/page-a", on_load=LeftPageState.load_a
+    )
+    wired_app.add_page(
+        lambda: rx.text("b"), route="/page-b", on_load=LeftPageState.load_b
+    )
+    assert real_base_state_processor._root_context is not None
+    state_manager = real_base_state_processor._root_context.state_manager
+    if isinstance(state_manager, StateManagerRedis):
+        # Oplock keeps the tree cached in the process, so deleting keys does not expire it.
+        state_manager._oplock_enabled = False
+
+    async with real_base_state_processor as processor:
+        await processor.enqueue(
+            token, _client_event(LeftPageState.bg(), _view("/page-a"))
+        )
+        await asyncio.wait_for(waiting.wait(), timeout=2)
+        # The client moves to /page-b, then the tree expires.
+        await _send(
+            processor, token, _client_event(LeftPageState.ping(), _view("/page-b"))
+        )
+        _expire_token(state_manager, token)
+        release.set()
+        await processor.join(5)
+        # The client's next event on /page-b.
+        await _send(
+            processor, token, _client_event(LeftPageState.ping(), _view("/page-b"))
+        )
+
+    async with _read_back(real_base_state_processor, token) as root:
+        assert (await root.get_state(LeftPageState)).loaded == ["b"], (
+            "the background rehydrate loaded the route the client had left"
+        )
+        assert root.router.url.path == "/page-b"
+
+
 async def test_background_yield_inside_context_flushes_delta_before_event(
     wired_app: App,
     real_base_state_processor: BaseStateEventProcessor,