ENG-9661 fix: cancel stale on_load chains via handler-declared supersession

news/6593.bugfix.md

@@ -0,0 +1 @@
+Cancel the previous unfinished `on_load` event chain when a newer page navigation arrives for the same client, instead of letting stale page-load work block and outlive the navigation.

packages/reflex-base/src/reflex_base/event/__init__.py

@@ -185,6 +185,7 @@ def from_event_type(
 )
 
 BACKGROUND_TASK_MARKER = "_reflex_background_task"
+SUPERSEDES_MARKER = "_reflex_supersedes"
 EVENT_ACTIONS_MARKER = "_rx_event_actions"
 UPLOAD_FILES_CLIENT_HANDLER = "uploadFiles"
 
@@ -442,6 +443,19 @@ def is_background(self) -> bool:
         """
         return getattr(self.fn, BACKGROUND_TASK_MARKER, False)
 
+    @property
+    def supersedes(self) -> bool:
+        """Whether a newer chain-root invocation supersedes an older one.
+
+        When True, enqueuing this handler as a chain root cancels the previous
+        unfinished event chain rooted at the same handler for the same client
+        token.
+
+        Returns:
+            True if the event handler is marked as superseding.
+        """
+        return getattr(self.fn, SUPERSEDES_MARKER, False)
+
     def __call__(self, *args: Any, **kwargs: Any) -> "EventSpec":
         """Pass arguments to the handler to get an event spec.
 
@@ -2778,6 +2792,7 @@ class EventNamespace:
 
     # Constants
     BACKGROUND_TASK_MARKER = BACKGROUND_TASK_MARKER
+    SUPERSEDES_MARKER = SUPERSEDES_MARKER
     EVENT_ACTIONS_MARKER = EVENT_ACTIONS_MARKER
     _EVENT_FIELDS = _EVENT_FIELDS
     FORM_DATA = FORM_DATA

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

@@ -120,6 +120,11 @@ class EventProcessor:
     _futures: dict[str, EventFuture] = dataclasses.field(
         default_factory=dict, init=False
     )
