diff --git a/packages/reflex-base/news/6793.feature.md b/packages/reflex-base/news/6793.feature.md new file mode 100644 index 00000000000..db9f3af82ce --- /dev/null +++ b/packages/reflex-base/news/6793.feature.md @@ -0,0 +1 @@ +Added argument to rx.event to cancel previously scheduled events. diff --git a/packages/reflex-base/src/reflex_base/event/__init__.py b/packages/reflex-base/src/reflex_base/event/__init__.py index 919e1f2fd1f..4a4f131db31 100644 --- a/packages/reflex-base/src/reflex_base/event/__init__.py +++ b/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) diff --git a/packages/reflex-base/src/reflex_base/event/processor/event_processor.py b/packages/reflex-base/src/reflex_base/event/processor/event_processor.py index 7d9296fe4dd..cbe65efe83a 100644 --- a/packages/reflex-base/src/reflex_base/event/processor/event_processor.py +++ b/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: diff --git a/tests/units/test_state.py b/tests/units/test_state.py index 112aa1e8d43..ebcc20f8b52 100644 --- a/tests/units/test_state.py +++ b/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."""