Add rx.event arg to cancel previous background tasks

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.
             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 = "_cancel_key"
 _StreamItemT = TypeVar("_StreamItemT")
 
 
@@ -579,6 +580,16 @@ 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)
+            for existing in self._tasks.values():
+                if (
+                    getattr(existing, CANCEL_KEY, None) == cancel_key
+                    and not existing.done()
+                ):
+                    existing.cancel()
         task = asyncio.create_task(
             self._process_event_queue_entry(
                 entry=entry, registered_handler=registered_handler
@@ -587,6 +598,8 @@ 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._tasks[entry.ctx.txid] = task
         task.add_done_callback(self._finish_task)
         return task

tests/units/test_state.py

@@ -2618,6 +2618,74 @@ 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]
+
+        # 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]
+
+        # Let the surviving run finish.
+        cancel_prev_release[0].set()
+        await asyncio.sleep(settle)
+
+    # 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."""