_link_to returns the newly linked state

_modify_linked_states context can now release the locks of newly linked states
and send updates for changes in newly linked states.

rehydrating after linking is no longer necessary.

Author: masenf

reflex/istate/manager/__init__.py

@@ -132,7 +132,7 @@ async def modify_state_with_links(
             The state for the token with linked states patched in.
         """
         async with self.modify_state(token, **context) as root_state:
-            if getattr(root_state, "_reflex_internal_links", None):
+            if getattr(root_state, "_reflex_internal_links", None) is not None:
                 from reflex.istate.shared import SharedStateBaseInternal
 
                 shared_state = await root_state.get_state(SharedStateBaseInternal)

reflex/istate/shared.py

@@ -3,12 +3,14 @@
 import asyncio
 import contextlib
 from collections.abc import AsyncIterator
+from typing import Self, TypeVar
 
 from reflex.event import Event, get_hydrate_event
 from reflex.state import BaseState, State, _override_base_method, _substate_key
 from reflex.utils.exceptions import ReflexRuntimeError
 
 UPDATE_OTHER_CLIENT_TASKS: set[asyncio.Task] = set()
+LINKED_STATE = TypeVar("LINKED_STATE", bound="SharedStateBaseInternal")
 
 
 def _do_update_other_tokens(
@@ -92,15 +94,19 @@ async def _patch_state(
 class SharedStateBaseInternal(State):
     """The private base state for all shared states."""
 
+    _exit_stack: contextlib.AsyncExitStack | None = None
+    _held_locks: dict[str, dict[type[BaseState], BaseState]] = {}
+
     def __getstate__(self):
         """Override redis serialization to remove temporary fields.
 
         Returns:
             The state dictionary without temporary fields.
         """
         s = super().__getstate__()
-        # Don't want to persist the cached substates
         s.pop("_previous_dirty_vars", None)
+        s.pop("_exit_stack", None)
+        s.pop("_held_locks", None)
         return s
 
     @_override_base_method
@@ -124,6 +130,8 @@ def _mark_dirty(self):
         state to be considered dirty either.
         """
         self.dirty_vars.discard("_previous_dirty_vars")
+        self.dirty_vars.discard("_exit_stack")
+        self.dirty_vars.discard("_held_locks")
         # Only mark dirty if there are still dirty vars, or any substate is dirty
         if self.dirty_vars or any(
             substate.dirty_vars for substate in self.substates.values()
@@ -144,7 +152,7 @@ def _rehydrate(self):
             State.set_is_hydrated(True),
         ]
 
-    async def _link_to(self, token: str):
+    async def _link_to(self, token: str) -> Self:
         """Link this shared state to a token.
 
         After linking, subsequent access to this shared state will affect the
@@ -155,15 +163,16 @@ async def _link_to(self, token: str):
             token: The token to link to.
 
         Returns:
-            The events to rehydrate the state after linking (these should be returned/yielded).
-        """
-        from reflex.istate.manager import get_state_manager
+            The newly linked state.
 
+        Raises:
+            ReflexRuntimeError: If linking fails or token is invalid.
+        """
         if not token:
             msg = "Cannot link shared state to empty token."
             raise ReflexRuntimeError(msg)
         if self._linked_to == token:
-            return None  # already linked to this token
+            return self  # already linked to this token
         if self._linked_to and self._linked_to != token:
             # Disassociate from previous linked token since unlink will not be called.
             self._linked_from.discard(self.router.session.client_token)
@@ -174,21 +183,10 @@ async def _link_to(self, token: str):
 
         # Associate substate with the given link token.
         state_name = self.get_full_name()
+        if self._reflex_internal_links is None:
+            self._reflex_internal_links = {}
         self._reflex_internal_links[state_name] = token
-
-        # Get the newly linked state and update pointers/delta for subsequent events.
-        async with get_state_manager().modify_state(
-            _substate_key(token, type(self))
-        ) as linked_root_state:
-            linked_state = await linked_root_state.get_state(type(self))
-            linked_state._linked_from.add(self.router.session.client_token)
-            linked_state._linked_to = token
-            async with _patch_state(
-                original_state=self,
-                linked_state=linked_state,
-                full_delta=True,
-            ):
-                return self._rehydrate()
+        return await self._internal_patch_linked_state(token, full_delta=True)
 
     async def _unlink(self):
         """Unlink this shared state from its linked token.
@@ -199,7 +197,10 @@ async def _unlink(self):
         from reflex.istate.manager import get_state_manager
 
         state_name = self.get_full_name()
-        if state_name not in self._reflex_internal_links:
+        if (
+            not self._reflex_internal_links
+            or state_name not in self._reflex_internal_links
+        ):
             msg = f"State {state_name} is not linked and cannot be unlinked."
             raise ReflexRuntimeError(msg)
 
@@ -219,6 +220,68 @@ async def _unlink(self):
         ):
             return self._rehydrate()
 
