Cancel stale on_load event chains

news/6593.bugfix.md

@@ -0,0 +1 @@
+Cancel stale page load events when a newer page navigation starts.

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

@@ -20,6 +20,7 @@
 from reflex.app_mixins.middleware import MiddlewareMixin
 from reflex.istate.manager import StateManager
 from reflex.utils import console
+from reflex_base import constants
 from reflex_base.event.context import EventContext
 from reflex_base.event.processor.compat import as_completed
 from reflex_base.event.processor.future import EventFuture
@@ -120,6 +121,12 @@ class EventProcessor:
     _futures: dict[str, EventFuture] = dataclasses.field(
         default_factory=dict, init=False
     )
+    _active_on_load_futures: dict[str, EventFuture] = dataclasses.field(
+        default_factory=dict, init=False
+    )
+    _active_on_load_tokens: dict[str, str] = dataclasses.field(
+        default_factory=dict, init=False
+    )
     _token_queues: dict[
         str,
         collections.deque[tuple[EventQueueEntry, RegisteredEventHandler]],
@@ -311,6 +318,8 @@ 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._active_on_load_futures.clear()
+        self._active_on_load_tokens.clear()
         # Cancel any remaining unresolved futures.
         for future in self._futures.values():
             if not future.done():
@@ -396,6 +405,8 @@ async def enqueue(
         # If this context has a parent, register as a child of the parent's future.
         if parent_future is not None:
             parent_future.add_child(tracked)
+        if self._is_on_load_internal_event(event):
+            self._cancel_stale_on_load(token=token, current=tracked)
         await queue.put(EventQueueEntry(event=event, ctx=ev_ctx))
         return tracked
 
@@ -496,9 +507,49 @@ def _try_clean_future(self, future: EventFuture) -> None:  # type: ignore[overri
             return
         parent = future.parent
         self._futures.pop(future.txid, None)
+        self._try_clean_active_on_load_future(future)
         if parent is not None and parent.txid:
             self._try_clean_future(parent)
 
+    @staticmethod
+    def _is_on_load_internal_event(event: Event) -> bool:
+        """Check whether an event is the internal page-load dispatcher.
+
+        Args:
+            event: The event to check.
+
+        Returns:
+            True if the event dispatches page on_load handlers.
+        """
+        return event.name.endswith(f".{constants.CompileVars.ON_LOAD_INTERNAL}")
+
+    def _cancel_stale_on_load(self, *, token: str, current: EventFuture) -> None:
+        """Cancel any older on_load chain for the token.
+
+        Args:
+            token: The client token for the page load event.
+            current: The newly enqueued on_load future.
+        """
+        previous = self._active_on_load_futures.get(token)
+        if previous is not None and previous is not current:
+            self._active_on_load_tokens.pop(previous.txid, None)
+            if not previous.all_done():
+                previous.cancel()
+        self._active_on_load_futures[token] = current
+        self._active_on_load_tokens[current.txid] = token
+
+    def _try_clean_active_on_load_future(self, future: EventFuture) -> None:
+        """Remove a completed active on_load future.
+
+        Args:
+            future: The future to remove if it is no longer active.
+        """
+        if not future.all_done():
+            return
+        token = self._active_on_load_tokens.pop(future.txid, None)
+        if token is not None and self._active_on_load_futures.get(token) is future:
+            self._active_on_load_futures.pop(token, None)
+
     def _on_future_done(self, future: EventFuture) -> None:  # type: ignore[override]
         """Callback invoked when an enqueued future completes.
 

tests/units/reflex_base/event/processor/test_event_processor.py

@@ -5,13 +5,14 @@
 from typing import Any
 
 import pytest
+from reflex_base import constants
 from reflex_base.event.context import EventContext
 from reflex_base.event.processor.event_processor import (
     EventProcessor,
     QueueShutDown,
     _stream_queue_until_done,
 )
-from reflex_base.registry import RegistrationContext
+from reflex_base.registry import RegisteredEventHandler, RegistrationContext
 
 from reflex.event import Event, EventHandler
 
@@ -445,6 +446,65 @@ async def test_error_does_not_stop_queue(
     assert _CALL_LOG == [{"value": "after_error"}]
 
 
+async def test_new_on_load_internal_cancels_previous_chain(
+    processor: EventProcessor,
+    token: str,
+):
+    """A newer navigation cancels stale on_load work for the same token.
+
+    Args:
+        processor: The event processor fixture.
+        token: The client token.
+    """
+    stale_started = asyncio.Event()
+    on_load_event_name = f"state.{constants.CompileVars.ON_LOAD_INTERNAL}"
+    load_handler_event_name = "state.load_page"
+
+    async def on_load_internal_handler(value: str):
+        ctx = EventContext.get()
+        await ctx.enqueue(Event(name=load_handler_event_name, payload={"value": value}))
+
+    async def load_handler(value: str):
+        if value == "stale":
+            stale_started.set()
+            try:
+                await asyncio.sleep(10)
+            except asyncio.CancelledError:
+                _CALL_LOG.append({"value": "stale_cancelled"})
+                raise
+        _CALL_LOG.append({"value": value})
+
+    registration_context = RegistrationContext.get()
+    registration_context.event_handlers[on_load_event_name] = RegisteredEventHandler(
+        handler=EventHandler(fn=on_load_internal_handler),
+        states=(),
+    )
+    registration_context.event_handlers[load_handler_event_name] = (
+        RegisteredEventHandler(
+            handler=EventHandler(fn=load_handler),
+            states=(),
+        )
+    )
+
+    processor.configure()
+    async with processor as ep:
+        stale = await ep.enqueue(
+            token,
+            Event(name=on_load_event_name, payload={"value": "stale"}),
+        )
+        await asyncio.wait_for(stale_started.wait(), timeout=1)
+        assert stale.done()
+        assert not stale.all_done()
+
+        current = await ep.enqueue(
+            token,
+            Event(name=on_load_event_name, payload={"value": "fresh"}),
+        )
+        await asyncio.wait_for(current.wait_all(), timeout=1)
+
+    assert _CALL_LOG == [{"value": "stale_cancelled"}, {"value": "fresh"}]
+
+
 async def test_chained_event_processed(token: str):
     """An event handler that enqueues another event via ctx.enqueue succeeds.