From 2301a72741b5787822f7d0e9a421519607c93ae8 Mon Sep 17 00:00:00 2001 From: Masen Furer Date: Wed, 14 Oct 2026 18:16:28 -0700 Subject: [PATCH 1/3] fix(state): deliver hydration override mutation dependents Preserve fields-only replay output while resolving dependents invalidated by real assignments. Commit uncached delivery records against the final delta and retain shared-state cleanup bookkeeping. Fixes #7571 --- news/+hydrate-override-assignments.bugfix.md | 1 + reflex/istate/delta.py | 26 +- reflex/state.py | 108 ++++- tests/units/istate/test_delta.py | 55 ++- tests/units/istate/test_shared.py | 58 ++- .../processor/test_base_state_processor.py | 445 ++++++++++++++++++ 6 files changed, 679 insertions(+), 14 deletions(-) create mode 100644 news/+hydrate-override-assignments.bugfix.md diff --git a/news/+hydrate-override-assignments.bugfix.md b/news/+hydrate-override-assignments.bugfix.md new file mode 100644 index 00000000000..ebd88425e79 --- /dev/null +++ b/news/+hydrate-override-assignments.bugfix.md @@ -0,0 +1 @@ +Deliver cached computed values invalidated by real state assignments in client-storage hydration overrides. Output-only storage rewrites remain separate from backend state. diff --git a/reflex/istate/delta.py b/reflex/istate/delta.py index a1a7e0eb712..bc83d5ba2e1 100644 --- a/reflex/istate/delta.py +++ b/reflex/istate/delta.py @@ -69,6 +69,26 @@ class _DeltaRecord(NamedTuple): ) +_deferred_delta_records: ContextVar[list[_DeltaRecord] | None] = ContextVar( + "_deferred_delta_records", default=None +) + + +@contextlib.contextmanager +def _defer_delta_recording() -> Iterator[list[_DeltaRecord]]: + """Collect delivery records until several deltas have been merged. + + Yields: + Records to commit against the final outgoing delta. + """ + pending: list[_DeltaRecord] = [] + token = _deferred_delta_records.set(pending) + try: + yield pending + finally: + _deferred_delta_records.reset(token) + + @contextlib.contextmanager def _suppress_delta_recording() -> Iterator[None]: """Stop delta values built in this block from counting as sent to the client. @@ -314,7 +334,11 @@ async def resolve_delta(state: BaseState) -> Delta: finally: _pending_delta_records.reset(records_token) if pending: - _commit_delta_records(pending, delta) + deferred = _deferred_delta_records.get() + if deferred is None: + _commit_delta_records(pending, delta) + else: + deferred.extend(pending) return delta diff --git a/reflex/state.py b/reflex/state.py index d27a57df0ea..2d67fcb0473 100644 --- a/reflex/state.py +++ b/reflex/state.py @@ -89,6 +89,8 @@ from reflex.istate.delta import ( Delta, DeltaMapping, + _commit_delta_records, + _defer_delta_recording, _resolve_delta, _send_fields_only, build_delta, @@ -2482,17 +2484,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 +2496,100 @@ async def hydrate_and_load( return [OnLoadInternalState.on_load_internal] +def _dirty_states(state: BaseState) -> Iterator[BaseState]: + """Yield a state and its dirty descendants. + + Args: + state: The root of the dirty traversal. + + Yields: + States whose dirty fields may need delivery or cleanup. + """ + yield state + for name in state.dirty_substates | state._always_dirty_substates: + yield from _dirty_states(state.substates[name]) + + +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 any newly invalidated dependent values. + """ + # Snapshot computed values were already delivered. Preserve their cleanup + # marks without letting them masquerade as mutations during the replay. + saved = [(state, state.dirty_vars) for state in _dirty_states(root)] + try: + replay = [] + for state, dirty_vars in saved: + state.dirty_vars = dirty_vars.intersection( + state._frontend_var_names + ).difference(state.computed_vars) + replay.extend((state, name) for name in state.dirty_vars) + replay.extend(applied) + before = { + (state, name): _serialize_var(state.get_value(name)) + for state, name in replay + } + for state, name in applied: + state.dirty_vars.add(name) + state._mark_ancestors_dirty() + with _defer_delta_recording() as records: + token = _send_fields_only.set(True) + try: + delta = await root._get_resolved_delta() + finally: + _send_fields_only.reset(token) + # Only discard synthetic replay marks. A returned delta may deliberately + # rewrite or withhold a field without changing its backend value. + for (state, name), value in before.items(): + if _serialize_var(state.get_value(name)) == value: + state.dirty_vars.discard(name) + # Fields-only traversal never evaluates computed vars. An override can + # still insert a placeholder for one, invalidating its previous record. + for state in _dirty_states(root): + subdelta = delta.get(state.get_full_name(), {}) + for name in state._always_dirty_computed_vars: + if name + FIELD_MARKER in subdelta: + attr = state.computed_vars[name]._last_delta_key_attr + if hasattr(state, attr): + delattr(state, attr) + state._was_touched = True + if any(state.dirty_vars for state in _dirty_states(root)): + # The replay's output remains authoritative, even when an override + # rewrote or withheld a field while also making a real assignment. + replay_keys: dict[str, set[str]] = {} + for field in before: + state, name = field + replay_keys.setdefault(state.get_full_name(), set()).add( + name + FIELD_MARKER + ) + for name, changes in (await root._get_resolved_delta()).items(): + if name in replay_keys and name not in delta: + continue + subdelta = delta.setdefault(name, {}) + for key, value in changes.items(): + if key not in replay_keys.get(name, ()): + subdelta.setdefault(key, value) + _commit_delta_records(records, delta) + return delta + finally: + # SharedState captures its fanout dirtiness in _clean: keep a single + # cleanup point and include fields consumed by either replay pass. + for state, dirty_vars in saved: + state.dirty_vars.update(dirty_vars) + for state, name in applied: + state.dirty_vars.add(name) + state._mark_ancestors_dirty() + root._clean() + + T = TypeVar("T", bound=BaseState) diff --git a/tests/units/istate/test_delta.py b/tests/units/istate/test_delta.py index e844182fd20..360ccfd703e 100644 --- a/tests/units/istate/test_delta.py +++ b/tests/units/istate/test_delta.py @@ -1,8 +1,20 @@ """Tests for building state deltas.""" +import asyncio from typing import Any -from reflex.istate.delta import _DROP_FROM_DELTA, _resolve_delta +import pytest + +import reflex as rx +from reflex.istate.delta import ( + _DROP_FROM_DELTA, + _commit_delta_records, + _defer_delta_recording, + _deferred_delta_records, + _resolve_delta, + _suppress_delta_recording, +) +from reflex.state import BaseState async def _coro(value: Any) -> Any: # noqa: RUF029 - a trivial coroutine value for the delta @@ -38,3 +50,44 @@ async def test_resolve_delta_pops_subdict_when_all_keys_drop(): } resolved = await _resolve_delta(delta) assert resolved == {"s2": {"keep": 1}} + + +@pytest.mark.parametrize("exception", [RuntimeError, asyncio.CancelledError]) +def test_deferred_records_reset_after_failure(exception: type[BaseException]): + """A failed or cancelled merged delta cannot leak its delivery collector. + + Args: + exception: The failure that interrupts the deferred scope. + """ + with _defer_delta_recording() as outer: + with pytest.raises(exception), _defer_delta_recording() as inner: + assert _deferred_delta_records.get() is inner + raise exception + assert _deferred_delta_records.get() is outer + assert _deferred_delta_records.get() is None + + +async def test_deferred_records_respect_suppression(): + """Suppressed traversals stay unrecorded inside a deferred delivery scope.""" + + class RecordingState(BaseState): + @rx.var(cache=False) + def value(self) -> str: + """Return an uncached value. + + Returns: + The test value. + """ + return "value" + + state = RecordingState(_reflex_internal_init=True) # pyright: ignore[reportCallIssue] + attr = RecordingState.computed_vars["value"]._last_delta_key_attr + with _defer_delta_recording() as pending: + with _suppress_delta_recording(): + await state._get_resolved_delta() + assert not pending + delta = await state._get_resolved_delta() + assert pending + assert not hasattr(state, attr) + _commit_delta_records(pending, delta) + assert hasattr(state, attr) diff --git a/tests/units/istate/test_shared.py b/tests/units/istate/test_shared.py index e7cb2e73010..30363a6c4a9 100644 --- a/tests/units/istate/test_shared.py +++ b/tests/units/istate/test_shared.py @@ -6,6 +6,7 @@ from unittest.mock import AsyncMock, Mock, patch import pytest +from reflex_base.registry import RegistrationContext import reflex as rx from reflex.istate.shared import ( @@ -13,7 +14,12 @@ _do_update_other_tokens, _patch_state, ) -from reflex.state import BaseState, State +from reflex.state import ( + BaseState, + State, + _get_client_storage_delta, + _override_base_method, +) from reflex.utils.token_manager import ( LocalTokenManager, RedisTokenManager, @@ -173,3 +179,53 @@ async def test_patch_state_recomputes_readers_after_restoring(): async with _patch_state(original_state=original, linked_state=linked): assert await reader.greeting == "linked" # pyright: ignore[reportAttributeAccessIssue] assert await reader.greeting == "private" # pyright: ignore[reportAttributeAccessIssue] + + +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._previous_dirty_vars >= {"snapshot_write", "value"} + assert not storage.dirty_vars 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 7504c27e669..00d94099352 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 @@ -2527,3 +2527,448 @@ async def slow_load(self): async with _read_back(real_base_state_processor, token) as root: assert (await root.get_state(SlowLoadState)).seen != "stale-finished" assert (await root.get_state(State)).is_hydrated is True + + +@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() + + 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 var_name == "value" and assignment == "after" and client_value == "stale" + 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_delivers_nonstorage_override_assignment( + wired_app: App, + real_base_state_processor: BaseStateEventProcessor, + emitted_deltas: list[tuple[str, Mapping[str, Mapping[str, Any]]]], + token: str, +): + """An override's real non-storage writes also reach their cached dependents. + + Args: + wired_app: The configured app. + real_base_state_processor: The event processor. + emitted_deltas: Captured client updates. + token: The client token. + """ + + class StorageState(State): + value: str = rx.LocalStorage("") + status: str = "initial" + + @_override_base_method + def get_delta(self) -> Delta: + """Assign a separate backend field after building the storage delta. + + Returns: + The unmodified base delta. + """ + delta = super().get_delta() + if "value" in self.dirty_vars and self.status == "initial": + self.status = "reconciled" + return delta + + class DependentState(State): + @rx.var(cache=True, auto_deps=False) + async def status(self) -> str: + """Read the separate changed field. + + Returns: + The backend reconciliation status. + """ + return (await self.get_state(StorageState)).status + + DependentState.computed_vars["status"].add_dependency( + DependentState, + StorageState.status, # pyright: ignore[reportArgumentType] + ) + wired_app.add_page( + lambda: rx.text(StorageState.value, DependentState.status), route="/" + ) + wired_app._compile_page("index") + boot_name = Event.from_event_type(State.hydrate_and_load())[0].name + payload = { + "vars": {f"{StorageState.get_full_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() + values = [ + delta[DependentState.get_full_name()]["status" + FIELD_MARKER] + for _, delta in emitted_deltas + if "status" + FIELD_MARKER in delta.get(DependentState.get_full_name(), {}) + ] + assert values[-1] == "reconciled" + + +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("constant", [False, True]) +@pytest.mark.parametrize("mode", ["placeholder", "drop_state"]) +async def test_hydrate_records_only_delivered_computed_values( + wired_app: App, + real_base_state_processor: BaseStateEventProcessor, + emitted_deltas: list[tuple[str, Mapping[str, Mapping[str, Any]]]], + token: str, + constant: bool, + mode: str, +): + """A filtered correction must not record an unseen computed value as delivered. + + Args: + wired_app: The configured app. + real_base_state_processor: The event processor. + emitted_deltas: Captured client updates. + token: The client token. + constant: Whether the computed value matches its snapshot value. + mode: Whether the override replaces a value or hides the whole state. + """ + computed_key = "computed" + FIELD_MARKER + + class StorageState(State): + value: str = rx.LocalStorage("") + status: str = "initial" + + @rx.var(cache=False) + def computed(self) -> str: + """Return a value whose delivery history must match the final delta. + + Returns: + The current value, possibly unchanged from the snapshot. + """ + return "constant" if constant else self.status + + @_override_base_method + def get_delta(self) -> Delta: + """Make a real assignment while filtering the first storage replay. + + Returns: + The explicitly filtered delta. + """ + first = "value" in self.dirty_vars and self.status == "initial" + if first: + self.status = "changed" + delta = super().get_delta() + if first: + if mode == "placeholder": + delta.setdefault(self.get_full_name(), {})[computed_key] = ( + "placeholder" + ) + else: + delta.pop(self.get_full_name(), None) + return delta + + @event + def noop(self): + """Allow a later ordinary delta to deliver the actual computed value.""" + + name = StorageState.get_full_name() + wired_app.add_page( + lambda: rx.text(StorageState.value, StorageState.computed), 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 ( + await processor.enqueue( + token, + _boot_event( + boot_name, {"vars": {f"{name}.value{FIELD_MARKER}": "browser"}} + ), + ) + ).wait_all() + await _send(processor, token, Event.from_event_type(StorageState.noop())[0]) + values = [ + delta[name][computed_key] + for _, delta in emitted_deltas + if computed_key in delta.get(name, {}) + ] + if mode == "placeholder": + assert "placeholder" in values + assert values[-1] == ("constant" if constant else "changed") From 2efe6e46e9114d612b15af4863706adfb22f5fbc Mon Sep 17 00:00:00 2001 From: Masen Furer Date: Wed, 14 Oct 2026 18:33:39 -0700 Subject: [PATCH 2/3] fix(state): replay hydration overrides only once Mark cached storage dependents for delivery without invalidating their snapshot caches, then use one public get_delta traversal. Preserve output filters and existing synchronous/after-super ordering semantics without repeating override side effects or merging separately filtered deltas. --- reflex/istate/delta.py | 40 +---- reflex/state.py | 111 +++--------- tests/units/istate/test_delta.py | 55 +----- .../processor/test_base_state_processor.py | 161 ++++++------------ 4 files changed, 85 insertions(+), 282 deletions(-) diff --git a/reflex/istate/delta.py b/reflex/istate/delta.py index bc83d5ba2e1..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): @@ -69,26 +71,6 @@ class _DeltaRecord(NamedTuple): ) -_deferred_delta_records: ContextVar[list[_DeltaRecord] | None] = ContextVar( - "_deferred_delta_records", default=None -) - - -@contextlib.contextmanager -def _defer_delta_recording() -> Iterator[list[_DeltaRecord]]: - """Collect delivery records until several deltas have been merged. - - Yields: - Records to commit against the final outgoing delta. - """ - pending: list[_DeltaRecord] = [] - token = _deferred_delta_records.set(pending) - try: - yield pending - finally: - _deferred_delta_records.reset(token) - - @contextlib.contextmanager def _suppress_delta_recording() -> Iterator[None]: """Stop delta values built in this block from counting as sent to the client. @@ -253,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 @@ -334,11 +316,7 @@ async def resolve_delta(state: BaseState) -> Delta: finally: _pending_delta_records.reset(records_token) if pending: - deferred = _deferred_delta_records.get() - if deferred is None: - _commit_delta_records(pending, delta) - else: - deferred.extend(pending) + _commit_delta_records(pending, delta) return delta diff --git a/reflex/state.py b/reflex/state.py index 2d67fcb0473..6cb2d0d0c6a 100644 --- a/reflex/state.py +++ b/reflex/state.py @@ -89,10 +89,8 @@ from reflex.istate.delta import ( Delta, DeltaMapping, - _commit_delta_records, - _defer_delta_recording, + _replay_client_storage, _resolve_delta, - _send_fields_only, build_delta, clean_state, resolve_delta, @@ -2445,9 +2443,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. @@ -2496,20 +2494,6 @@ async def hydrate_and_load( return [OnLoadInternalState.on_load_internal] -def _dirty_states(state: BaseState) -> Iterator[BaseState]: - """Yield a state and its dirty descendants. - - Args: - state: The root of the dirty traversal. - - Yields: - States whose dirty fields may need delivery or cleanup. - """ - yield state - for name in state.dirty_substates | state._always_dirty_substates: - yield from _dirty_states(state.substates[name]) - - async def _get_client_storage_delta( root: BaseState, applied: list[tuple[BaseState, str]] ) -> Delta: @@ -2520,73 +2504,32 @@ async def _get_client_storage_delta( applied: Browser storage fields applied by this hydration. Returns: - Override-filtered storage and any newly invalidated dependent values. + Override-filtered storage and cached dependent values. """ - # Snapshot computed values were already delivered. Preserve their cleanup - # marks without letting them masquerade as mutations during the replay. - saved = [(state, state.dirty_vars) for state in _dirty_states(root)] + # 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. + pending = list(applied) + seen: set[tuple[str, str]] = set() + while pending: + state, name = pending.pop() + for dependent in state._var_dependencies.get(name, ()): + if dependent in seen: + continue + seen.add(dependent) + state_name, cvar_name = dependent + target = root.get_substate(state_name.split(".")) + target.dirty_vars.add(cvar_name) + target._mark_ancestors_dirty() + pending.append((target, cvar_name)) + for state, name in applied: + state.dirty_vars.add(name) + state._mark_ancestors_dirty() + token = _replay_client_storage.set(True) try: - replay = [] - for state, dirty_vars in saved: - state.dirty_vars = dirty_vars.intersection( - state._frontend_var_names - ).difference(state.computed_vars) - replay.extend((state, name) for name in state.dirty_vars) - replay.extend(applied) - before = { - (state, name): _serialize_var(state.get_value(name)) - for state, name in replay - } - for state, name in applied: - state.dirty_vars.add(name) - state._mark_ancestors_dirty() - with _defer_delta_recording() as records: - token = _send_fields_only.set(True) - try: - delta = await root._get_resolved_delta() - finally: - _send_fields_only.reset(token) - # Only discard synthetic replay marks. A returned delta may deliberately - # rewrite or withhold a field without changing its backend value. - for (state, name), value in before.items(): - if _serialize_var(state.get_value(name)) == value: - state.dirty_vars.discard(name) - # Fields-only traversal never evaluates computed vars. An override can - # still insert a placeholder for one, invalidating its previous record. - for state in _dirty_states(root): - subdelta = delta.get(state.get_full_name(), {}) - for name in state._always_dirty_computed_vars: - if name + FIELD_MARKER in subdelta: - attr = state.computed_vars[name]._last_delta_key_attr - if hasattr(state, attr): - delattr(state, attr) - state._was_touched = True - if any(state.dirty_vars for state in _dirty_states(root)): - # The replay's output remains authoritative, even when an override - # rewrote or withheld a field while also making a real assignment. - replay_keys: dict[str, set[str]] = {} - for field in before: - state, name = field - replay_keys.setdefault(state.get_full_name(), set()).add( - name + FIELD_MARKER - ) - for name, changes in (await root._get_resolved_delta()).items(): - if name in replay_keys and name not in delta: - continue - subdelta = delta.setdefault(name, {}) - for key, value in changes.items(): - if key not in replay_keys.get(name, ()): - subdelta.setdefault(key, value) - _commit_delta_records(records, delta) - return delta + return await root._get_resolved_delta() finally: - # SharedState captures its fanout dirtiness in _clean: keep a single - # cleanup point and include fields consumed by either replay pass. - for state, dirty_vars in saved: - state.dirty_vars.update(dirty_vars) - for state, name in applied: - state.dirty_vars.add(name) - state._mark_ancestors_dirty() + _replay_client_storage.reset(token) root._clean() diff --git a/tests/units/istate/test_delta.py b/tests/units/istate/test_delta.py index 360ccfd703e..e844182fd20 100644 --- a/tests/units/istate/test_delta.py +++ b/tests/units/istate/test_delta.py @@ -1,20 +1,8 @@ """Tests for building state deltas.""" -import asyncio from typing import Any -import pytest - -import reflex as rx -from reflex.istate.delta import ( - _DROP_FROM_DELTA, - _commit_delta_records, - _defer_delta_recording, - _deferred_delta_records, - _resolve_delta, - _suppress_delta_recording, -) -from reflex.state import BaseState +from reflex.istate.delta import _DROP_FROM_DELTA, _resolve_delta async def _coro(value: Any) -> Any: # noqa: RUF029 - a trivial coroutine value for the delta @@ -50,44 +38,3 @@ async def test_resolve_delta_pops_subdict_when_all_keys_drop(): } resolved = await _resolve_delta(delta) assert resolved == {"s2": {"keep": 1}} - - -@pytest.mark.parametrize("exception", [RuntimeError, asyncio.CancelledError]) -def test_deferred_records_reset_after_failure(exception: type[BaseException]): - """A failed or cancelled merged delta cannot leak its delivery collector. - - Args: - exception: The failure that interrupts the deferred scope. - """ - with _defer_delta_recording() as outer: - with pytest.raises(exception), _defer_delta_recording() as inner: - assert _deferred_delta_records.get() is inner - raise exception - assert _deferred_delta_records.get() is outer - assert _deferred_delta_records.get() is None - - -async def test_deferred_records_respect_suppression(): - """Suppressed traversals stay unrecorded inside a deferred delivery scope.""" - - class RecordingState(BaseState): - @rx.var(cache=False) - def value(self) -> str: - """Return an uncached value. - - Returns: - The test value. - """ - return "value" - - state = RecordingState(_reflex_internal_init=True) # pyright: ignore[reportCallIssue] - attr = RecordingState.computed_vars["value"]._last_delta_key_attr - with _defer_delta_recording() as pending: - with _suppress_delta_recording(): - await state._get_resolved_delta() - assert not pending - delta = await state._get_resolved_delta() - assert pending - assert not hasattr(state, attr) - _commit_delta_records(pending, delta) - assert hasattr(state, attr) 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 00d94099352..e892eae1842 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 @@ -2640,6 +2640,8 @@ async def sibling(self) -> str: 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"), @@ -2652,7 +2654,11 @@ async def sibling(self) -> str: ] expected = ( "stale" - if var_name == "value" and assignment == "after" and client_value == "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 @@ -2740,71 +2746,6 @@ def get_delta(self) -> Delta: ) -async def test_hydrate_delivers_nonstorage_override_assignment( - wired_app: App, - real_base_state_processor: BaseStateEventProcessor, - emitted_deltas: list[tuple[str, Mapping[str, Mapping[str, Any]]]], - token: str, -): - """An override's real non-storage writes also reach their cached dependents. - - Args: - wired_app: The configured app. - real_base_state_processor: The event processor. - emitted_deltas: Captured client updates. - token: The client token. - """ - - class StorageState(State): - value: str = rx.LocalStorage("") - status: str = "initial" - - @_override_base_method - def get_delta(self) -> Delta: - """Assign a separate backend field after building the storage delta. - - Returns: - The unmodified base delta. - """ - delta = super().get_delta() - if "value" in self.dirty_vars and self.status == "initial": - self.status = "reconciled" - return delta - - class DependentState(State): - @rx.var(cache=True, auto_deps=False) - async def status(self) -> str: - """Read the separate changed field. - - Returns: - The backend reconciliation status. - """ - return (await self.get_state(StorageState)).status - - DependentState.computed_vars["status"].add_dependency( - DependentState, - StorageState.status, # pyright: ignore[reportArgumentType] - ) - wired_app.add_page( - lambda: rx.text(StorageState.value, DependentState.status), route="/" - ) - wired_app._compile_page("index") - boot_name = Event.from_event_type(State.hydrate_and_load())[0].name - payload = { - "vars": {f"{StorageState.get_full_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() - values = [ - delta[DependentState.get_full_name()]["status" + FIELD_MARKER] - for _, delta in emitted_deltas - if "status" + FIELD_MARKER in delta.get(DependentState.get_full_name(), {}) - ] - assert values[-1] == "reconciled" - - async def test_hydrate_cleans_storage_when_correction_emit_fails( wired_app: App, real_base_state_processor: BaseStateEventProcessor, @@ -2889,86 +2830,80 @@ async def failing_emit( # noqa: RUF029 assert not any(key in delta.get(name, {}) for _, delta in emitted_deltas) -@pytest.mark.parametrize("constant", [False, True]) -@pytest.mark.parametrize("mode", ["placeholder", "drop_state"]) -async def test_hydrate_records_only_delivered_computed_values( +@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, - constant: bool, - mode: str, + hide_status: bool, ): - """A filtered correction must not record an unseen computed value as delivered. + """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. - constant: Whether the computed value matches its snapshot value. - mode: Whether the override replaces a value or hides the whole state. + hide_status: Whether the override omits its newly assigned status. """ - computed_key = "computed" + FIELD_MARKER class StorageState(State): value: str = rx.LocalStorage("") + revision: int = 0 status: str = "initial" - @rx.var(cache=False) - def computed(self) -> str: - """Return a value whose delivery history must match the final delta. + @rx.var(cache=True) + def upper(self) -> str: + """Return the cached normalized value. Returns: - The current value, possibly unchanged from the snapshot. + The uppercase storage value. """ - return "constant" if constant else self.status + return self.value.upper() @_override_base_method def get_delta(self) -> Delta: - """Make a real assignment while filtering the first storage replay. + """Normalize storage and retain the override's field filter. Returns: - The explicitly filtered delta. + The filtered delta. """ - first = "value" in self.dirty_vars and self.status == "initial" - if first: - self.status = "changed" + 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 first: - if mode == "placeholder": - delta.setdefault(self.get_full_name(), {})[computed_key] = ( - "placeholder" - ) - else: - delta.pop(self.get_full_name(), None) + 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 - @event - def noop(self): - """Allow a later ordinary delta to deliver the actual computed value.""" - - name = StorageState.get_full_name() wired_app.add_page( - lambda: rx.text(StorageState.value, StorageState.computed), route="/" + 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, {"vars": {f"{name}.value{FIELD_MARKER}": "browser"}} - ), - ) + await processor.enqueue(token, _boot_event(boot_name, payload)) ).wait_all() - await _send(processor, token, Event.from_event_type(StorageState.noop())[0]) - values = [ - delta[name][computed_key] - for _, delta in emitted_deltas - if computed_key in delta.get(name, {}) - ] - if mode == "placeholder": - assert "placeholder" in values - assert values[-1] == ("constant" if constant else "changed") + 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 From 7c8a05ca40b2317a827fb2350c857867494106bd Mon Sep 17 00:00:00 2001 From: Masen Furer Date: Fri, 23 Oct 2026 16:48:43 +0000 Subject: [PATCH 3/3] fix(state): share the dependent walk between invalidation and hydration replay Fold the delivery-only marking into _mark_dirty_computed_vars behind an invalidate flag, so the hydration replay reuses the one dependency walk instead of carrying a second copy of it. Move the shared-cleanup test next to the module it covers and assert that the override's correction took effect. Drop the news fragment: #7568 is unreleased, so this fix changes nothing for any released version. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_014CNQyjqBL56phnEuga6U6E --- news/+hydrate-override-assignments.bugfix.md | 1 - reflex/state.py | 32 +++++------ tests/units/istate/test_shared.py | 58 +------------------- tests/units/test_state.py | 53 ++++++++++++++++++ 4 files changed, 69 insertions(+), 75 deletions(-) delete mode 100644 news/+hydrate-override-assignments.bugfix.md diff --git a/news/+hydrate-override-assignments.bugfix.md b/news/+hydrate-override-assignments.bugfix.md deleted file mode 100644 index ebd88425e79..00000000000 --- a/news/+hydrate-override-assignments.bugfix.md +++ /dev/null @@ -1 +0,0 @@ -Deliver cached computed values invalidated by real state assignments in client-storage hydration overrides. Output-only storage rewrites remain separate from backend state. diff --git a/reflex/state.py b/reflex/state.py index 8abe2f9cdde..1487ed14e06 100644 --- a/reflex/state.py +++ b/reflex/state.py @@ -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)) @@ -2509,20 +2519,8 @@ async def _get_client_storage_delta( # 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. - pending = list(applied) - seen: set[tuple[str, str]] = set() - while pending: - state, name = pending.pop() - for dependent in state._var_dependencies.get(name, ()): - if dependent in seen: - continue - seen.add(dependent) - state_name, cvar_name = dependent - target = root.get_substate(state_name.split(".")) - target.dirty_vars.add(cvar_name) - target._mark_ancestors_dirty() - pending.append((target, cvar_name)) 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) diff --git a/tests/units/istate/test_shared.py b/tests/units/istate/test_shared.py index 30363a6c4a9..e7cb2e73010 100644 --- a/tests/units/istate/test_shared.py +++ b/tests/units/istate/test_shared.py @@ -6,7 +6,6 @@ from unittest.mock import AsyncMock, Mock, patch import pytest -from reflex_base.registry import RegistrationContext import reflex as rx from reflex.istate.shared import ( @@ -14,12 +13,7 @@ _do_update_other_tokens, _patch_state, ) -from reflex.state import ( - BaseState, - State, - _get_client_storage_delta, - _override_base_method, -) +from reflex.state import BaseState, State from reflex.utils.token_manager import ( LocalTokenManager, RedisTokenManager, @@ -179,53 +173,3 @@ async def test_patch_state_recomputes_readers_after_restoring(): async with _patch_state(original_state=original, linked_state=linked): assert await reader.greeting == "linked" # pyright: ignore[reportAttributeAccessIssue] assert await reader.greeting == "private" # pyright: ignore[reportAttributeAccessIssue] - - -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._previous_dirty_vars >= {"snapshot_write", "value"} - assert not storage.dirty_vars 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