From 473b4b6e62c004e8b72a2182f9dedb198146db4e Mon Sep 17 00:00:00 2001 From: Krishna Date: Mon, 5 Oct 2026 00:53:42 +0530 Subject: [PATCH 1/3] fix(events): cancel disconnected page-load chains --- docs/events/page_load_events.md | 11 ++ .../+cancel-disconnected-page-loads.bugfix.md | 1 + .../+cancel-disconnected-page-loads.bugfix.md | 1 + .../event/processor/base_state_processor.py | 37 ++++ .../event/processor/event_processor.py | 9 + pyproject.toml | 2 +- reflex/app.py | 8 + .../processor/test_base_state_processor.py | 171 +++++++++++++++++- tests/units/test_app.py | 31 ++++ tests/units/test_check_min_deps.py | 4 +- 10 files changed, 271 insertions(+), 4 deletions(-) create mode 100644 news/+cancel-disconnected-page-loads.bugfix.md create mode 100644 packages/reflex-base/news/+cancel-disconnected-page-loads.bugfix.md diff --git a/docs/events/page_load_events.md b/docs/events/page_load_events.md index 851e09ccce0..72899f41ae9 100644 --- a/docs/events/page_load_events.md +++ b/docs/events/page_load_events.md @@ -41,6 +41,17 @@ def index(): return rx.text("A Beautiful App") ``` +## Cancellation + +Unfinished `on_load` event chains are cancelled when the client disconnects, +navigates to another page or reconnects. This includes events returned or yielded +by the page-load handler. A reconnect starts the new page's load events again. +Independent background events started outside the page-load chain continue +running after a disconnect. + +Use `try`/`finally` to release resources when an async handler is cancelled. Avoid +suppressing `asyncio.CancelledError`, which prevents the handler from stopping. + ## Handling Loading and Errors Avoid heavy synchronous work directly in an `on_load` handler — the page renders before the handler finishes, so a slow handler leaves the user staring at stale or empty data with no feedback. Instead: diff --git a/news/+cancel-disconnected-page-loads.bugfix.md b/news/+cancel-disconnected-page-loads.bugfix.md new file mode 100644 index 00000000000..816760396be --- /dev/null +++ b/news/+cancel-disconnected-page-loads.bugfix.md @@ -0,0 +1 @@ +Cancel unfinished page-load event chains when the client disconnects, including hydration that has not yet started. Independent background events continue running. diff --git a/packages/reflex-base/news/+cancel-disconnected-page-loads.bugfix.md b/packages/reflex-base/news/+cancel-disconnected-page-loads.bugfix.md new file mode 100644 index 00000000000..f5bbec43a14 --- /dev/null +++ b/packages/reflex-base/news/+cancel-disconnected-page-loads.bugfix.md @@ -0,0 +1 @@ +Cancel pending hydration and page-load event chains when their client disconnects. diff --git a/packages/reflex-base/src/reflex_base/event/processor/base_state_processor.py b/packages/reflex-base/src/reflex_base/event/processor/base_state_processor.py index 82f3aacc447..338592251f9 100644 --- a/packages/reflex-base/src/reflex_base/event/processor/base_state_processor.py +++ b/packages/reflex-base/src/reflex_base/event/processor/base_state_processor.py @@ -377,6 +377,15 @@ class BaseStateEventProcessor(EventProcessor): frontend. """ + @functools.cached_property + def _pending_hydrates(self) -> dict[str, dict[str, EventFuture]]: + """Hydrations cancellable on disconnect, but never by a navigation. + + Returns: + Pending hydration futures indexed by client token and transaction id. + """ + return {} + async def _rehydrate(self, root_state: BaseState): """Rehydrate the state by calling the hydrate event handler. @@ -438,8 +447,36 @@ def _supersede_previous( boot_name, on_load_name = _connect_supersedes() if event.name == boot_name: self._cancel_older_chains((on_load_name, token), tracked.root_gen) + self._pending_hydrates.setdefault(token, {})[tracked.txid] = tracked + tracked.add_done_callback(functools.partial(self._forget_hydrate, token)) return super()._supersede_previous(token=token, event=event, tracked=tracked) + def _forget_hydrate(self, token: str, future: EventFuture) -> None: + """Stop tracking a completed hydration invocation. + + Its page-load descendants have their own supersession registration. + + Args: + token: The client token. + future: The completed hydration future. + """ + if (pending := self._pending_hydrates.get(token)) is not None: + pending.pop(future.txid, None) + if not pending: + del self._pending_hydrates[token] + + def _on_disconnect(self, token: str) -> None: + """Cancel hydration and unfinished page-load chains for a disconnected client. + + Args: + token: The disconnected client token. + """ + _, on_load_name = _connect_supersedes() + for future in self._pending_hydrates.pop(token, {}).values(): + future.cancel() + for future in self._superseded.pop((on_load_name, token), {}).values(): + future.cancel() + async def _execute_event( self, *, entry: EventQueueEntry, registered_handler: RegisteredEventHandler ) -> None: 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 1318bc33327..e48f89bdc83 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 @@ -458,6 +458,15 @@ async def enqueue( await queue.put(EventQueueEntry(event=event, ctx=ev_ctx)) return tracked + def _on_disconnect(self, token: str) -> None: + """Handle a client disconnect without cancelling independent event work. + + Subclasses may cancel work tied to the disconnected page's lifecycle. + + Args: + token: The disconnected client token. + """ + async def enqueue_many(self, token: str, *events: Event) -> Sequence[EventFuture]: """Enqueue multiple events to be processed. diff --git a/pyproject.toml b/pyproject.toml index 8b2081a5042..d09c3a1a296 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -32,7 +32,7 @@ dependencies = [ "starlette >=1.3.1", "typing_extensions >=4.13.0", "wrapt >=1.17.0,<2.4", - "reflex-base >= 0.9.12.dev0", + "reflex-base >= 0.9.13.dev0", "reflex-components-code >= 0.9.0", "reflex-components-core >= 0.9.6", "reflex-components-dataeditor >= 0.9.0", diff --git a/reflex/app.py b/reflex/app.py index 508e649ed27..7332128d2f2 100644 --- a/reflex/app.py +++ b/reflex/app.py @@ -2261,6 +2261,14 @@ def on_disconnect(self, sid: str) -> asyncio.Task | None: # Get token before cleaning up disconnect_token = self.sid_to_token.get(sid) if disconnect_token: + socket_record = self._token_manager.token_to_socket.get(disconnect_token) + if ( + socket_record is not None + and socket_record.sid == sid + and socket_record.instance_id == self._token_manager.instance_id + and (processor := self.app._event_processor) is not None + ): + processor._on_disconnect(disconnect_token) # Use async cleanup through token manager task = asyncio.create_task( self._token_manager.disconnect_token(disconnect_token, sid), diff --git a/tests/units/reflex_base/event/processor/test_base_state_processor.py b/tests/units/reflex_base/event/processor/test_base_state_processor.py index a14dbe6780a..f14c8cbf30b 100644 --- a/tests/units/reflex_base/event/processor/test_base_state_processor.py +++ b/tests/units/reflex_base/event/processor/test_base_state_processor.py @@ -22,7 +22,7 @@ import reflex as rx from reflex import event -from reflex.app import App +from reflex.app import App, EventNamespace from reflex.event import Event, EventSpec from reflex.istate.manager import StateManager from reflex.istate.manager.disk import StateManagerDisk @@ -1805,3 +1805,172 @@ async def slow_load(self): async with _read_back(real_base_state_processor, token) as root: assert (await root.get_state(SlowLoadState)).seen != "stale-finished" assert (await root.get_state(State)).is_hydrated is True + + +@pytest.mark.parametrize( + "processor_state_manager", ["in_process", "redis"], indirect=True +) +@pytest.mark.parametrize("self_chain", [False, True]) +async def test_disconnect_cancels_page_load_chain( + wired_app: App, + real_base_state_processor: BaseStateEventProcessor, + token: str, + self_chain: bool, +): + """Disconnect cancels a running load even after its root handler has returned. + + Args: + wired_app: The app sharing the processor's state manager. + real_base_state_processor: The real event processor. + token: The client token. + self_chain: Whether to reach the suspended handler through a self-chain. + """ + started = asyncio.Event() + cancelled = asyncio.Event() + invocations = 0 + + class DisconnectState(State): + @event + async def load(self): + """Suspend a load or its self-chained successor until cancellation. + + Yields: + A successor event or an empty update before suspending. + """ + nonlocal invocations + invocations += 1 + if self_chain and invocations == 1: + yield DisconnectState.load + return + yield + started.set() + try: + await asyncio.Event().wait() + finally: + cancelled.set() + + wired_app.add_page(lambda: rx.text("load"), route="/", on_load=DisconnectState.load) + wired_app._event_processor = real_base_state_processor + namespace = EventNamespace("/_event", wired_app) + await namespace._token_manager.link_token_to_sid(token, "sid") + + async with real_base_state_processor as processor: + stale = await processor.enqueue( + token, _client_event(OnLoadInternalState.on_load_internal(), _view("/")) + ) + await asyncio.wait_for(started.wait(), timeout=5) + cleanup = namespace.on_disconnect("sid") + assert cleanup is not None + await cleanup + await asyncio.wait_for(cancelled.wait(), timeout=1) + await asyncio.wait_for(stale.wait_all(), timeout=5) + assert invocations == (2 if self_chain else 1) + # Cancellation releases the lock so a later connection can use the state. + async with _read_back(processor, token) as root: + assert await root.get_state(DisconnectState) is not None + + +@pytest.mark.parametrize("disconnect", [False, True]) +async def test_disconnect_cancels_queued_hydrate( + wired_app: App, + real_base_state_processor: BaseStateEventProcessor, + token: str, + disconnect: bool, +): + """A disconnect before hydration starts must not launch a page-load chain. + + Args: + wired_app: The app sharing the processor's state manager. + real_base_state_processor: The real event processor. + token: The client token. + disconnect: Whether the client disconnects before hydration runs. + """ + loads = [] + + class DisconnectState(State): + @event + def load(self): + """Record that the page-load handler ran.""" + loads.append("loaded") + + wired_app.add_page(lambda: rx.text("load"), route="/", on_load=DisconnectState.load) + wired_app._event_processor = real_base_state_processor + namespace = EventNamespace("/_event", wired_app) + await namespace._token_manager.link_token_to_sid(token, "sid") + async with real_base_state_processor as processor: + future = await processor.enqueue( + token, _client_event(State.hydrate_and_load(), _view("/")) + ) + if disconnect: + cleanup = namespace.on_disconnect("sid") + assert cleanup is not None + await cleanup + assert future.cancelled() + else: + await asyncio.wait_for(future.wait_all(), timeout=5) + await namespace._token_manager.disconnect_all() + assert loads == ([] if disconnect else ["loaded"]) + assert not processor._pending_hydrates + + +async def test_disconnect_preserves_independent_events( + wired_app: App, + real_base_state_processor: BaseStateEventProcessor, + token: str, +): + """Independent foreground and background work survives another chain's disconnect. + + Args: + wired_app: The app sharing the processor's state manager. + real_base_state_processor: The real event processor. + token: The client token. + """ + started = {name: asyncio.Event() for name in ("foreground", "background", "other")} + release = asyncio.Event() + finished = [] + + class IndependentState(State): + @event + async def foreground(self, name: str): + """Wait for release and record the named foreground invocation. + + Args: + name: The invocation to record. + """ + started[name].set() + await release.wait() + finished.append(name) + + @event(background=True) + async def background(self): + """Wait for release independently of the page-load lifecycle.""" + started["background"].set() + await release.wait() + finished.append("background") + + async with real_base_state_processor as processor: + # Start the background event before the foreground handler holds the lock. + background = await processor.enqueue( + token, Event.from_event_type(IndependentState.background())[0] + ) + await asyncio.wait_for(started["background"].wait(), timeout=5) + foreground = await processor.enqueue( + token, Event.from_event_type(IndependentState.foreground("foreground"))[0] + ) + other = await processor.enqueue( + "other-token", + Event.from_event_type(IndependentState.foreground("other"))[0], + ) + await asyncio.wait_for( + asyncio.gather(*(signal.wait() for signal in started.values())), timeout=5 + ) + processor._on_disconnect(token) + assert not any(future.cancelled() for future in (background, foreground, other)) + release.set() + await asyncio.wait_for( + asyncio.gather( + *(future.wait_all() for future in (background, foreground, other)) + ), + timeout=5, + ) + assert sorted(finished) == ["background", "foreground", "other"] diff --git a/tests/units/test_app.py b/tests/units/test_app.py index 85eaeca933b..207dc683095 100644 --- a/tests/units/test_app.py +++ b/tests/units/test_app.py @@ -4836,6 +4836,37 @@ async def event_namespace_with_processor_mock() -> AsyncGenerator[EventNamespace await event_namespace._token_manager.disconnect_all() +@pytest.mark.parametrize("owner", ["current", "replacement", "remote", "missing"]) +async def test_disconnect_only_cancels_current_local_socket( + event_namespace_with_processor_mock: EventNamespace, + owner: str, +): + """A stale disconnect must not cancel a replacement connection's page load. + + Args: + event_namespace_with_processor_mock: Namespace with a mocked processor. + owner: The socket currently owning the client token. + """ + namespace = event_namespace_with_processor_mock + manager = namespace._token_manager + manager.sid_to_token["old-sid"] = "token" + if owner != "missing": + manager.token_to_socket["token"] = SocketRecord( + instance_id="remote" if owner == "remote" else manager.instance_id, + sid="new-sid" if owner == "replacement" else "old-sid", + ) + manager.disconnect_token = AsyncMock() + cleanup = namespace.on_disconnect("old-sid") + assert cleanup is not None + await cleanup + cancel = cast(Mock, namespace.app.event_processor._on_disconnect) + if owner == "current": + cancel.assert_called_once_with("token") + else: + cancel.assert_not_called() + assert namespace.on_disconnect("unknown-sid") is None + + def _connect_environ(token: str) -> dict[str, Any]: return { "QUERY_STRING": f"token={token}", diff --git a/tests/units/test_check_min_deps.py b/tests/units/test_check_min_deps.py index 34081b559a5..78a074cbe12 100644 --- a/tests/units/test_check_min_deps.py +++ b/tests/units/test_check_min_deps.py @@ -434,7 +434,7 @@ def test_build_wheelhouse_builds_every_source_once( fake_run = _FakeRun( built={ # Both satisfy what the root package declares, so neither is rebuilt. - "reflex-base": "0.9.12.post1.dev0+abc1234", + "reflex-base": "0.9.13.dev0+abc1234", "reflex-hosting-cli": "0.1.71.post1.dev0+abc1234", } ) @@ -451,7 +451,7 @@ def test_build_wheelhouse_builds_every_source_once( ) assert detail is None assert check_min_deps._workspace_pins(package, versions)[0] == [ - "reflex-base==0.9.12.post1.dev0+abc1234", + "reflex-base==0.9.13.dev0+abc1234", "reflex-hosting-cli==0.1.71.post1.dev0+abc1234", ] From fcb2ce913589681e627d5269ff01294687cac08a Mon Sep 17 00:00:00 2001 From: Krishna Date: Mon, 5 Oct 2026 02:48:32 +0530 Subject: [PATCH 2/3] fix(events): track rehydrated page-load chains --- .../event/processor/base_state_processor.py | 14 ++++++- .../event/processor/event_processor.py | 18 ++++----- .../processor/test_base_state_processor.py | 38 ++++++++++++++++++- 3 files changed, 58 insertions(+), 12 deletions(-) diff --git a/packages/reflex-base/src/reflex_base/event/processor/base_state_processor.py b/packages/reflex-base/src/reflex_base/event/processor/base_state_processor.py index 338592251f9..614fc64cc04 100644 --- a/packages/reflex-base/src/reflex_base/event/processor/base_state_processor.py +++ b/packages/reflex-base/src/reflex_base/event/processor/base_state_processor.py @@ -392,6 +392,7 @@ async def _rehydrate(self, root_state: BaseState): Args: root_state: The root state to rehydrate. """ + from reflex.event import Event from reflex.state import OnLoadInternalState, State if type(root_state) is not State: @@ -418,12 +419,23 @@ async def _rehydrate(self, root_state: BaseState): root_state=root_state, ) return + ctx = EventContext.get() + parent = self._futures[ctx.txid] + first_child = len(parent.children) + handler = OnLoadInternalState.event_handlers["on_load_internal"] await process_event( - handler=OnLoadInternalState.event_handlers["on_load_internal"], + handler=handler, payload={}, state=await root_state.get_state(OnLoadInternalState), root_state=root_state, ) + # Inline rehydration bypasses enqueue's page-load registration. Track + # only its children so the ordinary event remains independent. + load_event = Event(name=format_event_handler(handler)) + for index in range(first_child, len(parent.children)): + self._supersede_previous( + token=ctx.token, event=load_event, tracked=parent.children[index] + ) def _supersede_previous( self, *, token: str, event: Event, tracked: EventFuture 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 e48f89bdc83..364a3a12141 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 @@ -592,15 +592,15 @@ def _try_clean_future(self, future: EventFuture) -> None: # type: ignore[overri return parent = future.parent self._futures.pop(future.txid, None) - if ( - (key := future.supersede_key) is not None - and (slot := self._superseded.get(key)) is not None - and slot.get(future.txid) is future - and future.all_done() - ): - del slot[future.txid] - if not slot: - del self._superseded[key] + if future.supersede_key is not None and future.all_done(): + # Inline page loads can also have their own supersession key. + for key in future.covered_supersede_keys: + if (slot := self._superseded.get(key)) is not None and slot.get( + future.txid + ) is future: + del slot[future.txid] + if not slot: + del self._superseded[key] if parent is None or not parent.txid: return future = parent diff --git a/tests/units/reflex_base/event/processor/test_base_state_processor.py b/tests/units/reflex_base/event/processor/test_base_state_processor.py index f14c8cbf30b..6c2f65162e2 100644 --- a/tests/units/reflex_base/event/processor/test_base_state_processor.py +++ b/tests/units/reflex_base/event/processor/test_base_state_processor.py @@ -1811,11 +1811,15 @@ async def slow_load(self): "processor_state_manager", ["in_process", "redis"], indirect=True ) @pytest.mark.parametrize("self_chain", [False, True]) +@pytest.mark.parametrize("rehydrate", [False, True]) +@pytest.mark.parametrize("supersedes", [False, True]) async def test_disconnect_cancels_page_load_chain( wired_app: App, real_base_state_processor: BaseStateEventProcessor, token: str, self_chain: bool, + rehydrate: bool, + supersedes: bool, ): """Disconnect cancels a running load even after its root handler has returned. @@ -1824,6 +1828,8 @@ async def test_disconnect_cancels_page_load_chain( real_base_state_processor: The real event processor. token: The client token. self_chain: Whether to reach the suspended handler through a self-chain. + rehydrate: Whether an ordinary routed event triggers the page load. + supersedes: Whether the user's load handler also supersedes earlier calls. """ started = asyncio.Event() cancelled = asyncio.Event() @@ -1831,6 +1837,10 @@ async def test_disconnect_cancels_page_load_chain( class DisconnectState(State): @event + def ordinary(self): + """Trigger compatibility hydration on a fresh client state.""" + + @event(supersedes=supersedes) async def load(self): """Suspend a load or its self-chained successor until cancellation. @@ -1856,7 +1866,13 @@ async def load(self): async with real_base_state_processor as processor: stale = await processor.enqueue( - token, _client_event(OnLoadInternalState.on_load_internal(), _view("/")) + token, + _client_event( + DisconnectState.ordinary() + if rehydrate + else OnLoadInternalState.on_load_internal(), + _view("/"), + ), ) await asyncio.wait_for(started.wait(), timeout=5) cleanup = namespace.on_disconnect("sid") @@ -1868,6 +1884,8 @@ async def load(self): # Cancellation releases the lock so a later connection can use the state. async with _read_back(processor, token) as root: assert await root.get_state(DisconnectState) is not None + await asyncio.sleep(0) + assert not processor._superseded @pytest.mark.parametrize("disconnect", [False, True]) @@ -1913,10 +1931,12 @@ def load(self): assert not processor._pending_hydrates +@pytest.mark.parametrize("routed", [False, True]) async def test_disconnect_preserves_independent_events( wired_app: App, real_base_state_processor: BaseStateEventProcessor, token: str, + routed: bool, ): """Independent foreground and background work survives another chain's disconnect. @@ -1924,12 +1944,19 @@ async def test_disconnect_preserves_independent_events( wired_app: The app sharing the processor's state manager. real_base_state_processor: The real event processor. token: The client token. + routed: Whether the foreground event also schedules compatibility page loads. """ started = {name: asyncio.Event() for name in ("foreground", "background", "other")} release = asyncio.Event() finished = [] + loads = [] class IndependentState(State): + @event + def load(self): + """Record whether the disconnected page-load work ran.""" + loads.append("loaded") + @event async def foreground(self, name: str): """Wait for release and record the named foreground invocation. @@ -1948,6 +1975,9 @@ async def background(self): await release.wait() finished.append("background") + wired_app.add_page( + lambda: rx.text("load"), route="/", on_load=IndependentState.load + ) async with real_base_state_processor as processor: # Start the background event before the foreground handler holds the lock. background = await processor.enqueue( @@ -1955,7 +1985,10 @@ async def background(self): ) await asyncio.wait_for(started["background"].wait(), timeout=5) foreground = await processor.enqueue( - token, Event.from_event_type(IndependentState.foreground("foreground"))[0] + token, + _client_event(IndependentState.foreground("foreground"), _view("/")) + if routed + else Event.from_event_type(IndependentState.foreground("foreground"))[0], ) other = await processor.enqueue( "other-token", @@ -1974,3 +2007,4 @@ async def background(self): timeout=5, ) assert sorted(finished) == ["background", "foreground", "other"] + assert not loads From a846d358d3195d73bf9d2e55eff55392eb6e26a9 Mon Sep 17 00:00:00 2001 From: Krishna Date: Mon, 5 Oct 2026 03:00:26 +0530 Subject: [PATCH 3/3] fix(events): guard page loads during hydration --- .../event/processor/base_state_processor.py | 66 +++++++++------- .../processor/test_base_state_processor.py | 76 +++++++++++++++++++ 2 files changed, 115 insertions(+), 27 deletions(-) diff --git a/packages/reflex-base/src/reflex_base/event/processor/base_state_processor.py b/packages/reflex-base/src/reflex_base/event/processor/base_state_processor.py index 614fc64cc04..ee01b5daae7 100644 --- a/packages/reflex-base/src/reflex_base/event/processor/base_state_processor.py +++ b/packages/reflex-base/src/reflex_base/event/processor/base_state_processor.py @@ -404,38 +404,50 @@ async def _rehydrate(self, root_state: BaseState): if routeless and root_state.is_hydrated: return - await process_event( - handler=State.event_handlers["hydrate"], - payload={}, - state=root_state, - root_state=root_state, - ) - if routeless: - # No page to load, but hydration still has to finish. + ctx = EventContext.get() + # Disconnect cancels this marker without cancelling the ordinary event. + pending = EventFuture(txid=ctx.txid) + self._pending_hydrates.setdefault(ctx.token, {})[pending.txid] = pending + try: await process_event( - handler=State.event_handlers["set_is_hydrated"], - payload={"value": True}, + handler=State.event_handlers["hydrate"], + payload={}, state=root_state, root_state=root_state, ) - return - ctx = EventContext.get() - parent = self._futures[ctx.txid] - first_child = len(parent.children) - handler = OnLoadInternalState.event_handlers["on_load_internal"] - await process_event( - handler=handler, - payload={}, - state=await root_state.get_state(OnLoadInternalState), - root_state=root_state, - ) - # Inline rehydration bypasses enqueue's page-load registration. Track - # only its children so the ordinary event remains independent. - load_event = Event(name=format_event_handler(handler)) - for index in range(first_child, len(parent.children)): - self._supersede_previous( - token=ctx.token, event=load_event, tracked=parent.children[index] + if routeless: + # No page to load, but hydration still has to finish. + await process_event( + handler=State.event_handlers["set_is_hydrated"], + payload={"value": True}, + state=root_state, + root_state=root_state, + ) + return + if pending.cancelled(): + return + parent = self._futures[ctx.txid] + first_child = len(parent.children) + handler = OnLoadInternalState.event_handlers["on_load_internal"] + await process_event( + handler=handler, + payload={}, + state=await root_state.get_state(OnLoadInternalState), + root_state=root_state, ) + # Inline rehydration bypasses enqueue's page-load registration. Track + # only its children so the ordinary event remains independent. + load_event = Event(name=format_event_handler(handler)) + for index in range(first_child, len(parent.children)): + child = parent.children[index] + if pending.cancelled(): + child.cancel() + else: + self._supersede_previous( + token=ctx.token, event=load_event, tracked=child + ) + finally: + self._forget_hydrate(ctx.token, pending) def _supersede_previous( self, *, token: str, event: Event, tracked: EventFuture diff --git a/tests/units/reflex_base/event/processor/test_base_state_processor.py b/tests/units/reflex_base/event/processor/test_base_state_processor.py index 6c2f65162e2..80ca5a389fa 100644 --- a/tests/units/reflex_base/event/processor/test_base_state_processor.py +++ b/tests/units/reflex_base/event/processor/test_base_state_processor.py @@ -1888,6 +1888,82 @@ async def load(self): assert not processor._superseded +@pytest.mark.parametrize( + "processor_state_manager", ["in_process", "redis"], indirect=True +) +@pytest.mark.parametrize("disconnect", [False, True]) +async def test_disconnect_during_inline_hydration( + wired_app: App, + real_base_state_processor: BaseStateEventProcessor, + token: str, + disconnect: bool, +): + """A disconnect during the hydration emit prevents late page-load work. + + Args: + wired_app: The app sharing the processor's state manager. + real_base_state_processor: The real event processor. + token: The client token. + disconnect: Whether to disconnect while sending the hydration update. + """ + sending = asyncio.Event() + release = asyncio.Event() + invocations = [] + + class InlineHydrationState(State): + @event + def ordinary(self): + """Record that independent work still runs after hydration.""" + invocations.append("ordinary") + + @event + def load(self): + """Record whether the page load ran.""" + invocations.append("load") + + async def pause_emit(token: str, delta: Mapping[str, Mapping[str, Any]]) -> None: + """Pause sending the hydration snapshot. + + Args: + token: The client token. + delta: The state update. + """ + sending.set() + await release.wait() + + processor = real_base_state_processor + assert processor._root_context is not None + processor._root_context = dataclasses.replace( + processor._root_context, emit_delta_impl=pause_emit + ) + wired_app.add_page( + lambda: rx.text("load"), route="/", on_load=InlineHydrationState.load + ) + wired_app._event_processor = processor + namespace = EventNamespace("/_event", wired_app) + await namespace._token_manager.link_token_to_sid(token, "sid") + + async with processor: + ordinary = await processor.enqueue( + token, _client_event(InlineHydrationState.ordinary(), _view("/")) + ) + try: + await asyncio.wait_for(sending.wait(), timeout=5) + if disconnect: + cleanup = namespace.on_disconnect("sid") + assert cleanup is not None + await cleanup + await namespace._token_manager.link_token_to_sid(token, "new-sid") + finally: + release.set() + await asyncio.wait_for(ordinary.wait_all(), timeout=5) + await asyncio.sleep(0) + assert not processor._pending_hydrates + assert not processor._superseded + + assert invocations == (["ordinary"] if disconnect else ["ordinary", "load"]) + + @pytest.mark.parametrize("disconnect", [False, True]) async def test_disconnect_cancels_queued_hydrate( wired_app: App,