fix(events): cancel disconnected page-load chains (#7415)

* fix(events): cancel disconnected page-load chains

* fix(events): track rehydrated page-load chains

* fix(events): guard page loads during hydration

---------

Co-authored-by: Masen Furer <m_github@0x26.net>
Co-authored-by: Farhan Ali Raza <62690310+FarhanAliRaza@users.noreply.github.com>

Author: 3 people

docs/events/page_load_events.md

@@ -41,6 +41,17 @@ def index():
     return rx.text("A Beautiful App")
 ```
 
+## Cancellation
+
+Unfinished `on_load` event chains are cancelled when the client disconnects,
+navigates to another page or reconnects. This includes events returned or yielded
+by the page-load handler. A reconnect starts the new page's load events again.
+Independent background events started outside the page-load chain continue
+running after a disconnect.
+
+Use `try`/`finally` to release resources when an async handler is cancelled. Avoid
+suppressing `asyncio.CancelledError`, which prevents the handler from stopping.
+
 ## Handling Loading and Errors
 
 Avoid heavy synchronous work directly in an `on_load` handler — the page renders before the handler finishes, so a slow handler leaves the user staring at stale or empty data with no feedback. Instead:

news/+cancel-disconnected-page-loads.bugfix.md

@@ -0,0 +1 @@
+Cancel unfinished page-load event chains when the client disconnects, including hydration that has not yet started. Independent background events continue running.

packages/reflex-base/news/+cancel-disconnected-page-loads.bugfix.md

@@ -0,0 +1 @@
+Cancel pending hydration and page-load event chains when their client disconnects.

packages/reflex-base/src/reflex_base/event/processor/base_state_processor.py

@@ -31,7 +31,7 @@
 logger = logging.getLogger(__name__)
 
 if TYPE_CHECKING:
-    from reflex.state import BaseState
+    from reflex.state import BaseState, State
 
 # Resolved once at import: find_spec on a missing package scans sys.path (~90us),
 # far too slow for the per-event-argument path below.
@@ -394,44 +394,88 @@ class BaseStateEventProcessor(EventProcessor):
     """
 
     @staticmethod
-    async def _rehydrate(root_state: BaseState):
-        """Rehydrate the state by calling the hydrate event handler.
+    async def _rehydrate_snapshot(root_state: State) -> None:
+        """Send hydration data, finishing hydration when no page is available.
 
         Args:
             root_state: The root state to rehydrate.
         """
-        from reflex.state import OnLoadInternalState, State
-
-        if type(root_state) is not State:
-            return
+        from reflex.state import State
 
-        # A backend-initiated event carries no route, so nothing sets
-        # router_data and every one of them would rehydrate again.
         routeless = not root_state.router_data
         if routeless and root_state.is_hydrated:
             return
-
         await process_event(
             handler=State.event_handlers["hydrate"],
             payload={},
             state=root_state,
             root_state=root_state,
         )
         if routeless:
-            # No page to load, but hydration still has to finish.
             await process_event(
                 handler=State.event_handlers["set_is_hydrated"],
                 payload={"value": True},
                 state=root_state,
                 root_state=root_state,
             )
+
+    @functools.cached_property
+    def _pending_hydrates(self) -> dict[str, dict[str, EventFuture]]:
+        """Hydrations cancellable on disconnect, but never by a navigation.
+
+        Returns:
+            Pending hydration futures indexed by client token and transaction id.
+        """
+        return {}
+
+    async def _rehydrate(self, root_state: BaseState):
+        """Rehydrate the state by calling the hydrate event handler.
+
+        Args:
+            root_state: The root state to rehydrate.
+        """
+        from reflex.event import Event
+        from reflex.state import OnLoadInternalState, State
+
+        if type(root_state) is not State:
+            return
+
+        # A backend-initiated event carries no route, so nothing sets
+        # router_data and every one of them would rehydrate again.
+        if not root_state.router_data:
+            await self._rehydrate_snapshot(root_state)
             return
-        await process_event(
-            handler=OnLoadInternalState.event_handlers["on_load_internal"],
-            payload={},
-            state=await root_state.get_state(OnLoadInternalState),
-            root_state=root_state,
-        )
+
+        ctx = EventContext.get()
+        # Disconnect cancels this marker without cancelling the ordinary event.
+        pending = EventFuture(txid=ctx.txid)
+        self._pending_hydrates.setdefault(ctx.token, {})[pending.txid] = pending
+        try:
+            await self._rehydrate_snapshot(root_state)
+            if pending.cancelled():
+                return
+            parent = self._futures[ctx.txid]
+            first_child = len(parent.children)
+            handler = OnLoadInternalState.event_handlers["on_load_internal"]
+            await process_event(
+                handler=handler,
+                payload={},
+                state=await root_state.get_state(OnLoadInternalState),
+                root_state=root_state,
+            )
+            # Inline rehydration bypasses enqueue's page-load registration. Track
+            # only its children so the ordinary event remains independent.
+            load_event = Event(name=format_event_handler(handler))
+            for index in range(first_child, len(parent.children)):
+                child = parent.children[index]
+                if pending.cancelled():
+                    child.cancel()
+                else:
+                    self._supersede_previous(
+                        token=ctx.token, event=load_event, tracked=child
+                    )
+        finally:
+            self._forget_hydrate(ctx.token, pending)
 
     def _supersede_previous(
         self, *, token: str, event: Event, tracked: EventFuture
@@ -455,8 +499,36 @@ def _supersede_previous(
         boot_name, on_load_name = _connect_supersedes()
         if event.name == boot_name:
             self._cancel_older_chains((on_load_name, token), tracked.root_gen)
+            self._pending_hydrates.setdefault(token, {})[tracked.txid] = tracked
+            tracked.add_done_callback(functools.partial(self._forget_hydrate, token))
         return super()._supersede_previous(token=token, event=event, tracked=tracked)
 
+    def _forget_hydrate(self, token: str, future: EventFuture) -> None:
+        """Stop tracking a completed hydration invocation.
+
+        Its page-load descendants have their own supersession registration.
+
+        Args:
+            token: The client token.
+            future: The completed hydration future.
+        """
+        if (pending := self._pending_hydrates.get(token)) is not None:
+            pending.pop(future.txid, None)
+            if not pending:
+                del self._pending_hydrates[token]
+
+    def _on_disconnect(self, token: str) -> None:
+        """Cancel hydration and unfinished page-load chains for a disconnected client.
+
+        Args:
+            token: The disconnected client token.
+        """
+        _, on_load_name = _connect_supersedes()
+        for future in self._pending_hydrates.pop(token, {}).values():
+            future.cancel()
+        for future in self._superseded.pop((on_load_name, token), {}).values():
+            future.cancel()
+
     async def _execute_event(
         self, *, entry: EventQueueEntry, registered_handler: RegisteredEventHandler
     ) -> None:

packages/reflex-base/src/reflex_base/event/processor/event_processor.py

@@ -460,6 +460,15 @@ async def enqueue(
         await queue.put(EventQueueEntry(event=event, ctx=ev_ctx))
         return tracked
 
+    def _on_disconnect(self, token: str) -> None:
+        """Handle a client disconnect without cancelling independent event work.
+
+        Subclasses may cancel work tied to the disconnected page's lifecycle.
+
+        Args:
+            token: The disconnected client token.
+        """
+
     async def enqueue_many(self, token: str, *events: Event) -> Sequence[EventFuture]:
         """Enqueue multiple events to be processed.
 
@@ -585,15 +594,15 @@ def _try_clean_future(self, future: EventFuture) -> None:  # type: ignore[overri
                 return
             parent = future.parent
             self._futures.pop(future.txid, None)
-            if (
-                (key := future.supersede_key) is not None
-                and (slot := self._superseded.get(key)) is not None
-                and slot.get(future.txid) is future
-                and future.all_done()
-            ):
-                del slot[future.txid]
-                if not slot:
-                    del self._superseded[key]
+            if future.supersede_key is not None and future.all_done():
+                # Inline page loads can also have their own supersession key.
+                for key in future.covered_supersede_keys:
+                    if (slot := self._superseded.get(key)) is not None and slot.get(
+                        future.txid
+                    ) is future:
+                        del slot[future.txid]
+                        if not slot:
+                            del self._superseded[key]
             if parent is None or not parent.txid:
                 return
             future = parent

reflex/app.py

@@ -2319,6 +2319,14 @@ def on_disconnect(self, sid: str) -> asyncio.Task | None:
         # Get token before cleaning up
         disconnect_token = self.sid_to_token.get(sid)
         if disconnect_token:
+            socket_record = self._token_manager.token_to_socket.get(disconnect_token)
+            if (
+                socket_record is not None
+                and socket_record.sid == sid
+                and socket_record.instance_id == self._token_manager.instance_id
+                and (processor := self.app._event_processor) is not None
+            ):
+                processor._on_disconnect(disconnect_token)
             # Use async cleanup through token manager
             task = asyncio.create_task(
                 self._token_manager.disconnect_token(disconnect_token, sid),

reflex/istate/proxy.py

@@ -281,7 +281,7 @@ async def __aenter__(self) -> Self:
                 # next event loads its current page.
                 from reflex_base.event.processor import BaseStateEventProcessor
 
-                await BaseStateEventProcessor._rehydrate(mutable_state)
+                await BaseStateEventProcessor._rehydrate_snapshot(mutable_state)
         except (Exception, asyncio.CancelledError):
             # Restore the proxy to a consistent state since __aexit__ will not be called when __aenter__ raises.
             await self.__aexit__(*sys.exc_info())

tests/units/conftest.py

@@ -350,6 +350,9 @@ def mock_base_state_event_processor_obj(
         A fresh BaseState event processor.
     """
     monkeypatch.setattr(BaseStateEventProcessor, "_rehydrate", mock.AsyncMock())
+    monkeypatch.setattr(
+        BaseStateEventProcessor, "_rehydrate_snapshot", mock.AsyncMock()
+    )
 
     def handle_backend_exception(ex: Exception) -> None:
         formatted_exc = "\n".join(traceback.format_exception(ex))

tests/units/istate/test_proxy.py

@@ -198,7 +198,9 @@ async def failing_rehydrate(root_state: BaseState):  # noqa: RUF029
         msg = "Simulated rehydrate failure"
         raise RuntimeError(msg)
 
-    monkeypatch.setattr(BaseStateEventProcessor, "_rehydrate", failing_rehydrate)
+    monkeypatch.setattr(
+        BaseStateEventProcessor, "_rehydrate_snapshot", failing_rehydrate
+    )
     with pytest.raises(RuntimeError, match="Simulated rehydrate failure"):
         async with state_proxy:
             pass