fix: sync shared-name storage vars in the writing tab (#7567)

news/+sync-same-tab-shared-name.bugfix.md

@@ -0,0 +1 @@
+Synced `LocalStorage` vars that share one storage `name` now update in the tab that writes, not only in other tabs.

reflex/state.py

@@ -1112,6 +1112,7 @@ def _init_var_dependency_dicts(cls):
         # change to any state invalidates every cached root.
         cls._to_schema.cache_clear()
         _frontend_schema.cache_clear()
+        _shared_storage_groups.cache_clear()
 
     @classmethod
     def _iter_functions(cls) -> Iterator[tuple[str, FunctionType]]:
@@ -1953,6 +1954,8 @@ async def _get_resolved_delta(self) -> Delta:
         Returns:
             The resolved delta.
         """
+        if not _replay_client_storage.get():
+            await _sync_shared_storage(self)
         return await resolve_delta(self)
 
     def _mark_dirty(self, var_names: Iterable[str] | None = None) -> None:
@@ -2283,6 +2286,92 @@ def _deserialize(
         return state
 
 
+@functools.lru_cache
+def _shared_storage_groups(
+    root: type[BaseState], context: RegistrationContext
+) -> tuple[tuple[tuple[type[BaseState], str, bool], ...], ...]:
+    """Group the localStorage vars by storage name.
+
+    Only names that a synced var shares with another var are included.
+
+    Args:
+        root: The state at the top of the tree.
+        context: The registration context holding the tree.
+
+    Returns:
+        The state, var name and sync flag of each var, one group per name.
+    """
+    # The compiler imports this module.
+    from reflex.compiler.utils import _compile_client_storage_recursive
+
+    by_name: dict[str, list[tuple[type[BaseState], str, bool]]] = {}
+    for state_key, options in _compile_client_storage_recursive(root)[1].items():
+        if name := options.get("name"):
+            state_name, _, var_name = state_key.removesuffix(FIELD_MARKER).rpartition(
+                "."
+            )
+            by_name.setdefault(name, []).append((
+                root.get_class_substate(state_name),
+                var_name,
+                bool(options.get("sync")),
+            ))
+    # Sorted, so conflicting writes to one name set the other vars the same way
+    # in every process.
+    return tuple(
+        tuple(sorted(members, key=lambda m: (m[0].get_full_name(), m[1])))
+        for members in by_name.values()
+        if len(members) > 1 and any(sync for _, _, sync in members)
+    )
+
+
+async def _sync_shared_storage(state: BaseState) -> None:
+    """Set the synced vars sharing the storage name of a var this event wrote.
+
+    The browser raises the storage event in other tabs only, so the tab that
+    stores the value would keep the other vars with its name stale.
+
+    Args:
+        state: The root of the states to get the delta for.
+    """
+    for members in _shared_storage_groups(
+        state.get_root_state(), RegistrationContext.get()
+    ):
+        written = {
+            (state_cls, var_name): source.get_value(var_name)
+            for state_cls, var_name, _ in members
+            if (source := _loaded_substate(state, state_cls.get_full_name()))
+            is not None
+            and var_name in source.dirty_vars
+        }
+        if not written:
+            continue
+        # The vars this event wrote keep their values; the others take the first
+        # written value in name order.
+        value = next(iter(written.values()))
+        for state_cls, var_name, sync in members:
+            if sync and (state_cls, var_name) not in written:
+                target = await state.get_state(state_cls)
+                if target.get_value(var_name) != value:
+                    setattr(target, var_name, value)
+
+
+def _loaded_substate(root: BaseState, state_name: str) -> BaseState | None:
+    """Get a substate by its full name only if it is already loaded.
+
+    Args:
+        root: The root state.
+        state_name: The full name of the substate.
+
+    Returns:
+        The substate, or None if it is not loaded.
+    """
+    state = root
+    for name in state_name.split(".")[1:]:
+        if (state := state.substates.get(name)) is None:
+            return None
+    return state
+
+
 @functools.lru_cache
 def _frontend_schema(root: type[BaseState], context: RegistrationContext) -> str:
     """Hash the frontend schema of the state tree below ``root``.

tests/integration/tests_playwright/test_hydration_storage.py

@@ -165,6 +165,7 @@ def get_delta(self):
 
     class SharedNameState(rx.State):
         synced: str = rx.LocalStorage("", name="hydrate-shared", sync=True)
+        synced2: str = rx.LocalStorage("", name="hydrate-shared", sync=True)
         plain: str = rx.LocalStorage("", name="hydrate-shared")
 
         @rx.event
@@ -176,6 +177,41 @@ def set_plain(self, value: str):
             """
             self.plain = value
 
+        @rx.event
+        def set_synced(self, value: str):
+            """Store a value through the synced var.
+
+            Args:
+                value: The value to store.
+            """
+            self.synced = value
+
+    class SharedOverrideState(rx.State):
+        first: str = rx.LocalStorage("", name="hydrate-shared-override", sync=True)
+        second: str = rx.LocalStorage("", name="hydrate-shared-override", sync=True)
+
+        @rx.event
+        def set_second(self, value: str):
+            """Store a value through the second var.
+
+            Args:
+                value: The value to store.
+            """
+            self.second = value
+
+        @rx.state._override_base_method
+        def get_delta(self):
+            """Replace the first var's "x".
+
+            Returns:
+                The delta, with the value replaced.
+            """
+            delta = super().get_delta()
+            subdelta = delta.get(self.get_full_name(), {})
+            if subdelta.get("first" + FIELD_MARKER) == "x":
+                subdelta["first" + FIELD_MARKER] = "x-replaced"
+            return delta
+
     def synced():
         """Display a browser storage var synced across tabs.
 
@@ -190,6 +226,20 @@ def synced():
             rx.text(ReplaceState.plain, id="replace-plain"),
             rx.button("Keep", on_click=ReplaceState.keep_sent, id="keep-sent"),
             rx.text(SharedNameState.synced, id="shared-value"),
+            rx.text(SharedNameState.synced2, id="shared-value-2"),
+            rx.button(
+                "Synced new",
+                on_click=SharedNameState.set_synced("new"),
+                id="set-synced-new",
+            ),
+            rx.text(SharedOverrideState.first, id="override-first"),
+            rx.text(SharedOverrideState.second, id="override-second"),
+            rx.button(
+                "Second x", on_click=SharedOverrideState.set_second("x"), id="set-x"
+            ),
+            rx.button(
+                "Second y", on_click=SharedOverrideState.set_second("y"), id="set-y"
+            ),
             rx.button(
                 "Plain old",
                 on_click=SharedNameState.set_plain("old"),
@@ -664,6 +714,47 @@ def test_synced_storage_echo_after_own_write_is_written(
     assert page.evaluate("key => localStorage.getItem(key)", key) == "sent"
 
 
+def test_synced_storage_shared_name_syncs_siblings_in_writing_tab(
+    hydration_storage_app: AppHarness, page: Page
+):
+    """A write to one var sharing a storage name updates the others in the writing tab.
+
+    Args:
+        hydration_storage_app: The running app.
+        page: A fresh browser page.
+    """
+    socket = EventSocket(page)
+    open_synced(hydration_storage_app, page)
+    page.locator("#set-synced-new").click()
+    expect(page.locator("#shared-value")).to_have_text("new")
+    expect(page.locator("#shared-value-2")).to_have_text("new")
+    page.wait_for_timeout(500)
+    assert page.evaluate("localStorage.getItem('hydrate-shared')") == "new"
+    assert not [m for m in socket.sent if "update_vars_internal" in m]
+
+
+def test_synced_storage_shared_name_override_keeps_syncing(
+    hydration_storage_app: AppHarness, page: Page
+):
+    """A get_delta override replacing one shared-name value leaves later writes synced.
+
+    Args:
+        hydration_storage_app: The running app.
+        page: A fresh browser page.
+    """
+    socket = EventSocket(page)
+    open_synced(hydration_storage_app, page)
+    page.locator("#set-x").click()
+    expect(page.locator("#override-first")).to_have_text("x-replaced")
+    expect(page.locator("#override-second")).to_have_text("x")
+    page.locator("#set-y").click()
+    expect(page.locator("#override-first")).to_have_text("y")
+    expect(page.locator("#override-second")).to_have_text("y")
+    page.wait_for_timeout(500)
+    assert page.evaluate("localStorage.getItem('hydrate-shared-override')") == "y"
+    assert not [m for m in socket.sent if "update_vars_internal" in m]
+
+
 @pytest.mark.parametrize("other_tab_stores", [False, True], ids=["alone", "other_tab"])
 @pytest.mark.parametrize("synced", [True, False], ids=["synced", "plain"])
 def test_replaced_storage_echo_does_not_hide_later_change(

tests/units/test_state.py

@@ -7539,3 +7539,110 @@ def _clean(self):
     assert storage.value == "fresh"
     assert storage._previous_dirty_vars >= {"snapshot_write", "value"}
     assert not storage.dirty_vars
+
+
+async def test_shared_storage_name_syncs_siblings(
+    clean_registration_context: RegistrationContext,
+):
+    """A write to a var sharing a storage name sets the synced vars with that name.
+
+    Args:
+        clean_registration_context: An isolated state registry.
+    """
+
+    class SharedFirst(State):
+        value: str = rx.LocalStorage("", name="shared", sync=True)
+
+    class SharedSecond(State):
+        value: str = rx.LocalStorage("", name="shared", sync=True)
+
+    class SharedPlain(State):
+        value: str = rx.LocalStorage("", name="shared")
+
+    class OtherName(State):
+        value: str = rx.LocalStorage("", name="other", sync=True)
+
+    key = "value" + FIELD_MARKER
+    root = State()
+    first = await root.get_state(SharedFirst)
+    first.value = "new"
+    delta = await root._get_resolved_delta()
+    root._clean()
+    assert delta[SharedFirst.get_full_name()][key] == "new"
+    assert delta[SharedSecond.get_full_name()][key] == "new"
+    assert SharedPlain.get_full_name() not in delta
+    assert OtherName.get_full_name() not in delta
+
+    # An unsynced var's write raises the storage event too, so it syncs them.
+    plain = await root.get_state(SharedPlain)
+    plain.value = "plain"
+    delta = await root._get_resolved_delta()
+    root._clean()
+    assert delta[SharedFirst.get_full_name()][key] == "plain"
+    assert delta[SharedSecond.get_full_name()][key] == "plain"
+
+    # An unchanged sibling is not sent again.
+    first.value = "plain"
+    assert await root._get_resolved_delta() == {
+        SharedFirst.get_full_name(): {key: "plain"}
+    }
+    root._clean()
+
+    # A name no other var shares sets nothing else.
+    other = await root.get_state(OtherName)
+    other.value = "alone"
+    assert await root._get_resolved_delta() == {
+        OtherName.get_full_name(): {key: "alone"}
+    }
+
+
+@pytest.mark.parametrize("second_sync", [True, False])
+async def test_shared_storage_name_keeps_event_writes(
+    clean_registration_context: RegistrationContext,
+    second_sync: bool,
+):
+    """Vars sharing a storage name that one event writes keep their own values.
+
+    Args:
+        clean_registration_context: An isolated state registry.
+        second_sync: Whether the second var syncs.
+    """
+
+    class WrittenFirst(State):
+        value: str = rx.LocalStorage("", name="shared", sync=True)
+
+    class WrittenSecond(State):
+        value: str = rx.LocalStorage("", name="shared", sync=second_sync)
+
+    root = State()
+    first = await root.get_state(WrittenFirst)
+    second = await root.get_state(WrittenSecond)
+    first.value = "one"
+    second.value = "two"
+    await root._get_resolved_delta()
+    assert (first.value, second.value) == ("one", "two")
+
+
+async def test_shared_storage_name_conflicting_writes_pick_by_name(
+    clean_registration_context: RegistrationContext,
+):
+    """Conflicting writes to one storage name set the other synced vars in name order.
+
+    Args:
+        clean_registration_context: An isolated state registry.
+    """
+
+    class ZWriter(State):
+        value: str = rx.LocalStorage("", name="shared", sync=True)
+
+    class AWriter(State):
+        value: str = rx.LocalStorage("", name="shared", sync=True)
+
+    class MSibling(State):
+        value: str = rx.LocalStorage("", name="shared", sync=True)
+
+    root = State()
+    (await root.get_state(ZWriter)).value = "z"
+    (await root.get_state(AWriter)).value = "a"
+    await root._get_resolved_delta()
+    assert (await root.get_state(MSibling)).value == "a"