ENG-9350: App.modify_state can directly modify SharedState tokens (#6336)

* ENG-9350: App.modify_state can directly modify SharedState tokens

When a SharedState is modified directly by its shared token (e.g. from an
API route webhook), propagate dirty vars to all clients linked to that
shared state — matching the behavior that already exists when modifying
shared state through a private client token.

- Add _collect_shared_token_updates on SharedStateBaseInternal to detect
  the shared-token case inside _modify_linked_states and propagate to
  linked clients via the existing _do_update_other_tokens mechanism
- Add App.set_contexts / App._set_contexts_internal to centralize pushing
  RegistrationContext and EventContext into the current contextvars scope,
  replacing the old _registration_context_middleware
- App.modify_state now calls set_contexts so out-of-band callers (API
  routes, webhooks) get the necessary contexts automatically
- Integration test: API endpoint modifies two SharedState subclasses by
  shared token, asserts both propagate to two linked browser tabs, and
  verifies normal event handlers still work afterward
- Unit tests for set_contexts covering all combinations of pre-existing /
  absent contexts, no-event-processor, and reset-on-exit behavior

Closes #6335

* deal with race condition due to slower execution in CI environment

set an initial value for the shared state to provide affirmative proof that the
state was linked _before_ the API call was made.

if only one state was linked when the HTTP request was made, then only one
state would be updated by such request resulting in the observed behavior.

* followup from previous commit

but actually set the value on the newly linked state, not the old private
state. derpy easy to make mistake with this _link_to mechanism =/

* Add test case for fetching previously unfetched linked state

Using BaseState.get_state(...) with a linked state that _is_ linked, but not
currently cached should fetch the linked state instance, not the private state
instance.

* Handle `.get_state` when directly modifying a shared token

In an `app.modify_state(...)` context where the passed token is a shared token,
handle the case where `.get_state(...)` is used to fetch _another_ shared state
with the same token as the original shared state token.

Basically if some code is modifying "tokenA" for "Shared1" and calls
`shared_1.get_state(Shared2)`, then the retrieved state will also be associated
with "tokenA", as expected.

* recursively _collect_shared_token_updates

Starting from the SharedStateBaseInternal, recurse into each pre-cached
instance and collect affected tokens and dirty vars.

* Fixturify the new test_app set_contexts test case

* remove superfluous `run_isolated` fixture

just use `isolated_context.run` for simplicity

Author: masenf

reflex/app.py

@@ -16,6 +16,7 @@
 import traceback
 import urllib.parse
 from collections.abc import AsyncIterator, Callable, Coroutine, Mapping, Sequence
+from contextvars import Token
 from datetime import datetime
 from itertools import chain
 from pathlib import Path
@@ -31,6 +32,7 @@
     evaluate_style_namespaces,
 )
 from reflex_base.config import get_config
+from reflex_base.context.base import BaseContext
 from reflex_base.environment import ExecutorType, environment
 from reflex_base.event import (
     _EVENT_FIELDS,
@@ -40,6 +42,7 @@
     IndividualEventType,
     noop,
 )
+from reflex_base.event.context import EventContext
 from reflex_base.event.processor import BaseStateEventProcessor, EventProcessor
 from reflex_base.registry import RegistrationContext
 from reflex_base.utils import console
@@ -577,8 +580,65 @@ async def modified_send(message: Message):
         # Ensure the event processor starts and stops with the server.
         self.register_lifespan_task(self._setup_event_processor)
 
-    def _registration_context_middleware(self, app: ASGIApp) -> ASGIApp:
-        """Ensure the RegistrationContext is attached to the ASGI app.
+    def _set_contexts_internal(self) -> dict[type[BaseContext], Token]:
+        """Set Reflex contexts if not already present, returning reset tokens.
+
+        Returns:
+            A dict mapping context class to the contextvars Token for each
+            context that was set. Empty if all contexts were already present.
+        """
+        tokens: dict[type[BaseContext], Token] = {}
+
+        if self._registration_context is not None:
+            try:
+                RegistrationContext.get()
+            except LookupError:
+                tokens[RegistrationContext] = RegistrationContext.set(
+                    self._registration_context
+                )
+
+        if (
+            self._event_processor is not None
+            and self._event_processor._root_context is not None
+        ):
+            try:
+                EventContext.get()
+            except LookupError:
+                tokens[EventContext] = EventContext.set(
+                    self._event_processor._root_context
+                )
+
+        return tokens
+
+    def set_contexts(self) -> contextlib.AbstractContextManager:
+        """Set Reflex contexts needed for state and event processing.
+
+        Pushes RegistrationContext and EventContext into the current
+        contextvars scope, but only if they are not already set.
+
+        Can be used as a context manager::
+
+            with app.set_contexts():
+                async with app.modify_state(token) as state:
+                    ...
+
+        Returns:
+            A context manager that resets any contexts that were set on exit.
+        """
+        tokens = self._set_contexts_internal()
+        if not tokens:
+            return contextlib.nullcontext()
+        stack = contextlib.ExitStack()
+        for ctx_cls, tok in tokens.items():
+            stack.callback(ctx_cls.reset, tok)
+        return stack
+
+    def _context_middleware(self, app: ASGIApp) -> ASGIApp:
+        """Ensure Reflex contexts are attached for each ASGI request.
+
+        Many ASGI servers start each request with a fresh contextvars scope,
+        so this middleware re-applies the RegistrationContext and EventContext
+        that are needed for Reflex state and event processing.
 
         Args:
             app: The ASGI app to attach the middleware to.
