fix: rehydrate state trees that expire during a background task (#7585)

* fix: rehydrate state trees that expire during a background task

A background task releases the state lock while it runs. If the session
expires in that time, `async with self` gets a recreated tree, and the
delta carries only the handler's changes. The browser keeps the values
of the expired session.

A recreated `rx.State` root has empty `router_data`. A foreground event
already uses this signal to rehydrate. `StateProxy.__aenter__` now uses
the same signal: it restores the router from the event context and runs
`BaseStateEventProcessor._rehydrate`. The hydrate snapshot carries
`is_hydrated=false`, so the browser keeps its client storage.

Closes #7478

* fix: release the state lock when rehydrating an expired tree fails

Mark the proxy mutable before the rehydrate, so a failure or a
cancellation goes through __aexit__ and closes the state manager
context.

* docs: scope the rehydrate news to the loaded state tree

* test: disable oplock when simulating Redis expiry

Oplock keeps the tree cached in the process, so deleting the Redis keys
does not expire it.

* fix: keep a background rehydrate from loading the page the client left

The background task's route is the one it started on. After the tree
expires, the client may be on another page, so the rehydrate stays
routeless and the client's next event loads its current page.

Author: FarhanAliRaza

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,