Compute client storage dependents once per hydrate

news/+hydrate-storage-dependents.bugfix.md

@@ -0,0 +1 @@
+Cached computed vars that depend on client storage (`rx.LocalStorage`, `rx.SessionStorage`, `rx.Cookie`) are now evaluated once per page load instead of twice.

reflex/istate/delta.py

@@ -42,6 +42,12 @@
 )
 
 
+# True while a delta is built for dirty fields whose computed dependents the
+# client already holds current, so the traversal neither invalidates nor
+# evaluates any computed var.
+_send_fields_only: ContextVar[bool] = ContextVar("_send_fields_only", default=False)
+
+
 class _DeltaRecord(NamedTuple):
     """An uncached var value that counts as sent once the delta delivers it."""
 
@@ -227,8 +233,12 @@ def build_delta(state: BaseState) -> Delta:
     """
     delta = {}
 
-    state._mark_dirty_computed_vars()
-    delta_vars = state.dirty_vars & state._frontend_var_names
+    if _send_fields_only.get():
+        delta_vars = state.dirty_vars & state._frontend_var_names
+        delta_vars.difference_update(state.computed_vars)
+    else:
+        state._mark_dirty_computed_vars()
+        delta_vars = state.dirty_vars & state._frontend_var_names
 
     always_dirty_computed_vars = state._always_dirty_computed_vars
     # Where to leave the values this traversal sends, for whoever delivers

reflex/state.py

@@ -90,6 +90,7 @@
     Delta,
     DeltaMapping,
     _resolve_delta,
+    _send_fields_only,
     build_delta,
     clean_state,
     resolve_delta,
@@ -2439,10 +2440,12 @@ async def hydrate_and_load(
         storage reset, the browser-provided client storage values and the
         state snapshot under one state lock avoids separate load/persist
         cycles of the state tree. The snapshot does not pass through
-        ``get_delta`` and may not write browser storage, so this event's own
-        delta carries the browser-provided values again, as
+        ``get_delta`` and may not write browser storage, so a second delta
+        carries the browser-provided values again, as
         ``update_vars_internal``'s does: ``get_delta`` overrides see them
-        there, and the browser stores what they return. A page with on_load
+        there, and the browser stores what they return. That delta holds only
+        those values, because the snapshot already sent the computed vars
+        derived from them. A page with on_load
         handlers then gets ``on_load_internal``, which only locks its leaf
         substate.
 
@@ -2478,11 +2481,19 @@ async def hydrate_and_load(
             await ctx.emit_delta(delta=delta)
             # Follow-up corrections must be allowed to write browser storage.
             self.dirty_vars.discard(constants.CompileVars.IS_HYDRATED)
-        # The browser's values only: the reset defaults stay clean, so they are
-        # not written back to the browser.
-        for var_state, var_name in applied:
-            var_state.dirty_vars.add(var_name)
-            var_state._mark_dirty((var_name,))
+            if applied:
+                # Pass the browser's values through get_delta overrides. The
+                # snapshot already sent their dependents, computed from them.
+                for var_state, var_name in applied:
+                    var_state.dirty_vars.add(var_name)
+                    var_state._mark_ancestors_dirty()
+                token = _send_fields_only.set(True)
+                try:
+                    delta = await self._get_resolved_delta()
+                finally:
+                    _send_fields_only.reset(token)
+                    self._clean()
+                await ctx.emit_delta(delta=delta)
         if not RegistrationContext.get().app.get_load_events(self.rx_router_url.path):
             self.is_hydrated = True
             return None

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

@@ -1,6 +1,7 @@
 """Tests for BaseStateEventProcessor, specifically the _rehydrate path."""
 
 import asyncio
+import collections
 import contextlib
 import dataclasses
 import gc
@@ -2247,6 +2248,152 @@ def load(self):
     assert not any("unset" + FIELD_MARKER in subdelta for subdelta in persisted)
 
 
+@pytest.mark.parametrize(
+    "processor_state_manager", ["in_process", "redis"], indirect=True
+)
+async def test_hydrate_computes_client_storage_dependents_once(
+    wired_app: App,
+    real_base_state_processor: BaseStateEventProcessor,
+    emitted_deltas: list[tuple[str, Mapping[str, Mapping[str, Any]]]],
+    token: str,
+):
+    """Computed vars depending on browser storage are evaluated once per hydrate.
+
+    The snapshot already computes them from the browser's values, so passing
+    those values through get_delta again must not recompute them.
+
+    Args:
+        wired_app: The app wired to the state manager.
+        real_base_state_processor: The event processor.
+        emitted_deltas: Captured deltas.
+        token: The client token.
+    """
+    calls: collections.Counter[str] = collections.Counter()
+
+    class StorageState(State):
+        value: str = rx.LocalStorage("")
+
+        @rx.var(cache=True)
+        def upper(self) -> str:
+            calls["upper"] += 1
+            return self.value.upper()
+
+        @rx.var(cache=False)
+        def length(self) -> int:
+            calls["length"] += 1
+            return len(self.value)
+
+    class DependentState(State):
+        @rx.var(cache=True)
+        async def copy(self) -> str:
+            calls["copy"] += 1
+            return (await self.get_state(StorageState)).value
+
+    wired_app.add_page(
+        lambda: rx.text(StorageState.upper, StorageState.length, DependentState.copy),
+        route="/",
+    )
+    wired_app._compile_page("index")
+    boot_name = Event.from_event_type(State.hydrate_and_load())[0].name  # pyright: ignore[reportCallIssue]
+    payload = {"vars": {f"{StorageState.get_full_name()}.value{FIELD_MARKER}": "abc"}}
+
+    async def hydrate(payload: dict[str, Any]) -> collections.Counter[str]:
+        """Hydrate and count the computed var evaluations.
+
+        Args:
+            payload: The boot event payload.
+
+        Returns:
+            The evaluations per computed var.
+        """
+        calls.clear()
+        async with real_base_state_processor as processor:
+            await (
+                await processor.enqueue(token, _boot_event(boot_name, payload))
+            ).wait_all()
+        return calls.copy()
+
+    without_storage = await hydrate({})
+    with_storage = await hydrate(payload)
+
+    assert with_storage["upper"] == 1
+    assert with_storage["copy"] == 1
+    assert with_storage["length"] == without_storage["length"]
+    sent = {
+        key: value
+        for _, delta in emitted_deltas
+        for key, value in {
+            **delta.get(StorageState.get_full_name(), {}),
+            **delta.get(DependentState.get_full_name(), {}),
+        }.items()
+    }
+    assert sent["upper" + FIELD_MARKER] == "ABC"
+    assert sent["copy" + FIELD_MARKER] == "abc"
+
+
+async def test_hydrate_cleans_client_storage_when_get_delta_raises(
+    wired_app: App,
+    real_base_state_processor: BaseStateEventProcessor,
+    emitted_deltas: list[tuple[str, Mapping[str, Mapping[str, Any]]]],
+    token: str,
+):
+    """A get_delta override raising on browser storage leaves no dirty fields behind.
+
+    Args:
+        wired_app: The app wired to the state manager.
+        real_base_state_processor: The event processor.
+        emitted_deltas: Captured deltas.
+        token: The client token.
+    """
+    key = "value" + FIELD_MARKER
+    seen: list[str] = []
+    handled: list[Exception] = []
+
+    class StorageState(State):
+        value: str = rx.LocalStorage("")
+
+        @_override_base_method
+        def get_delta(self) -> Delta:
+            """Raise the first time the browser's value passes through.
+
+            Returns:
+                The delta.
+
+            Raises:
+                RuntimeError: On the first browser value.
+            """
+            delta = super().get_delta()
+            if (value := delta.get(self.get_full_name(), {}).get(key)) is not None:
+                seen.append(value)
+                if len(seen) == 1:
+                    msg = "boom"
+                    raise RuntimeError(msg)
+            return delta
+
+        @event
+        def noop(self):
+            """Change nothing."""
+
+    real_base_state_processor.backend_exception_handler = handled.append
+    wired_app.add_page(lambda: rx.text(StorageState.value), route="/")
+    wired_app._compile_page("index")
+    boot_name = Event.from_event_type(State.hydrate_and_load())[0].name  # pyright: ignore[reportCallIssue]
+    payload = {"vars": {f"{StorageState.get_full_name()}.{key}": "abc"}}
+
+    async with real_base_state_processor as processor:
+        await processor.enqueue(token, _boot_event(boot_name, payload))
+        await processor.join(5)
+        emitted_deltas.clear()
+        await _send(processor, token, Event.from_event_type(StorageState.noop())[0])
+
+    assert [type(ex) for ex in handled] == [RuntimeError]
+    assert seen == ["abc"]
+    assert not any(
+        key in delta.get(StorageState.get_full_name(), {})
+        for _, delta in emitted_deltas
+    )
+
+
 @pytest.mark.parametrize(
     "processor_state_manager", ["in_process", "redis"], indirect=True
 )