@@ -587,14 +647,11 @@ def _registration_context_middleware(self, app: ASGIApp) -> ASGIApp:
             The ASGI app with the middleware attached.
         """
 
-        async def registration_context_middleware(
-            scope: Scope, receive: Receive, send: Send
-        ):
-            if self._registration_context is not None:
-                RegistrationContext.set(self._registration_context)
+        async def context_middleware(scope: Scope, receive: Receive, send: Send):
+            self._set_contexts_internal()
             await app(scope, receive, send)
 
-        return registration_context_middleware
+        return context_middleware
 
     @contextlib.asynccontextmanager
     async def _setup_event_processor(self) -> AsyncIterator[None]:
@@ -672,10 +729,10 @@ def __call__(self) -> ASGIApp:
                     asgi_app = api_transformer(asgi_app)
 
         top_asgi_app = Starlette(lifespan=self._run_lifespan_tasks)
-        # Make sure the RegistrationContext is attached.
+        # Make sure Reflex contexts are attached for each request.
         top_asgi_app.mount(
             "",
-            self._registration_context_middleware(asgi_app),
+            self._context_middleware(asgi_app),
         )
         App._add_cors(top_asgi_app)
         return top_asgi_app
@@ -1615,20 +1672,22 @@ async def modify_state(
         if isinstance(token, str):
             token = BaseStateToken.from_legacy_token(token, root_state=self._state)
 
-        # Get exclusive access to the state.
-        async with self.state_manager.modify_state_with_links(
-            token, previous_dirty_vars=previous_dirty_vars, **context
-        ) as state:
-            # No other event handler can modify the state while in this context.
-            yield state
-            delta = await state._get_resolved_delta()
-            state._clean()
-            if delta:
-                # When the frontend vars are modified emit the delta to the frontend.
-                await self.event_namespace.emit_update(
-                    update=StateUpdate(delta=delta),
-                    token=token.ident,
-                )
+        # Ensure Reflex contexts are available (e.g. when called from an API route).
+        with self.set_contexts():
+            # Get exclusive access to the state.
+            async with self.state_manager.modify_state_with_links(
+                token, previous_dirty_vars=previous_dirty_vars, **context
+            ) as state:
+                # No other event handler can modify the state while in this context.
+                yield state
+                delta = await state._get_resolved_delta()
+                state._clean()
+                if delta:
+                    # When the frontend vars are modified emit the delta to the frontend.
+                    await self.event_namespace.emit_update(
+                        update=StateUpdate(delta=delta),
+                        token=token.ident,
+                    )
 
     def _validate_exception_handlers(self):
         """Validate the custom event exception handlers for front- and backend.

reflex/istate/shared.py

@@ -173,6 +173,38 @@ def _rehydrate(self):
             State.set_is_hydrated(True),
         ]
 
