EventProcessor.stop() partially cancels same-token event bursts instead of draining them

[FarhanAliRaza]

Problem

During graceful shutdown, EventProcessor.stop() only drains events whose handler tasks are already running. Events still waiting in a per-token sequential queue are discarded and their futures cancelled — even when graceful_shutdown_timeout has plenty of budget left.

With 10 same-token no-op events enqueued and a generous drain budget, only ~2-3 events complete; the rest are cancelled. With 10 independent tokens, all 10 drain cleanly.

Mechanism

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

Same-token (non-background) events are serialized through _token_queues: only the front entry per token has a task in self._tasks at any moment (_enqueue_for_token / _dispatch_next_for_token). The next task is only created by _finish_task when the previous one completes.

stop() drains in phases:

1. join() — drains the main asyncio queue, but that only moves entries into _token_queues, so it returns almost immediately.
2. _stop_tasks(timeout=...) — iterates as_completed(self._tasks.values(), ...) over a snapshot of currently-running tasks. Tasks dispatched afterwards by _finish_task chaining are not awaited; each _stop_tasks pass effectively lets ~1 more event per token finish.
3. self._token_queues.clear() — all still-queued same-token entries are discarded.
4. Remaining futures are cancelled.

So the drain budget is never actually applied to the per-token backlog.

Repro

import asyncio

from reflex_base.event.processor import EventProcessor
from reflex_base.registry import RegistrationContext

from reflex.event import Event, EventHandler


async def noop() -> None: ...


ctx = RegistrationContext.get_or_create_for_testing()
ctx.__enter__()
handler = EventHandler(fn=noop)
ctx.register_event_handler(handler)  # however registration is spelled; see tests/benchmarks/test_event_loop.py::_register_handlers


async def main():
    processor = EventProcessor(graceful_shutdown_timeout=60).configure()
    async with processor:
        futures = [
            await processor.enqueue("token", Event.from_event_type(handler())[0])
            for _ in range(10)
        ]
    done = sum(f.done() and not f.cancelled() for f in futures)
    print(f"{done}/10 drained")  # ~3/10; expected 10/10


asyncio.run(main())

A working registration setup lives in tests/benchmarks/test_event_loop.py (_register_handlers). test_event_shutdown_drain there deliberately uses independent tokens and documents this bug as the reason.

Expected

While the drain budget has time remaining, stop() should keep processing per-token queues (e.g. loop _stop_tasks while self._tasks or _token_queues are non-empty and the deadline hasn't passed) before clearing _token_queues and cancelling futures.

Found while adding drain assertions to the event-loop benchmarks in #6751.

[linear-code]

ENG-10195