diff --git a/news/+hydrate-storage-dependents.bugfix.md b/news/+hydrate-storage-dependents.bugfix.md new file mode 100644 index 00000000000..bd6d7246b2f --- /dev/null +++ b/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. diff --git a/reflex/istate/delta.py b/reflex/istate/delta.py index d312f52e99b..a1a7e0eb712 100644 --- a/reflex/istate/delta.py +++ b/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 diff --git a/reflex/state.py b/reflex/state.py index f064b7c9d07..d27a57df0ea 100644 --- a/reflex/state.py +++ b/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 diff --git a/tests/units/reflex_base/event/processor/test_base_state_processor.py b/tests/units/reflex_base/event/processor/test_base_state_processor.py index 11c6516912c..7504c27e669 100644 --- a/tests/units/reflex_base/event/processor/test_base_state_processor.py +++ b/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 )