diff --git a/reflex/istate/delta.py b/reflex/istate/delta.py index a1a7e0eb712..7c03df6634e 100644 --- a/reflex/istate/delta.py +++ b/reflex/istate/delta.py @@ -42,10 +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) +# Replay browser storage and cached dependents without blanket invalidation. +# Real assignments invalidate their caches eagerly. Uncached values are left +# to the normal event delta, avoiding an extra evaluation during hydration. +_replay_client_storage: ContextVar[bool] = ContextVar( + "_replay_client_storage", default=False +) class _DeltaRecord(NamedTuple): @@ -233,9 +235,9 @@ def build_delta(state: BaseState) -> Delta: """ delta = {} - if _send_fields_only.get(): + if _replay_client_storage.get(): delta_vars = state.dirty_vars & state._frontend_var_names - delta_vars.difference_update(state.computed_vars) + delta_vars.difference_update(state._always_dirty_computed_vars) else: state._mark_dirty_computed_vars() delta_vars = state.dirty_vars & state._frontend_var_names diff --git a/reflex/state.py b/reflex/state.py index bb7fbca3546..1487ed14e06 100644 --- a/reflex/state.py +++ b/reflex/state.py @@ -89,8 +89,8 @@ from reflex.istate.delta import ( Delta, DeltaMapping, + _replay_client_storage, _resolve_delta, - _send_fields_only, build_delta, clean_state, resolve_delta, @@ -1862,16 +1862,25 @@ async def get_var_value(self, var: Var[VAR_TYPE]) -> VAR_TYPE: return await value return value - def _mark_dirty_computed_vars(self, var_names: Iterable[str] | None = None) -> None: - """Invalidate the computed vars depending on changed vars of this state, transitively. + def _mark_dirty_computed_vars( + self, var_names: Iterable[str] | None = None, *, invalidate: bool = True + ) -> None: + """Mark the computed vars depending on changed vars of this state dirty, transitively. + Invalidates their cached values, so that the next read recomputes them. Uncached and expired computed vars recompute on access, so they count as changed along with any var. Args: var_names: The changed vars of this state; all its dirty vars if omitted. + invalidate: Whether the dependents' caches are invalidated. False marks + them for delivery only, leaving their caches and the uncached and + expired computed vars alone, so that an unchanged value is sent + as it is rather than recomputed. """ - if self._always_dirty_computed_vars or self._interval_computed_var_names: + if invalidate and ( + self._always_dirty_computed_vars or self._interval_computed_var_names + ): recomputed = ( self._always_dirty_computed_vars | self._expired_computed_vars() ) @@ -1900,7 +1909,8 @@ def _mark_dirty_computed_vars(self, var_names: Iterable[str] | None = None) -> N else: target = state._get_root_state().get_substate(state_name.split(".")) target._mark_ancestors_dirty() - target.computed_vars[cvar_name].mark_dirty(instance=target) + if invalidate: + target.computed_vars[cvar_name].mark_dirty(instance=target) target.dirty_vars.add(cvar_name) if cvar_name in target._var_dependencies: pending.append((target, cvar_name)) @@ -2443,9 +2453,9 @@ async def hydrate_and_load( ``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. That delta holds only - those values, because the snapshot already sent the computed vars - derived from them. A page with on_load + there, and the browser stores what they return. Cached dependents are + included without invalidating their snapshot caches; real assignments + in an override invalidate those caches normally. A page with on_load handlers then gets ``on_load_internal``, which only locks its leaf substate. @@ -2482,17 +2492,7 @@ async def hydrate_and_load( # Follow-up corrections must be allowed to write browser storage. self.dirty_vars.discard(constants.CompileVars.IS_HYDRATED) 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() + delta = await _get_client_storage_delta(self, applied) await ctx.emit_delta(delta=delta) if not RegistrationContext.get().app.get_load_events(self.rx_router_url.path): self.is_hydrated = True @@ -2504,6 +2504,33 @@ async def hydrate_and_load( return [OnLoadInternalState.on_load_internal] +async def _get_client_storage_delta( + root: BaseState, applied: list[tuple[BaseState, str]] +) -> Delta: + """Replay browser storage without discarding real override mutations. + + Args: + root: The hydrated root state. + applied: Browser storage fields applied by this hydration. + + Returns: + Override-filtered storage and cached dependent values. + """ + # Include cached storage dependents in the one public traversal, even when + # their branch is visited before the override changes the storage value. + # Mark delivery only: unchanged caches still hold the snapshot's values. + for state, name in applied: + state._mark_dirty_computed_vars((name,), invalidate=False) + state.dirty_vars.add(name) + state._mark_ancestors_dirty() + token = _replay_client_storage.set(True) + try: + return await root._get_resolved_delta() + finally: + _replay_client_storage.reset(token) + root._clean() + + T = TypeVar("T", bound=BaseState) 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 7cc5b9be3d3..efa118506b1 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 @@ -3023,3 +3023,383 @@ async def background(self): ) assert sorted(finished) == ["background", "foreground", "other"] assert not loads + + +@pytest.mark.parametrize("storage", [rx.LocalStorage, rx.SessionStorage, rx.Cookie]) +@pytest.mark.parametrize("assignment", ["before", "after", "after_output"]) +@pytest.mark.parametrize("with_load", [True, False]) +@pytest.mark.parametrize("with_hashes", [True, False]) +@pytest.mark.parametrize("client_value", ["stale", "fresh"]) +@pytest.mark.parametrize( + "processor_state_manager", ["in_process", "redis"], indirect=True +) +async def test_hydrate_delivers_storage_override_assignments( + wired_app: App, + real_base_state_processor: BaseStateEventProcessor, + emitted_deltas: list[tuple[str, Mapping[str, Mapping[str, Any]]]], + token: str, + storage: type[rx.LocalStorage] | type[rx.SessionStorage] | type[rx.Cookie], + assignment: str, + with_load: bool, + with_hashes: bool, + client_value: str, +): + """Deliver real storage assignments and their cached dependents during hydration. + + Args: + wired_app: The configured app. + real_base_state_processor: The event processor. + emitted_deltas: Captured client updates. + token: The client token. + storage: The storage implementation. + assignment: Whether the assignment precedes or follows delta creation. + with_load: Whether the page has an on-load handler. + with_hashes: Whether compiled defaults accompany hydration. + client_value: The browser value. + """ + calls = collections.Counter() + + class StorageState(State): + value: str = storage("") + unset: str = rx.Cookie("unset-default") + _reconciled: bool = False + + @_override_base_method + def get_delta(self) -> Delta: + """Normalize a stale value once, before or after the base traversal. + + Returns: + The unmodified base delta. + """ + correct = ( + "value" in self.dirty_vars + and self.value == "stale" + and not self._reconciled + ) + if correct and assignment == "before": + self.value = "fresh" + self._reconciled = True + delta = super().get_delta() + if correct and assignment != "before": + self.value = "fresh" + self._reconciled = True + if assignment == "after_output": + delta[self.get_full_name()]["value" + FIELD_MARKER] = self.value + return delta + + @rx.var(cache=True) + def same_state(self) -> str: + """Read the normalized storage value in the same state. + + Returns: + The current backend storage value. + """ + calls["same"] += 1 + return self.value + + @event + def load(self): + """Provide an on-load handler.""" + + class DependentState(State): + @rx.var(cache=True, auto_deps=False) + async def sibling(self) -> str: + """Read the normalized storage value from another branch. + + Returns: + The current backend storage value. + """ + calls["sibling"] += 1 + return (await self.get_state(StorageState)).value + + DependentState.computed_vars["sibling"].add_dependency( + DependentState, + StorageState.value, # pyright: ignore[reportArgumentType] + ) + wired_app.add_page( + lambda: rx.text( + StorageState.value, StorageState.same_state, DependentState.sibling + ), + route="/", + on_load=StorageState.load if with_load else None, + ) + wired_app._compile_page("index") + name = StorageState.get_full_name() + key = "value" + FIELD_MARKER + payload: dict[str, Any] = {"vars": {f"{name}.{key}": client_value}} + if with_hashes: + payload["hashes"] = state_snapshot_hashes(compile_state(State)) + calls.clear() + boot_name = Event.from_event_type(State.hydrate_and_load())[0].name + async with real_base_state_processor as processor: + await ( + await processor.enqueue(token, _boot_event(boot_name, payload)) + ).wait_all() + + # Preserve the public return contract: a synchronous value collected before + # an after-super assignment is not rewritten after the override returns. + for state_name, var_name in ( + (name, "value"), + (name, "same_state"), + (DependentState.get_full_name(), "sibling"), + ): + values = [ + delta[state_name][var_name + FIELD_MARKER] + for _, delta in emitted_deltas + if var_name + FIELD_MARKER in delta.get(state_name, {}) + ] + expected = ( + "stale" + if client_value == "stale" + and ( + (var_name == "value" and assignment == "after") + or (var_name == "same_state" and assignment != "before") + ) + else "fresh" + ) + assert values[-1] == expected + assert not any( + "unset" + FIELD_MARKER in delta.get(name, {}) for _, delta in emitted_deltas[1:] + ) + if client_value == "fresh": + assert calls == {"same": 1, "sibling": 1} + async with _read_back(real_base_state_processor, token) as root: + stored = await root.get_state(StorageState) + dependent = await root.get_state(DependentState) + assert stored.value == stored.same_state == "fresh" + assert await dependent.sibling == "fresh" + + +@pytest.mark.parametrize("mode", ["rewrite", "drop_key", "drop_state"]) +@pytest.mark.parametrize("mutate", [False, True]) +async def test_hydrate_keeps_override_output_filters( + wired_app: App, + real_base_state_processor: BaseStateEventProcessor, + emitted_deltas: list[tuple[str, Mapping[str, Mapping[str, Any]]]], + token: str, + mode: str, + mutate: bool, +): + """Keep outgoing filters separate from actual backend assignments. + + Args: + wired_app: The configured app. + real_base_state_processor: The event processor. + emitted_deltas: Captured client updates. + token: The client token. + mode: How the override filters the output. + mutate: Whether the override also assigns a backend field. + """ + seen = [] + key = "value" + FIELD_MARKER + + class StorageState(State): + value: str = rx.LocalStorage("") + + @_override_base_method + def get_delta(self) -> Delta: + """Filter a storage value without copying the filtered value into state. + + Returns: + The filtered delta. + """ + filter_output = "value" in self.dirty_vars and self.value == "stale" + if mutate and filter_output: + self.value = "fresh" + delta = super().get_delta() + subdelta = delta.get(self.get_full_name(), {}) + if key in subdelta: + seen.append(subdelta[key]) + if key in subdelta and filter_output: + if mode == "rewrite": + subdelta[key] = "public-placeholder" + elif mode == "drop_key": + subdelta.pop(key) + else: + delta.pop(self.get_full_name()) + return delta + + wired_app.add_page(lambda: rx.text(StorageState.value), route="/") + wired_app._compile_page("index") + name = StorageState.get_full_name() + boot_name = Event.from_event_type(State.hydrate_and_load())[0].name + payload = {"vars": {f"{name}.{key}": "stale"}} + async with real_base_state_processor as processor: + await ( + await processor.enqueue(token, _boot_event(boot_name, payload)) + ).wait_all() + outgoing = [ + delta[name][key] + for _, delta in emitted_deltas[1:] + if key in delta.get(name, {}) + ] + assert outgoing == (["public-placeholder"] if mode == "rewrite" else []) + if not mutate: + assert seen == ["stale"] + async with _read_back(real_base_state_processor, token) as root: + assert (await root.get_state(StorageState)).value == ( + "fresh" if mutate else "stale" + ) + + +async def test_hydrate_cleans_storage_when_correction_emit_fails( + wired_app: App, + real_base_state_processor: BaseStateEventProcessor, + emitted_deltas: list[tuple[str, Mapping[str, Mapping[str, Any]]]], + token: str, +): + """A failed correction send leaves no synthetic dirtiness on the next event. + + Args: + wired_app: The configured app. + real_base_state_processor: The event processor. + emitted_deltas: Captured client updates. + token: The client token. + """ + seen = [] + handled = [] + key = "value" + FIELD_MARKER + + class StorageState(State): + value: str = rx.LocalStorage("") + + @_override_base_method + def get_delta(self) -> Delta: + """Record storage-bearing deltas. + + Returns: + The base delta. + """ + delta = super().get_delta() + if key in delta.get(self.get_full_name(), {}): + seen.append(self.value) + return delta + + @event + def noop(self): + """Leave the state unchanged.""" + + name = StorageState.get_full_name() + failed = False + + async def failing_emit( # noqa: RUF029 + client_token: str, delta: Mapping[str, Mapping[str, Any]] + ) -> None: + """Fail the first persist-eligible storage delta. + + Args: + client_token: The receiving client. + delta: The outgoing update. + + Raises: + RuntimeError: On the first storage correction. + """ + nonlocal failed + guarded = ( + delta.get(State.get_full_name(), {}).get("is_hydrated" + FIELD_MARKER) + is False + ) + if not failed and not guarded and key in delta.get(name, {}): + failed = True + message = "correction send failed" + raise RuntimeError(message) + emitted_deltas.append((client_token, delta)) + + assert real_base_state_processor._root_context is not None + real_base_state_processor._root_context = dataclasses.replace( + real_base_state_processor._root_context, emit_delta_impl=failing_emit + ) + 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 + async with real_base_state_processor as processor: + await processor.enqueue( + token, _boot_event(boot_name, {"vars": {f"{name}.{key}": "browser"}}) + ) + await processor.join(5) + emitted_deltas.clear() + await _send(processor, token, Event.from_event_type(StorageState.noop())[0]) + assert failed + assert [type(error) for error in handled] == [RuntimeError] + assert seen == ["browser"] + assert not any(key in delta.get(name, {}) for _, delta in emitted_deltas) + + +@pytest.mark.parametrize("hide_status", [False, True]) +async def test_hydrate_replays_override_side_effects_once( + wired_app: App, + real_base_state_processor: BaseStateEventProcessor, + emitted_deltas: list[tuple[str, Mapping[str, Mapping[str, Any]]]], + token: str, + hide_status: bool, +): + """Preserve one-shot effects and omissions for fields assigned during replay. + + Args: + wired_app: The configured app. + real_base_state_processor: The event processor. + emitted_deltas: Captured client updates. + token: The client token. + hide_status: Whether the override omits its newly assigned status. + """ + + class StorageState(State): + value: str = rx.LocalStorage("") + revision: int = 0 + status: str = "initial" + + @rx.var(cache=True) + def upper(self) -> str: + """Return the cached normalized value. + + Returns: + The uppercase storage value. + """ + return self.value.upper() + + @_override_base_method + def get_delta(self) -> Delta: + """Normalize storage and retain the override's field filter. + + Returns: + The filtered delta. + """ + dirty = "value" in self.dirty_vars + changed = dirty and self.value != self.value.strip() + if dirty: + self.value = self.value.strip() + self.revision += 1 + if changed: + self.status = "normalized" + delta = super().get_delta() + if hide_status and changed: + delta.get(self.get_full_name(), {}).pop("status" + FIELD_MARKER, None) + delta.get(self.get_full_name(), {}).pop("upper" + FIELD_MARKER, None) + return delta + + wired_app.add_page( + lambda: rx.text(StorageState.value, StorageState.revision, StorageState.upper), + route="/", + ) + wired_app._compile_page("index") + name = StorageState.get_full_name() + boot_name = Event.from_event_type(State.hydrate_and_load())[0].name + payload = {"vars": {f"{name}.value{FIELD_MARKER}": " browser "}} + async with real_base_state_processor as processor: + await ( + await processor.enqueue(token, _boot_event(boot_name, payload)) + ).wait_all() + corrections = [delta.get(name, {}) for _, delta in emitted_deltas[1:]] + assert [ + d["revision" + FIELD_MARKER] + for d in corrections + if "revision" + FIELD_MARKER in d + ] == [1] + if hide_status: + assert not any("status" + FIELD_MARKER in d for d in corrections) + assert not any("upper" + FIELD_MARKER in d for d in corrections) + async with _read_back(real_base_state_processor, token) as root: + state = await root.get_state(StorageState) + assert state.value == "browser" + assert state.revision == 1 diff --git a/tests/units/test_state.py b/tests/units/test_state.py index 791843c4a0a..b2061a26ef9 100644 --- a/tests/units/test_state.py +++ b/tests/units/test_state.py @@ -80,7 +80,9 @@ State, StateUpdate, _frontend_schema, + _get_client_storage_delta, _load_events_for_page, + _override_base_method, all_base_state_classes, is_serializable, state_snapshot_hashes, @@ -7412,3 +7414,54 @@ def test_non_var_annotation_with_undefined_name_is_allowed( state = type("UndefinedNonVarState", (BaseState,), namespace) assert name not in state.get_fields() + + +async def test_hydration_replay_preserves_shared_cleanup_marks( + monkeypatch: pytest.MonkeyPatch, clean_registration_context: RegistrationContext +): + """A single shared cleanup records snapshot and correction writes for fanout. + + Args: + monkeypatch: Restore root metadata changed by defining a shared state. + clean_registration_context: Isolate the temporary shared state class. + """ + links = State.get_fields()["_reflex_internal_links"] + monkeypatch.setattr(links, "default", links.default) + monkeypatch.setattr(links, "default_factory", links.default_factory) + monkeypatch.setattr( + State, "_always_dirty_substates", State._always_dirty_substates.copy() + ) + clean_calls = 0 + + class StorageState(rx.SharedState): + value: str = rx.LocalStorage("") + snapshot_write: str = "initial" + + @_override_base_method + def get_delta(self): + """Normalize a backend storage value before returning its delta. + + Returns: + The base delta. + """ + if "value" in self.dirty_vars and self.value == "stale": + self.value = "fresh" + return super().get_delta() + + @_override_base_method + def _clean(self): + """Count the cleanup calls that capture shared fanout marks.""" + nonlocal clean_calls + clean_calls += 1 + super()._clean() + + root = State.get_root_state()(_reflex_internal_init=True) + storage = await root.get_state(StorageState) + storage.snapshot_write = "snapshot" + storage.value = "stale" + storage._previous_dirty_vars = set() + await _get_client_storage_delta(root, [(storage, "value")]) + assert clean_calls == 1 + assert storage.value == "fresh" + assert storage._previous_dirty_vars >= {"snapshot_write", "value"} + assert not storage.dirty_vars