ENG-8258: API for linking and sharing states (#6024)

* ENG-8258: API for linking and sharing states

It works by defining a substate of SharedState and then calling
self._link_to(target_token) from some event handler.  from that point on,
whenever that user's state is loaded, the StateManager will patch in the linked
shared states.  whenever a linked state is modified, we explicitly load all of
the other linked tokens, patch in the modified states, and send a delta to
those clients

You can call ._unlink to remove the link association, which causes the substate
to be subsequently loaded from the client_token's tree as a private state

It is intended to work transparently with computed vars, background events, and
frontend rendering.

* WiP: move _links to root state to avoid extra redis load

* Fix test_mixin_state

* only import ShareStateBaseInternal if there are links present

* Fix delta calculation in certain scenarios

* Add integration test for linked states

* fix test_computed_var assertion

* remove background task deadlock updating linked state

perform the subsequent updates in an asyncio.Task to allow the original caller
to drop the lock for the other shared states.

* do not reuse _modify_linked_states in `_link_to`

the state might have multiple links, and we may have already entered the
context and are holding the lock already, so we don't want to take the lock
again which will hang.

instead, check if the state is already linked to the target token and avoid
doing extra work.

* abstract out _patch_state and use it when linking or unlinking

* _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.

* fix test cases

* mark router and router_data as dirty when patching state

definitely the token will be different. although private-dependent data should
use private states, it's common for llm generated code to define router_data
dependent vars in the linked state itself, so we make that special case work

* AI CR feedback

Author: masenf

pyi_hashes.json

@@ -1,5 +1,5 @@
 {
-  "reflex/__init__.pyi": "b304ed6f7a2fa028a194cad81bd83112",
+  "reflex/__init__.pyi": "0a3ae880e256b9fd3b960e12a2cb51a7",
   "reflex/components/__init__.pyi": "ac05995852baa81062ba3d18fbc489fb",
   "reflex/components/base/__init__.pyi": "16e47bf19e0d62835a605baa3d039c5a",
   "reflex/components/base/app_wrap.pyi": "22e94feaa9fe675bcae51c412f5b67f1",

reflex/__init__.py

@@ -336,6 +336,7 @@
         "State",
         "dynamic",
     ],
+    "istate.shared": ["SharedState"],
     "istate.wrappers": ["get_state"],
     "style": ["Style", "toggle_color_mode"],
     "utils.imports": ["ImportDict", "ImportVar"],

reflex/app.py

@@ -1562,13 +1562,17 @@ def all_routes(_request: Request) -> Response:
 
     @contextlib.asynccontextmanager
     async def modify_state(
-        self, token: str, background: bool = False
+        self,
+        token: str,
+        background: bool = False,
+        previous_dirty_vars: dict[str, set[str]] | None = None,
     ) -> AsyncIterator[BaseState]:
         """Modify the state out of band.
 
         Args:
             token: The token to modify the state for.
             background: Whether the modification is happening in a background task.
+            previous_dirty_vars: Vars that are considered dirty from a previous operation.
 
         Yields:
             The state to modify.
@@ -1581,7 +1585,9 @@ async def modify_state(
             raise RuntimeError(msg)
 
         # Get exclusive access to the state.
-        async with self.state_manager.modify_state(token) as state:
+        async with self.state_manager.modify_state_with_links(
+            token, previous_dirty_vars=previous_dirty_vars
+        ) as state:
             # No other event handler can modify the state while in this context.
             yield state
             delta = await state._get_resolved_delta()
@@ -1769,7 +1775,7 @@ async def process(
             constants.RouteVar.CLIENT_IP: client_ip,
         })
         # Get the state for the session exclusively.
-        async with app.state_manager.modify_state(
+        async with app.state_manager.modify_state_with_links(
             event.substate_token, event=event
         ) as state:
             # When this is a brand new instance of the state, signal the
@@ -2003,7 +2009,9 @@ async def _ndjson_updates():
                 Each state update as JSON followed by a new line.
             """
             # Process the event.
-            async with app.state_manager.modify_state(event.substate_token) as state:
+            async with app.state_manager.modify_state_with_links(
+                event.substate_token
+            ) as state:
                 async for update in state._process(event):
                     # Postprocess the event.
                     update = await app._postprocess(state, event, update)

reflex/istate/manager/__init__.py

@@ -114,6 +114,35 @@ async def modify_state(
         """
         yield self.state()
 
+    @contextlib.asynccontextmanager
+    async def modify_state_with_links(
+        self,
+        token: str,
+        previous_dirty_vars: dict[str, set[str]] | None = None,
+        **context: Unpack[StateModificationContext],
+    ) -> AsyncIterator[BaseState]:
+        """Modify the state for a token, including linked substates, while holding exclusive lock.
+
+        Args:
+            token: The token to modify the state for.
+            previous_dirty_vars: The previously dirty vars for linked states.
+            context: The state modification context.
+
+        Yields:
+            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) is not None:
+                from reflex.istate.shared import SharedStateBaseInternal
+
+                shared_state = await root_state.get_state(SharedStateBaseInternal)
+                async with shared_state._modify_linked_states(
+                    previous_dirty_vars=previous_dirty_vars
+                ) as _:
+                    yield root_state
+            else:
+                yield root_state
+
     async def close(self):  # noqa: B027
         """Close the state manager."""