+    # Latest-wins tracking for superseding handlers: (event name, token) -> the
+    # currently active chain root future.
+    _superseded: dict[tuple[str, str], EventFuture] = dataclasses.field(
+        default_factory=dict, init=False
+    )
     _token_queues: dict[
         str,
         collections.deque[tuple[EventQueueEntry, RegisteredEventHandler]],
@@ -311,6 +316,7 @@ async def stop(self, graceful_shutdown_timeout: float | None = None) -> None:
             self._queue_task = None
         # Discard any pending per-token queue entries.
         self._token_queues.clear()
+        self._superseded.clear()
         # Cancel any remaining unresolved futures.
         for future in self._futures.values():
             if not future.done():
@@ -372,6 +378,8 @@ async def enqueue(
 
         Returns:
             An EventFuture that resolves to the result of the associated task.
+            If the event was chained from an already-cancelled chain, the
+            returned future is already cancelled and the event is dropped.
         """
         if ev_ctx is None:
             try:
@@ -395,7 +403,14 @@ async def enqueue(
         tracked.add_done_callback(self._on_future_done)
         # If this context has a parent, register as a child of the parent's future.
         if parent_future is not None:
+            if parent_future.cancelled():
+                # The chain this event belongs to was cancelled; the event is
+                # stillborn and never enters the queue.
+                tracked.cancel()
+                return tracked
             parent_future.add_child(tracked)
+        if parent_future is None:
+            self._supersede_previous(token=token, event=event, tracked=tracked)
         await queue.put(EventQueueEntry(event=event, ctx=ev_ctx))
         return tracked
 
@@ -491,14 +506,51 @@ def _try_clean_future(self, future: EventFuture) -> None:  # type: ignore[overri
         """
         if not future.done():
             return
+        if future.txid in self._tasks:
+            # The handler task is still running or unwinding; keep the future
+            # so late-chained events can find their (possibly cancelled) parent.
+            return
         # Not checking future.all_done() to avoid waiting for grandchildren here.
         if not all(c.done() for c in future.children):
             return
         parent = future.parent
         self._futures.pop(future.txid, None)
+        if (
+            (key := future.supersede_key) is not None
+            and self._superseded.get(key) is future
+            and future.all_done()
+        ):
+            del self._superseded[key]
         if parent is not None and parent.txid:
             self._try_clean_future(parent)
 
+    def _supersede_previous(
+        self, *, token: str, event: Event, tracked: EventFuture
+    ) -> None:
+        """Cancel the previous unfinished chain of a superseding event handler.
+
+        Root handlers marked with ``supersedes`` (e.g. ``on_load_internal``)
+        use latest-wins semantics: enqueuing a new invocation cancels the
+        previous unfinished event chain for the same handler and client token.
+
+        Args:
+            token: The client token associated with the event.
+            event: The event being enqueued.
+            tracked: The future of the event being enqueued.
+        """
+        try:
+            registered = RegistrationContext.get().event_handlers.get(event.name)
+        except LookupError:
+            return
+        if registered is None or not registered.handler.supersedes:
+            return
+        key = (event.name, token)
+        previous = self._superseded.get(key)
+        if previous is not None and not previous.all_done():
+            previous.cancel()
+        self._superseded[key] = tracked
+        tracked.supersede_key = key
+
     def _on_future_done(self, future: EventFuture) -> None:  # type: ignore[override]
         """Callback invoked when an enqueued future completes.
 
@@ -625,10 +677,13 @@ def _dispatch_next_for_token(self, token: str) -> None:
         if not token_queue:
             return
         entry, registered_handler = token_queue[0]
-        # Skip cancelled futures.
+        # Skip cancelled futures. Before a task exists, the only way a future
+        # can be done (and thus already cleaned up) is cancellation, so a
+        # missing future also means the entry was cancelled.
         future = self._futures.get(entry.ctx.txid)
-        if future is not None and future.cancelled():
-            self._try_clean_future(future)
+        if future is None or future.cancelled():
+            if future is not None:
+                self._try_clean_future(future)
             token_queue.popleft()
             if token_queue:
                 self._dispatch_next_for_token(token)
@@ -645,10 +700,12 @@ async def _process_queue(self):
         with contextlib.suppress(QueueShutDown):
             while True:
                 entry = await queue.get()
-                if (
-                    future := self._futures.get(entry.ctx.txid)
-                ) is not None and future.cancelled():
-                    self._try_clean_future(future)
+                # A missing future means the entry was cancelled and already
+                # cleaned up (see _dispatch_next_for_token).
+                future = self._futures.get(entry.ctx.txid)
+                if future is None or future.cancelled():
+                    if future is not None:
+                        self._try_clean_future(future)
                     queue.task_done()
                     continue
                 try:
@@ -757,6 +814,10 @@ def _finish_task(self, task: asyncio.Task):
             else:
                 if future is not None and not future.done():
                     future.set_result(result)
+        if future is not None:
+            # The task is gone; clean up now in case the future resolved
+            # earlier (e.g. external cancellation) and cleanup was deferred.
+            self._try_clean_future(future)
 
 
 __all__ = [

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

@@ -31,6 +31,10 @@ class EventFuture(asyncio.Future):
         default_factory=asyncio.get_running_loop, repr=False
     )
 
+    # Key under which this future is registered for latest-wins supersession
+    # in the EventProcessor, if any.
+    supersede_key: tuple[str, str] | None = dataclasses.field(default=None, repr=False)
+
     def __post_init__(self) -> None:
         """Call Future.__init__ for the EventFuture."""
         super(EventFuture, self).__init__(loop=self.loop)

reflex/state.py

@@ -33,6 +33,7 @@
 from reflex_base.event import (
     BACKGROUND_TASK_MARKER,
     EVENT_ACTIONS_MARKER,
+    SUPERSEDES_MARKER,
     Event,
     EventHandler,
     EventSpec,
@@ -755,8 +756,9 @@ def _copy_fn(fn: Callable) -> Callable:
             closure=fn.__closure__,
         )
         newfn.__annotations__ = fn.__annotations__
-        if mark := getattr(fn, BACKGROUND_TASK_MARKER, None):
-            setattr(newfn, BACKGROUND_TASK_MARKER, mark)
+        for marker in (BACKGROUND_TASK_MARKER, SUPERSEDES_MARKER):
+            if mark := getattr(fn, marker, None):
+                setattr(newfn, marker, mark)
         # Preserve event_actions from @rx.event decorator
         if event_actions := getattr(fn, EVENT_ACTIONS_MARKER, None):
             object.__setattr__(newfn, EVENT_ACTIONS_MARKER, event_actions)
@@ -2482,6 +2484,15 @@ def on_load_internal(self) -> list[Event | EventSpec | event.EventCallback] | No
         ]
 
 
+# A newer navigation supersedes the previous unfinished on_load chain for the
+# same client token, cancelling its stale work (#6593).
+setattr(
+    OnLoadInternalState.event_handlers["on_load_internal"].fn,
+    SUPERSEDES_MARKER,
+    True,
+)
+
+
 class ComponentState(State, mixin=True):
     """Base class to allow for the creation of a state instance per component.