+    async def _internal_patch_linked_state(
+        self, token: str, full_delta: bool = False
+    ) -> Self:
+        """Load and replace this state with the linked state for a given token.
+
+        Must be called inside a `_modify_linked_states` context, to ensure locks are
+        released after the event is done processing.
+
+        Args:
+            token: The token of the linked state.
+            full_delta: If True, mark all Vars in linked_state dirty and resolve
+                delta to update cached computed vars
+
+        Returns:
+            The state that was linked into the tree.
+        """
+        from reflex.istate.manager import get_state_manager
+
+        if self._exit_stack is None:
+            msg = "Cannot link shared state outside of _modify_linked_states context."
+            raise ReflexRuntimeError(msg)
+
+        # Get the newly linked state and update pointers/delta for subsequent events.
+        if token not in self._held_locks:
+            linked_root_state = await self._exit_stack.enter_async_context(
+                get_state_manager().modify_state(_substate_key(token, type(self)))
+            )
+            self._held_locks.setdefault(token, {})
+        else:
+            linked_root_state = await get_state_manager().get_state(
+                _substate_key(token, type(self))
+            )
+        linked_state = await linked_root_state.get_state(type(self))
+        # Avoid unnecessary dirtiness of shared state when there are no changes.
+        if type(self) not in self._held_locks[token]:
+            self._held_locks[token][type(self)] = linked_state
+        if self.router.session.client_token not in linked_state._linked_from:
+            linked_state._linked_from.add(self.router.session.client_token)
+        if linked_state._linked_to != token:
+            linked_state._linked_to = token
+        await self._exit_stack.enter_async_context(
+            _patch_state(
+                original_state=self,
+                linked_state=linked_state,
+                full_delta=full_delta,
+            )
+        )
+        return linked_state
+
+    def _held_locks_linked_states(self) -> list["SharedState"]:
+        """Get all linked states currently held by this state.
+
+        Returns:
+            The list of linked states currently held.
+        """
+        return [
+            linked_state
+            for linked_state_cls_to_instance in self._held_locks.values()
+            for linked_state in linked_state_cls_to_instance.values()
+            if isinstance(linked_state, SharedState)
+        ]
+
     @contextlib.asynccontextmanager
     async def _modify_linked_states(
         self, previous_dirty_vars: dict[str, set[str]] | None = None
@@ -236,67 +299,54 @@ async def _modify_linked_states(
         Yields:
             None.
         """
-        from reflex.istate.manager import get_state_manager
-
-        exit_stack = contextlib.AsyncExitStack()
-        held_locks: set[str] = set()
-        linked_states: list[BaseState] = []
+        if self._exit_stack is not None:
+            msg = "Cannot nest _modify_linked_states contexts."
+            raise ReflexRuntimeError(msg)
+        if self._reflex_internal_links is None:
+            msg = "No linked states to modify."
+            raise ReflexRuntimeError(msg)
+        self._exit_stack = contextlib.AsyncExitStack()
+        self._held_locks = {}
         current_dirty_vars: dict[str, set[str]] = {}
         affected_tokens: set[str] = set()
         # Go through all linked states and patch them in if they are present in the tree
         for linked_state_name, linked_token in self._reflex_internal_links.items():
-            linked_state_cls = self.get_root_state().get_class_substate(
-                linked_state_name
+            linked_state_cls: type[SharedState] = (
+                self.get_root_state().get_class_substate(  # pyright: ignore[reportAssignmentType]
+                    linked_state_name
+                )
             )
             # TODO: Avoid always fetched linked states, it should be based on
             # whether the state is accessed, however then `get_state` would need
             # to know how to fetch in a linked state.
             original_state = await self.get_state(linked_state_cls)
-            if linked_token not in held_locks:
-                linked_root_state = await exit_stack.enter_async_context(
-                    get_state_manager().modify_state(
-                        _substate_key(linked_token, linked_state_cls)
-                    )
-                )
-                held_locks.add(linked_token)
-            else:
-                linked_root_state = await get_state_manager().get_state(
-                    _substate_key(linked_token, linked_state_cls)
-                )
-            linked_state = await linked_root_state.get_state(linked_state_cls)
-            linked_states.append(linked_state)
-            linked_state._linked_to = linked_token
-            linked_state._linked_from.add(self.router.session.client_token)
-            await exit_stack.enter_async_context(
-                _patch_state(original_state, linked_state)
+            linked_state = await original_state._internal_patch_linked_state(
+                linked_token
             )
             if (
                 previous_dirty_vars
                 and (dv := previous_dirty_vars.get(linked_state_name)) is not None
             ):
                 linked_state.dirty_vars.update(dv)
                 linked_state._mark_dirty()
-        async with exit_stack:
+        async with self._exit_stack:
             yield None
             # Collect dirty vars and other affected clients that need to be updated.
-            for linked_state in linked_states:
-                if (
-                    linked_state_previous_dirty_vars := getattr(
-                        linked_state, "_previous_dirty_vars", None
-                    )
-                ) is not None:
+            for linked_state in self._held_locks_linked_states():
+                if linked_state._previous_dirty_vars is not None:
                     current_dirty_vars[linked_state.get_full_name()] = set(
-                        linked_state_previous_dirty_vars
+                        linked_state._previous_dirty_vars
                     )
                 if (
                     linked_state._get_was_touched()
-                    or linked_state_previous_dirty_vars is not None
+                    or linked_state._previous_dirty_vars is not None
                 ):
                     affected_tokens.update(
                         token
                         for token in linked_state._linked_from
                         if token != self.router.session.client_token
                     )
+        self._exit_stack = None
 
         # Only propagate dirty vars when we are not already propagating from another state.
         if previous_dirty_vars is None:
@@ -324,3 +374,6 @@ def __init_subclass__(cls, **kwargs):
         kwargs["mixin"] = False
         cls._mixin = False
         super().__init_subclass__(**kwargs)
+        root_state = cls.get_root_state()
+        if root_state.backend_vars["_reflex_internal_links"] is None:
+            root_state.backend_vars["_reflex_internal_links"] = {}

reflex/state.py

@@ -2465,7 +2465,7 @@ class State(BaseState):
     # The hydrated bool.
     is_hydrated: bool = False
     # Maps the state full_name to an arbitrary token it is linked to for shared state.
-    _reflex_internal_links: dict[str, str] = {}
+    _reflex_internal_links: dict[str, str] | None = None
 
     @event
     def set_is_hydrated(self, value: bool) -> None:

tests/integration/test_linked_state.py

@@ -16,6 +16,7 @@
 
 def LinkedStateApp():
     """Test that linked state works as expected."""
+    import uuid
     from typing import Any
 
     import reflex as rx
@@ -36,23 +37,31 @@ def set_who(self, who: str) -> None:
 
         @rx.event
         async def link_to(self, token: str):
-            return await self._link_to(token)
+            await self._link_to(token)
+
+        @rx.event
+        async def link_to_and_increment(self):
+            linked_state = await self._link_to(f"arbitrary-token-{uuid.uuid4()}")
+            linked_state.counter += 1
 
         @rx.event
         async def unlink(self):
             return await self._unlink()
 
         @rx.event
         async def on_load_link_default(self):
-            return await self._link_to(self.room or "default")
+            linked_state = await self._link_to(self.room or "default")
+            if self.room:
+                assert linked_state._linked_to == self.room
+            else:
+                assert linked_state._linked_to == "default"
 
         @rx.event
         async def handle_submit(self, form_data: dict[str, Any]):
             if "who" in form_data:
                 self.set_who(form_data["who"])
             if "token" in form_data:
-                return await self.link_to(form_data["token"])
-            return None
+                await self.link_to(form_data["token"])
 
     class PrivateState(rx.State):
         @rx.var
@@ -126,6 +135,11 @@ def index() -> rx.Component:
                 on_click=PrivateState.bump_counter_yield,
                 id="yield-button",
             ),
+            rx.button(
+                "Link to arbitrary token and Increment n_changes",
+                on_click=SharedState.link_to_and_increment,
+                id="link-increment-button",
+            ),
         )
 
     app = rx.App()
@@ -355,3 +369,20 @@ def test_linked_state(
         == "Hello, world!"
     )
     assert linked_state.poll_for_content(counter_button_1, exp_not_equal="48") == "0"
+    counter_button_1.click()
+    assert linked_state.poll_for_content(counter_button_1, exp_not_equal="0") == "1"
+    counter_button_1.click()
+    assert linked_state.poll_for_content(counter_button_1, exp_not_equal="1") == "2"
+    counter_button_1.click()
+    assert linked_state.poll_for_content(counter_button_1, exp_not_equal="2") == "3"
+    # Ensure other tabs are unaffected
+    assert n_changes_2.text == "2"
+    assert greeting_2.text == "Hello, Diana!"
+    assert counter_button_2.text == "48"
+    assert n_changes_3.text == "2"
+    assert greeting_3.text == "Hello, Diana!"
+    assert counter_button_3.text == "48"
+
+    # Link to a new state and increment the counter in the same event
+    tab1.find_element(By.ID, "link-increment-button").click()
+    assert linked_state.poll_for_content(counter_button_1, exp_not_equal="3") == "1"