Fix client-storage writes during boot hydration (#7460)

* Preserve client-storage semantics during boot hydration

* Use normal hydration deltas for snapshot corrections

* Split hydration release notes into bugfix and feature

* Add bugfix documentation for computed var changes

Author: masenf

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

@@ -0,0 +1 @@
+Prevent hydration from persisting client-storage defaults.

news/+hydrate-computed-var-changes.bugfix.md

@@ -0,0 +1 @@
+Deliver changes to client-storage and other state vars made by computed vars during hydration to the browser.

reflex/state.py

@@ -2338,6 +2338,7 @@ async def hydrate_and_load(
         self._reset_client_storage()
         if vars:
             await _apply_client_storage_vars(self, vars)
+        self._clean()
         # The snapshot must carry is_hydrated=False: the frontend skips
         # writing client storage for a delta that is not yet hydrated, and
         # the reset defaults above must not be written back to the browser.
@@ -2347,8 +2348,12 @@ async def hydrate_and_load(
             delta = await _resolve_delta(self.dict())
             if hashes:
                 delta = await _diff_against_initial_state(type(self), delta, hashes)
+            # Include the guard and values changed while resolving the snapshot.
+            for state_name, changes in (await self._get_resolved_delta()).items():
+                delta.setdefault(state_name, {}).update(changes)
             await ctx.emit_delta(delta=delta)
-        self._clean()
+            # Follow-up corrections must be allowed to write browser storage.
+            self.dirty_vars.discard(constants.CompileVars.IS_HYDRATED)
         if not RegistrationContext.get().app.get_load_events(self.rx_router_url.path):
             self.is_hydrated = True
             return None

tests/integration/tests_playwright/test_hydration_storage.py

@@ -0,0 +1,114 @@
+"""Browser regressions for client-storage writes during hydration."""
+
+from collections.abc import Generator
+
+import pytest
+from playwright.sync_api import Page, expect
+
+from reflex.testing import AppHarness, AppHarnessProd
+
+
+def HydrationStorageApp():
+    """Create an app with mismatched defaults and hydration-time storage writes."""
+    import uuid
+
+    import reflex as rx
+
+    class StorageState(rx.State):
+        local: str = rx.LocalStorage("local-default", name="hydrate-local")
+        session: str = rx.SessionStorage("session-default", name="hydrate-session")
+        cookie: str = rx.Cookie("cookie-default", name="hydrate-cookie")
+        nonce: rx.Field[str] = rx.field(default_factory=lambda: uuid.uuid4().hex)
+
+        @rx.var
+        def checked(self) -> str:
+            """Clear rejected values in each storage type.
+
+            Returns:
+                The current values for display.
+            """
+            if self.local == "bad":
+                self.local = ""
+            if self.session == "bad":
+                self.session = ""
+            if self.cookie == "bad":
+                self.cookie = ""
+            return f"{self.local}|{self.session}|{self.cookie}"
+
+        @rx.event
+        def load(self):
+            """Provide an on-load event for the second page."""
+
+    def index():
+        """Display hydration and normalized storage values.
+
+        Returns:
+            The page component.
+        """
+        return rx.box(
+            rx.text(StorageState.checked, id="checked"),
+            rx.text(rx.cond(rx.State.is_hydrated, "true", "false"), id="hydrated"),
+        )
+
+    app = rx.App()
+    app.add_page(index)
+    app.add_page(index, route="/loaded", on_load=StorageState.load)
+
+
+@pytest.fixture(scope="module", params=[AppHarness, AppHarnessProd])
+def hydration_storage_app(
+    request: pytest.FixtureRequest, tmp_path_factory: pytest.TempPathFactory
+) -> Generator[AppHarness, None, None]:
+    """Run the same storage app in development and production.
+
+    Args:
+        request: The harness parameter.
+        tmp_path_factory: The temporary directory factory.
+
+    Yields:
+        The running app harness.
+    """
+    with request.param.create(
+        root=tmp_path_factory.mktemp("hydration_storage"),
+        app_source=HydrationStorageApp,
+    ) as harness:
+        yield harness
+
+
+@pytest.mark.parametrize("route", ["", "loaded"])
+def test_hydration_storage(hydration_storage_app: AppHarness, page: Page, route: str):
+    """Defaults stay absent, and computed-var corrections persist after reload.
+
+    Args:
+        hydration_storage_app: The running app.
+        page: A fresh browser page.
+        route: The page with or without on-load handlers.
+    """
+    assert hydration_storage_app.frontend_url is not None
+    page.goto(f"{hydration_storage_app.frontend_url.rstrip('/')}/{route}")
+    expect(page.locator("#hydrated")).to_have_text("true")
+    expect(page.locator("#checked")).to_have_text(
+        "local-default|session-default|cookie-default"
+    )
+    assert page.evaluate("localStorage.getItem('hydrate-local')") is None
+    assert page.evaluate("sessionStorage.getItem('hydrate-session')") is None
+    assert not any(
+        cookie.get("name") == "hydrate-cookie" for cookie in page.context.cookies()
+    )
+
+    page.evaluate("""() => {
+        localStorage.setItem('hydrate-local', 'bad');
+        sessionStorage.setItem('hydrate-session', 'bad');
+        document.cookie = 'hydrate-cookie=bad; path=/';
+    }""")
+    page.reload()
+    expect(page.locator("#hydrated")).to_have_text("true")
+    expect(page.locator("#checked")).to_have_text("||")
+    page.wait_for_function("""() =>
+        localStorage.getItem('hydrate-local') === '' &&
+        sessionStorage.getItem('hydrate-session') === '' &&
+        document.cookie.split('; ').includes('hydrate-cookie=')
+    """)
+    page.reload()
+    expect(page.locator("#hydrated")).to_have_text("true")
+    expect(page.locator("#checked")).to_have_text("||")

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

@@ -23,14 +23,21 @@
 import reflex as rx
 from reflex import event
 from reflex.app import App
+from reflex.compiler.utils import compile_state
 from reflex.event import Event, EventSpec
 from reflex.istate.manager import StateManager
 from reflex.istate.manager.disk import StateManagerDisk
 from reflex.istate.manager.memory import StateManagerMemory
 from reflex.istate.manager.redis import StateManagerRedis
 from reflex.istate.manager.token import BaseStateToken
 from reflex.middleware.middleware import Middleware
-from reflex.state import BaseState, OnLoadInternalState, State, StateUpdate
+from reflex.state import (
+    BaseState,
+    OnLoadInternalState,
+    State,
+    StateUpdate,
+    state_snapshot_hashes,
+)
 from reflex.utils import types as reflex_types
 from tests.units.conftest import metric_points
 from tests.units.mock_redis import mock_redis
@@ -1630,8 +1637,6 @@ async def test_hydrate_and_load_diffs_against_compiled_defaults(
         emitted_deltas: List to capture emitted deltas.
         token: The client token.
     """
-    from reflex.compiler.utils import compile_state
-    from reflex.state import state_snapshot_hashes
 
     class CookieState(State):
         flavor: str = rx.Cookie("plain")
@@ -1669,7 +1674,12 @@ def set_python_equal_values(self):
     assert set(snapshot) == {state_name, CookieState.get_full_name()}
     assert set(snapshot[state_name]) == {
         var + FIELD_MARKER
-        for var in ("rx_router_page", "rx_router_url", "rx_router_route_id")
+        for var in (
+            "rx_router_page",
+            "rx_router_url",
+            "rx_router_route_id",
+            CompileVars.IS_HYDRATED,
+        )
     }
     assert snapshot[CookieState.get_full_name()] == {
         "flavor" + FIELD_MARKER: "chocolate"
@@ -1718,6 +1728,178 @@ def set_python_equal_values(self):
     assert "loads" + FIELD_MARKER in snapshot[CookieState.get_full_name()]
 
 
+async def test_hydrate_keeps_storage_write_guard_for_mismatched_substate(
+    wired_app: App,
+    real_base_state_processor: BaseStateEventProcessor,
+    emitted_deltas: list[tuple[str, Mapping[str, Mapping[str, Any]]]],
+    token: str,
+):
+    """A full substate snapshot must not persist its client-storage defaults.
+
+    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.
+    """
+
+    class StorageState(State):
+        local: str = rx.LocalStorage("local-default")
+        session: str = rx.SessionStorage("session-default")
+        cookie: str = rx.Cookie("cookie-default")
+
+    wired_app.add_page(lambda: rx.text(StorageState.local), route="/")
+    wired_app._compile_page("index")
+    compiled = compile_state(State)
+    hashes = state_snapshot_hashes(compiled)
+    hashes[sorted(compiled).index(StorageState.get_full_name()) + 1] = "mismatch"
+    boot_name = Event.from_event_type(State.hydrate_and_load())[0].name  # pyright: ignore[reportCallIssue]
+
+    async with real_base_state_processor as processor:
+        await (
+            await processor.enqueue(token, _boot_event(boot_name, {"hashes": hashes}))
+        ).wait_all()
+
+    snapshot = emitted_deltas[0][1]
+    assert (
+        snapshot[State.get_full_name()][CompileVars.IS_HYDRATED + FIELD_MARKER] is False
+    )
+    assert snapshot[StorageState.get_full_name()] == {
+        "local" + FIELD_MARKER: "local-default",
+        "session" + FIELD_MARKER: "session-default",
+        "cookie" + FIELD_MARKER: "cookie-default",
+    }
+
+
+@pytest.mark.parametrize("storage", [rx.LocalStorage, rx.SessionStorage, rx.Cookie])
+@pytest.mark.parametrize("cached", [True, False])
+@pytest.mark.parametrize("asynchronous", [True, False])
+@pytest.mark.parametrize("with_load", [True, False])
+@pytest.mark.parametrize("with_hashes", [True, False])
+@pytest.mark.parametrize(
+    "processor_state_manager", ["in_process", "redis"], indirect=True
+)
+async def test_hydrate_delivers_computed_var_mutations(
+    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],
+    cached: bool,
+    asynchronous: bool,
+    with_load: bool,
+    with_hashes: bool,
+):
+    """Snapshot evaluation preserves storage and plain-var writes for delivery.
+
+    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.
+        storage: The client-storage type.
+        cached: Whether to cache the computed var.
+        asynchronous: Whether its evaluation is asynchronous.
+        with_load: Whether the page has an on-load handler.
+        with_hashes: Whether this is the first boot or a reconnect.
+    """
+
+    class StorageState(State):
+        value: str = storage("")
+        status: str = "initial"
+        loaded_value: str = "initial"
+
+        @rx.var
+        def before_validation(self) -> str:
+            """Expose values captured before the validating computed var runs.
+
+            Returns:
+                The status and storage value.
+            """
+            return f"{self.status}:{self.value}"
+
+        def _sanitize(self) -> str:
+            """Clear rejected storage and update a plain var.
+
+            Returns:
+                The current storage value.
+            """
+            if self.value == "bad":
+                self.value = ""
+                self.status = "cleared"
+            return self.value
+
+        if asynchronous:
+
+            @rx.var(cache=cached)
+            async def checked(self) -> str:  # pyright: ignore[reportRedeclaration]
+                """Validate storage asynchronously.
+
+                Returns:
+                    The current storage value.
+                """
+                await asyncio.sleep(0)
+                return self._sanitize()
+
+        else:
+
+            @rx.var(cache=cached)
+            def checked(self) -> str:
+                """Validate storage.
+
+                Returns:
+                    The current storage value.
+                """
+                return self._sanitize()
+
+        @event
+        def load(self):
+            """Record the normalized value seen by on-load handlers."""
+            self.loaded_value = self.value
+
+    wired_app.add_page(
+        lambda: rx.text(StorageState.checked, StorageState.before_validation),
+        route="/",
+        on_load=StorageState.load if with_load else None,
+    )
+    wired_app._compile_page("index")
+    hashes = state_snapshot_hashes(compile_state(State))
+    name = StorageState.get_full_name()
+    key = "value" + FIELD_MARKER
+    payload: dict[str, Any] = {"vars": {f"{name}.{key}": "bad"}}
+    if with_hashes:
+        payload["hashes"] = hashes
+    boot_name = Event.from_event_type(State.hydrate_and_load())[0].name  # pyright: ignore[reportCallIssue]
+
+    async with real_base_state_processor as processor:
+        await (
+            await processor.enqueue(token, _boot_event(boot_name, payload))
+        ).wait_all()
+
+    snapshot = emitted_deltas[0][1]
+    assert snapshot.get(name, {}).get(key, "") == ""
+    assert snapshot[name]["status" + FIELD_MARKER] == "cleared"
+    assert snapshot[name]["before_validation" + FIELD_MARKER] == "cleared:"
+    assert any(
+        delta.get(name, {}).get(key) == ""
+        and delta.get(State.get_full_name(), {}).get(
+            CompileVars.IS_HYDRATED + FIELD_MARKER
+        )
+        is not False
+        for _, delta in emitted_deltas[1:]
+    )
+    if with_load:
+        assert any(
+            delta.get(name, {}).get("loaded_value" + FIELD_MARKER) == ""
+            for _, delta in emitted_deltas[1:]
+        )
+    async with _read_back(real_base_state_processor, token) as root:
+        stored = await root.get_state(StorageState)
+        assert stored.value == ""
+        assert stored.status == "cleared"
+        assert stored.loaded_value == ("" if with_load else "initial")
+
+
 @pytest.mark.parametrize(
     "processor_state_manager", ["in_process", "redis"], indirect=True
 )