+    async def _resolve_linked_state(
+        self, state_cls: type["BaseState"], linked_token: str
+    ) -> "BaseState":
+        """Load and patch a linked state that was not pre-loaded in the tree.
+
+        Called by State._get_state_from_redis when a state in
+        _reflex_internal_links is not yet in the cache. This loads the
+        private copy into the tree first, then patches the linked version
+        on top of it via _internal_patch_linked_state.
+
+        Args:
+            state_cls: The shared state class to resolve.
+            linked_token: The shared token the state is linked to.
+
+        Returns:
+            The linked state instance, patched into the current tree.
+
+        Raises:
+            ReflexRuntimeError: If the resolved state is not a SharedState.
+        """
+        root_state = self._get_root_state()
+
+        # Load the private copy into the tree so _internal_patch_linked_state
+        # has an original to swap out (needed for unlink / restore).
+        original_state = await BaseState._get_state_from_redis(root_state, state_cls)
+
+        if isinstance(original_state, SharedStateBaseInternal):
+            return await original_state._internal_patch_linked_state(linked_token)
+
+        msg = f"Failed to resolve linked state {state_cls.get_full_name()} for token {linked_token}: state does not inherit from rx.SharedState"
+        raise ReflexRuntimeError(msg)
+
     async def _link_to(self, token: str) -> Self:
         """Link this shared state to a token.
 
@@ -194,7 +226,7 @@ async def _link_to(self, token: str) -> Self:
             raise ReflexRuntimeError(msg)
         if not isinstance(self, SharedState):
             msg = "Can only link SharedState instances."
-            raise RuntimeError(msg)
+            raise ReflexRuntimeError(msg)
         if self._linked_to == token:
             return self  # already linked to this token
         if self._linked_to and self._linked_to != token:
@@ -280,6 +312,18 @@ async def _internal_patch_linked_state(
                     BaseStateToken(ident=token, cls=type(self))
                 )
             )
+            # Set client_token on the linked root so that subsequent get_state
+            # calls when directly modifying a linked token will load the
+            # associated instance.
+            if linked_root_state.router.session.client_token != token:
+                import dataclasses as dc
+
+                linked_root_state.router = dc.replace(
+                    linked_root_state.router,
+                    session=dc.replace(
+                        linked_root_state.router.session, client_token=token
+                    ),
+                )
             self._held_locks.setdefault(token, {})
         else:
             linked_root_state = await get_state_manager().get_state(
@@ -386,6 +430,22 @@ async def _modify_linked_states(
                             for token in linked_state._linked_from
                             if token != self.router.session.client_token
                         )
+                # When modifying a shared token directly (empty _reflex_internal_links),
+                # the held locks will be empty. Check SharedState substates for linked
+                # clients that need to be notified.
+                if not self._reflex_internal_links:
+                    shared_state_base_internal = await self.get_state(
+                        SharedStateBaseInternal
+                    )
+                    if not isinstance(
+                        shared_state_base_internal, SharedStateBaseInternal
+                    ):
+                        msg = "Expected SharedStateBaseInternal in substates."
+                        raise ReflexRuntimeError(msg)
+                    # Collect affected tokens from all potentially linked states.
+                    shared_state_base_internal._collect_shared_token_updates(
+                        affected_tokens, current_dirty_vars
+                    )
         finally:
             self._exit_stack = None
 
@@ -397,6 +457,34 @@ async def _modify_linked_states(
                 state_type=type(self),
             )
 
+    def _collect_shared_token_updates(
+        self,
+        affected_tokens: set[str],
+        current_dirty_vars: dict[str, set[str]],
+    ) -> None:
+        """Recursively collect dirty vars and linked clients from SharedState substates.
+
+        When a shared state is modified directly by its shared token (rather than
+        through a private client token), the held locks are empty so the normal
+        collection loop above finds nothing. This method recursively checks
+        SharedState substates for linked clients that need to be notified.
+
+        Args:
+            affected_tokens: Set to update with client tokens that need notification.
+            current_dirty_vars: Dict to update with dirty var mappings per state.
+        """
+        for substate in self.substates.values():
+            if not isinstance(substate, SharedState):
+                continue
+            if substate._linked_from:
+                if substate._previous_dirty_vars:
+                    current_dirty_vars[substate.get_full_name()] = set(
+                        substate._previous_dirty_vars
+                    )
+                if substate._get_was_touched() or substate._previous_dirty_vars:
+                    affected_tokens.update(substate._linked_from)
+            substate._collect_shared_token_updates(affected_tokens, current_dirty_vars)
+
 
 class SharedState(SharedStateBaseInternal, mixin=True):
     """Mixin for defining new shared states."""

reflex/state.py

@@ -2146,7 +2146,6 @@ async def _get_state_from_redis(self, state_cls: type[T_STATE]) -> T_STATE:
         Returns:
             The instance of state_cls associated with this state's client_token.
         """
-        state_instance = await super()._get_state_from_redis(state_cls)
         if (
             self._reflex_internal_links
             and (
@@ -2155,15 +2154,12 @@ async def _get_state_from_redis(self, state_cls: type[T_STATE]) -> T_STATE:
                 )
             )
             is not None
-            and (
-                internal_patch_linked_state := getattr(
-                    state_instance, "_internal_patch_linked_state", None
-                )
-            )
-            is not None
         ):
-            return await internal_patch_linked_state(linked_token)
-        return state_instance
+            from reflex.istate.shared import SharedStateBaseInternal
+
+            shared_base = await self.get_state(SharedStateBaseInternal)
+            return await shared_base._resolve_linked_state(state_cls, linked_token)  # type: ignore[return-value]
+        return await super()._get_state_from_redis(state_cls)
 
     @event
     async def hydrate(self) -> None: