Add rx.event arg to cancel previous background tasks

packages/reflex-base/news/6793.feature.md

@@ -0,0 +1 @@
+Added argument to rx.event to cancel previously scheduled events.

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

@@ -291,6 +291,7 @@ def _scan_detach(value: Any, memo: dict[int, Any], active: set[int]) -> Any:
 
 
 BACKGROUND_TASK_MARKER = "_reflex_background_task"
+CANCEL_PREVIOUS_TASK_MARKER = "_reflex_cancel_previous_task"
 EVENT_ACTIONS_MARKER = "_rx_event_actions"
 UPLOAD_FILES_CLIENT_HANDLER = "uploadFiles"
 
@@ -548,6 +549,16 @@ def is_background(self) -> bool:
         """
         return getattr(self.fn, BACKGROUND_TASK_MARKER, False)
 
+    @property
+    def is_cancel_previous_task(self) -> bool:
+        """Whether starting this handler cancels its previous run.
+        Requires the handler to be a background task.
+
+        Returns:
+            True if the event handler is marked to cancel its previous run.
+        """
+        return getattr(self.fn, CANCEL_PREVIOUS_TASK_MARKER, False)
+
     def __call__(self, *args: Any, **kwargs: Any) -> "EventSpec":
         """Pass arguments to the handler to get an event spec.
 
@@ -2909,6 +2920,7 @@ def __new__(
         func: None = None,
         *,
         background: bool | None = None,
+        cancel_previous_task: bool | None = None,
         stop_propagation: bool | None = None,
         prevent_default: bool | None = None,
         throttle: int | None = None,
@@ -2924,6 +2936,7 @@ def __new__(
         func: "Callable[[BASE_STATE, Unpack[P]], Any]",
         *,
         background: bool | None = None,
+        cancel_previous_task: bool | None = None,
         stop_propagation: bool | None = None,
         prevent_default: bool | None = None,
         throttle: int | None = None,
@@ -2936,6 +2949,7 @@ def __new__(
         func: "Callable[[BASE_STATE, Unpack[P]], Any] | None" = None,
         *,
         background: bool | None = None,
+        cancel_previous_task: bool | None = None,
         stop_propagation: bool | None = None,
         prevent_default: bool | None = None,
         throttle: int | None = None,
@@ -2947,6 +2961,7 @@ def __new__(
         Args:
             func: The function to wrap.
             background: Whether the event should be run in the background. Defaults to False.
+            cancel_previous_task: Whether dispatching this handler cancels its still-running previous run for the same client token. Requires background=True. NOTE: This can cause partial state updates if previous task always modified state before cancellation.
             stop_propagation: Whether to stop the event from bubbling up the DOM tree.
             prevent_default: Whether to prevent the default behavior of the event.
             throttle: Throttle the event handler to limit calls (in milliseconds).
@@ -2958,6 +2973,7 @@ def __new__(
 
         Raises:
             TypeError: If background is True and the function is not a coroutine or async generator. # noqa: DAR402
+            ValueError: If cancel_previous_task is True but background is not True. # noqa: DAR402
         """
 
         def _build_event_actions():
@@ -2998,6 +3014,11 @@ def wrapper(
                     msg = "Background task must be async function or generator."
                     raise TypeError(msg)
                 setattr(func, BACKGROUND_TASK_MARKER, True)
+            if cancel_previous_task:
+                if background is not True:
+                    msg = "cancel_previous_task=True requires background=True."
+                    raise ValueError(msg)
+                setattr(func, CANCEL_PREVIOUS_TASK_MARKER, True)
             if getattr(func, "__name__", "").startswith("_"):
                 msg = "Event handlers cannot be private."
                 raise ValueError(msg)

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

@@ -41,6 +41,7 @@ class QueueShutDown(Exception):  # noqa: N818
         """Exception raised when trying to put an item into a shut down queue."""
 
 
+CANCEL_KEY = "_reflex_cancel_key"
 _StreamItemT = TypeVar("_StreamItemT")
 
 
@@ -98,6 +99,7 @@ class EventProcessor:
         _root_context: The root event context to use for events enqueued without an explicit context.
         _attached_root_context_token: The context variable token for the attached root context, used to reset the context variable on shutdown.
         _tasks: A mapping of active transaction ids to their corresponding event handler tasks, used for tracking and cancellation on shutdown.
+        _cancel_keys: A mapping of ``cancel_previous_task`` keys (token, event name) to the txid of the most recently dispatched task for that key.
     """
 
     middleware: MiddlewareMixin | None = None
@@ -124,6 +126,9 @@ class EventProcessor:
         str,
         collections.deque[tuple[EventQueueEntry, RegisteredEventHandler]],
     ] = dataclasses.field(default_factory=dict, init=False)
+    _cancel_keys: dict[tuple[str, str], str] = dataclasses.field(
+        default_factory=dict, init=False
+    )
 
     def configure(
         self,
@@ -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._cancel_keys.clear()
         # Cancel any remaining unresolved futures.
         for future in self._futures.values():
             if not future.done():
@@ -579,6 +585,14 @@ def _create_event_task(
         Returns:
             The created asyncio.Task.
         """
+        handler = registered_handler.handler
+        cancel_key: tuple[str, str] | None = None
+        if handler.is_background and handler.is_cancel_previous_task:
+            cancel_key = (entry.ctx.token, entry.event.name)
+            if (existing_txid := self._cancel_keys.get(cancel_key)) is not None:
+                existing = self._tasks.get(existing_txid)
+                if existing is not None and not existing.done():
+                    existing.cancel()
         task = asyncio.create_task(
             self._process_event_queue_entry(
                 entry=entry, registered_handler=registered_handler
@@ -587,6 +601,9 @@ def _create_event_task(
         )
         if sys.version_info < (3, 12):
             task._event_ctx = entry.ctx  # pyright: ignore[reportAttributeAccessIssue]
+        if cancel_key is not None:
+            setattr(task, CANCEL_KEY, cancel_key)
+            self._cancel_keys[cancel_key] = entry.ctx.txid
         self._tasks[entry.ctx.txid] = task
         task.add_done_callback(self._finish_task)
         return task
@@ -714,6 +731,12 @@ def _finish_task(self, task: asyncio.Task):
         else:
             task_ctx = task.get_context().run(EventContext.get)
         self._tasks.pop(task_ctx.txid, None)
+        cancel_key = getattr(task, CANCEL_KEY, None)
+        if (
+            cancel_key is not None
+            and self._cancel_keys.get(cancel_key) == task_ctx.txid
+        ):
+            del self._cancel_keys[cancel_key]
         # Chain the next sequential event for this token if applicable.
         token_queue = self._token_queues.get(task_ctx.token)
         if token_queue and token_queue[0][0].ctx.txid == task_ctx.txid:

tests/units/test_state.py

@@ -2618,6 +2618,80 @@ async def test_background_task_no_chain():
         await bts.bad_chain2()
 
 
+cancel_prev_started: list[int] = []
+cancel_prev_cancelled: list[int] = []
+cancel_prev_completed: list[int] = []
+cancel_prev_release: list[asyncio.Event] = []
+
+
+class CancelPreviousTaskState(BaseState):
+    """A state whose background task cancels its own previous in-flight run."""
+
+    runs: int = 0
+
+    @rx.event(background=True, cancel_previous_task=True)
+    async def slow_task(self):
+        """A slow background task that supersedes any still-running prior run."""
+        async with self:
+            self.runs += 1
+            run_id = self.runs
+        cancel_prev_started.append(run_id)
+        try:
+            await cancel_prev_release[0].wait()
+        except asyncio.CancelledError:
+            cancel_prev_cancelled.append(run_id)
+            raise
+        cancel_prev_completed.append(run_id)
+
+
+@pytest.mark.asyncio
+async def test_background_task_cancel_previous(
+    mock_app: rx.App,
+    token: str,
+    mock_base_state_event_processor: BaseStateEventProcessor,
+):
+    """Test that cancel_previous_task cancels the still-running prior run.
+
+    Args:
+        mock_app: An app that will be returned by `get_app()`
+        token: A token.
+        mock_base_state_event_processor: The event processor.
+    """
+    cancel_prev_started.clear()
+    cancel_prev_cancelled.clear()
+    cancel_prev_completed.clear()
+    cancel_prev_release.clear()
+    cancel_prev_release.append(asyncio.Event())
+
+    event_name = f"{CancelPreviousTaskState.get_full_name()}.slow_task"
+    settle = 0.5 if CI else 0.2
+    async with mock_base_state_event_processor as processor:
+        # First run starts and parks in its polling loop.
+        await processor.enqueue(token, Event(name=event_name, payload={}))
+        await asyncio.sleep(settle)
+        assert cancel_prev_started == [1]
+
+        assert (token, event_name) in processor._cancel_keys
+        assert len(processor._cancel_keys) == 1
+
+        # Second run supersedes the first; run 1 should be cancelled.
+        await processor.enqueue(token, Event(name=event_name, payload={}))
+        await asyncio.sleep(settle)
+        assert cancel_prev_started == [1, 2]
+        assert cancel_prev_cancelled == [1]
+        assert len(processor._cancel_keys) == 1
+
+        # Let the surviving run finish.
+        cancel_prev_release[0].set()
+        await asyncio.sleep(settle)
+
+        assert processor._cancel_keys == {}
+
+    # Only the second run completed; the first was cancelled mid-flight.
+    assert cancel_prev_cancelled == [1]
+    assert cancel_prev_completed == [2]
+
+
 class YieldFromBackgroundState(BaseState):
     """A state used to verify the type of `self` in a yielded event handler."""