From e85c14117ce066efda2c8e7e8eb6e4164979b4cc Mon Sep 17 00:00:00 2001 From: Krishna Date: Sun, 4 Oct 2026 22:16:02 +0530 Subject: [PATCH 1/3] fix(state): persist cancelled redis handler updates --- docs/events/yield_events.md | 2 + news/+persist-cancelled-redis-state.bugfix.md | 1 + reflex/istate/manager/redis.py | 49 +++++-- tests/units/istate/manager/test_redis.py | 121 +++++++++++++++++- .../processor/test_base_state_processor.py | 73 +++++++++++ 5 files changed, 236 insertions(+), 10 deletions(-) create mode 100644 news/+persist-cancelled-redis-state.bugfix.md diff --git a/docs/events/yield_events.md b/docs/events/yield_events.md index ce21804006d..f26a3bedb4c 100644 --- a/docs/events/yield_events.md +++ b/docs/events/yield_events.md @@ -8,6 +8,8 @@ A regular event handler will send a `StateUpdate` when it has finished running. To do so, we can use the Python keyword `yield`. For every yield inside the function, a `StateUpdate` will be sent to the frontend with the changes up to this point in the execution of the event handler. +When using Redis, cancelling a foreground event handler preserves its state changes as long as it still holds the state lock. This includes updates already yielded to the frontend before another invocation of an `@rx.event(supersedes=True)` handler cancels it. + This example below shows how to yield 100 updates to the UI. ```python demo exec diff --git a/news/+persist-cancelled-redis-state.bugfix.md b/news/+persist-cancelled-redis-state.bugfix.md new file mode 100644 index 00000000000..a03edbf02be --- /dev/null +++ b/news/+persist-cancelled-redis-state.bugfix.md @@ -0,0 +1 @@ +Preserve state changes made by cancelled foreground event handlers when using Redis, including updates already yielded to the client before a superseding event. diff --git a/reflex/istate/manager/redis.py b/reflex/istate/manager/redis.py index 36037348be7..06601b6b746 100644 --- a/reflex/istate/manager/redis.py +++ b/reflex/istate/manager/redis.py @@ -517,6 +517,32 @@ def _collect_state_writes( stack.extend(state.substates.values()) return writes + @contextlib.asynccontextmanager + async def _modify_state_with_lock( + self, + token: StateToken[TOKEN_TYPE], + lock_id: bytes, + **context: Unpack[StateModificationContext], + ) -> AsyncIterator[TOKEN_TYPE]: + """Load and save state under an existing lock, including on cancellation. + + Args: + token: The token to modify the state for. + lock_id: The lock that must still be held when saving. + context: The state modification context. + + Yields: + The state for the token. + """ + state = await self.get_state(token) + try: + yield state + except asyncio.CancelledError: + # A handler may already have yielded these writes to the client. + await self.set_state(token, state, lock_id=lock_id, **context) + raise + await self.set_state(token, state, lock_id=lock_id, **context) + @contextlib.asynccontextmanager async def _try_modify_state( self, token: StateToken[TOKEN_TYPE], **context: Unpack[StateModificationContext] @@ -533,10 +559,11 @@ async def _try_modify_state( event_name = event.name if (event := context.get("event")) is not None else None if not self._oplock_enabled: # OpLock is disabled, get a fresh lock, write, and release. - async with self._lock(token, event_name=event_name) as lock_id: - state = await self.get_state(token) + async with ( + self._lock(token, event_name=event_name) as lock_id, + self._modify_state_with_lock(token, lock_id, **context) as state, + ): yield state - await self.set_state(token, state, lock_id=lock_id, **context) return # Opportunistically reuse existing lock. @@ -566,10 +593,11 @@ async def _try_modify_state( logger.debug( f"{SMR} [{time.monotonic() - start:.3f}] {lock_key} has contention, not leasing" ) - async with lock_held_ctx: - state = await self.get_state(token) + async with ( + lock_held_ctx, + self._modify_state_with_lock(token, lock_id, **context) as state, + ): yield state - await self.set_state(token, state, lock_id=lock_id, **context) return # Create the lease break task since we got the lock. @@ -592,10 +620,13 @@ async def _try_modify_state( logger.debug( f"{SMR} [{time.monotonic() - start:.3f}] {lock_key} holding lock {lock_id.decode()}, {new_lease_task=} already exited, doing single update..." ) - async with lock_held_ctx: - state = await self.get_state(token) + async with ( + lock_held_ctx, + self._modify_state_with_lock( + token, lock_id, **context + ) as state, + ): yield state - await self.set_state(token, state, lock_id=lock_id, **context) return elif self._debug_enabled: logger.debug( diff --git a/tests/units/istate/manager/test_redis.py b/tests/units/istate/manager/test_redis.py index 692a446fb2a..56c6d3d3e75 100644 --- a/tests/units/istate/manager/test_redis.py +++ b/tests/units/istate/manager/test_redis.py @@ -13,7 +13,7 @@ import pytest import pytest_asyncio -from reflex_base.utils.exceptions import EnvironmentVarValueError +from reflex_base.utils.exceptions import EnvironmentVarValueError, LockExpiredError from reflex.istate.manager.redis import ( StateManagerRedis, @@ -35,6 +35,8 @@ class RedisTestState(BaseState): class SubState1(RedisTestState): """A test substate for redis state manager tests.""" + child_value: str = "" + class SubState2(RedisTestState): """A test substate for redis state manager tests.""" @@ -230,6 +232,123 @@ async def test_modify( assert final_state.count == 3 +@pytest.mark.parametrize( + "locking_path", ["disabled", "contended", "lease_ended", "cached"] +) +async def test_modify_cancelled_persists_state( + state_manager_redis: StateManagerRedis, + monkeypatch: pytest.MonkeyPatch, + locking_path: str, +): + """Keep mutations and cancellation cleanup visible to the next handler. + + Args: + state_manager_redis: The StateManagerRedis to test. + monkeypatch: The pytest monkeypatch fixture. + locking_path: The Redis locking path to exercise. + """ + manager = state_manager_redis + manager._oplock_enabled = locking_path != "disabled" + if locking_path == "contended": + monkeypatch.setattr(manager, "_n_lock_contenders", AsyncMock(return_value=1)) + elif locking_path == "lease_ended": + monkeypatch.setattr( + manager, "_create_lease_break_task", AsyncMock(return_value=None) + ) + elif locking_path == "cached": + await _subscribed(manager) + + token = BaseStateToken(ident=str(uuid.uuid4()), cls=RedisTestState) + started = asyncio.Event() + + async def modify(): + """Mutate state before cancellation and again during handler cleanup.""" + async with manager.modify_state(token) as state: + state.count = 1 + state.substates[SubState1.get_name()].child_value = "child update" + started.set() + try: + await asyncio.Event().wait() + finally: + state.foo = "cancelled" + + task = asyncio.create_task(modify()) + try: + await asyncio.wait_for(started.wait(), timeout=5) + task.cancel("superseded") + with pytest.raises(asyncio.CancelledError, match="superseded"): + await task + assert task.cancelled() + + async with manager.modify_state(token) as state: + assert isinstance(state, RedisTestState) + assert state.count == 1 + assert state.foo == "cancelled" + child = state.substates[SubState1.get_name()] + assert isinstance(child, SubState1) + assert child.child_value == "child update" + state.count += 1 + + # Closing flushes a cached lease; a fresh read must see persisted state. + await manager.close() + persisted = await manager.get_state(token) + assert isinstance(persisted, RedisTestState) + assert persisted.count == 2 + assert persisted.foo == "cancelled" + persisted_child = persisted.substates[SubState1.get_name()] + assert isinstance(persisted_child, SubState1) + assert persisted_child.child_value == "child update" + finally: + task.cancel() + with contextlib.suppress(asyncio.CancelledError): + await task + + +async def test_modify_exception_does_not_persist_state( + state_manager_redis: StateManagerRedis, +): + """Keep the existing write-discarding behavior for non-cancellation errors. + + Args: + state_manager_redis: The StateManagerRedis to test. + """ + manager = state_manager_redis + manager._oplock_enabled = False + token = BaseStateToken(ident=str(uuid.uuid4()), cls=RedisTestState) + with pytest.raises(RuntimeError, match="handler failed"): + async with manager.modify_state(token) as state: + state.count = 1 + msg = "handler failed" + raise RuntimeError(msg) + persisted = await manager.get_state(token) + assert isinstance(persisted, RedisTestState) + assert persisted.count == 0 + + +async def test_modify_cancelled_does_not_save_after_lock_expiry( + state_manager_redis: StateManagerRedis, +): + """Cancellation must not bypass the lock fence when saving state. + + Args: + state_manager_redis: The StateManagerRedis to test. + """ + manager = state_manager_redis + manager._oplock_enabled = False + token = BaseStateToken(ident=str(uuid.uuid4()), cls=RedisTestState) + async with manager.modify_state(token) as state: + state.count = 1 + + with pytest.raises(LockExpiredError): + async with manager.modify_state(token) as state: + state.count = 2 + await manager.redis.delete(manager._lock_key(token)) + raise asyncio.CancelledError + persisted = await manager.get_state(token) + assert isinstance(persisted, RedisTestState) + assert persisted.count == 1 + + async def test_get_state_discards_unpicklable_state( state_manager_redis: StateManagerRedis, app_classes_module: ModuleType, 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..b18ed742b0e 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 @@ -1024,6 +1024,79 @@ def ping(self): assert (await root.get_state(State)).is_hydrated is True +@pytest.mark.parametrize( + "processor_state_manager", ["in_process", "redis"], indirect=True +) +async def test_superseded_handler_preserves_yielded_state( + wired_app: App, + real_base_state_processor: BaseStateEventProcessor, + processor_state_manager: StateManager, + emitted_deltas: list[tuple[str, Mapping[str, Mapping[str, Any]]]], + token: str, +): + """A replacement handler keeps writes already yielded by its cancelled predecessor. + + Args: + wired_app: The App wired to the processor's state manager. + real_base_state_processor: The unmocked event processor. + processor_state_manager: The state manager backing the processor. + emitted_deltas: The deltas sent by the processor. + token: The client token. + """ + if isinstance(processor_state_manager, StateManagerRedis): + processor_state_manager._oplock_enabled = False + started = asyncio.Event() + + class RefreshState(State): + """Track writes from superseding foreground handlers.""" + + entries: list[str] = [] + + @event(supersedes=True) + async def refresh(self, tag: str): + """Record work before yielding and after completion or cancellation. + + Args: + tag: The invocation to record. + + Yields: + The state update before waiting. + """ + self.entries = [*self.entries, f"{tag}:start"] + yield + if tag == "A": + started.set() + try: + await asyncio.Event().wait() + except asyncio.CancelledError: + self.entries = [*self.entries, f"{tag}:cancelled"] + raise + self.entries = [*self.entries, f"{tag}:done"] + + wired_app.add_page(lambda: rx.text("refresh"), route="/") + async with real_base_state_processor as processor: + first = await processor.enqueue( + token, _client_event(RefreshState.refresh("A"), _view("/")) + ) + await asyncio.wait_for(started.wait(), timeout=5) + await _send( + processor, token, _client_event(RefreshState.refresh("B"), _view("/")) + ) + assert first.done() + + expected = ["A:start", "A:cancelled", "B:start", "B:done"] + async with _read_back(real_base_state_processor, token) as root: + assert (await root.get_state(RefreshState)).entries == expected + entry_deltas = [ + fields["entries" + FIELD_MARKER] + for _, delta in emitted_deltas + if "entries" + FIELD_MARKER + in (fields := delta.get(RefreshState.get_full_name(), {})) + ] + assert ["A:start"] in entry_deltas + assert entry_deltas[-1] == expected + + async def test_rehydrate_after_expiry_does_not_reload_the_previous_route( wired_app: App, real_base_state_processor: BaseStateEventProcessor, From 055d695183963a7a8df435dc2d380595e5a12943 Mon Sep 17 00:00:00 2001 From: Krishna Date: Mon, 5 Oct 2026 02:45:16 +0530 Subject: [PATCH 2/3] fix(state): finish saves through repeated cancellation --- reflex/istate/manager/redis.py | 9 +++- tests/units/istate/manager/test_redis.py | 57 ++++++++++++++++++++++++ 2 files changed, 65 insertions(+), 1 deletion(-) diff --git a/reflex/istate/manager/redis.py b/reflex/istate/manager/redis.py index 06601b6b746..05fa732fb7e 100644 --- a/reflex/istate/manager/redis.py +++ b/reflex/istate/manager/redis.py @@ -539,7 +539,14 @@ async def _modify_state_with_lock( yield state except asyncio.CancelledError: # A handler may already have yielded these writes to the client. - await self.set_state(token, state, lock_id=lock_id, **context) + save = asyncio.create_task( + self.set_state(token, state, lock_id=lock_id, **context) + ) + # Keep the lock until the save finishes, even during shutdown. + while not save.done(): + with contextlib.suppress(asyncio.CancelledError): + await asyncio.shield(save) + save.result() raise await self.set_state(token, state, lock_id=lock_id, **context) diff --git a/tests/units/istate/manager/test_redis.py b/tests/units/istate/manager/test_redis.py index 56c6d3d3e75..ed727a1afd7 100644 --- a/tests/units/istate/manager/test_redis.py +++ b/tests/units/istate/manager/test_redis.py @@ -304,6 +304,63 @@ async def modify(): await task +async def test_repeated_cancellation_waits_for_state_save( + state_manager_redis: StateManagerRedis, + monkeypatch: pytest.MonkeyPatch, +): + """Keep the Redis lock until a repeatedly cancelled handler finishes saving. + + Args: + state_manager_redis: The StateManagerRedis to test. + monkeypatch: The pytest monkeypatch fixture. + """ + manager = state_manager_redis + manager._oplock_enabled = False + token = BaseStateToken(ident=str(uuid.uuid4()), cls=RedisTestState) + started = asyncio.Event() + saving = asyncio.Event() + release_save = asyncio.Event() + original_set_state = manager.set_state + + async def save(*args, **kwargs): + """Pause the state write while the handler is cancelled again.""" + saving.set() + await release_save.wait() + await original_set_state(*args, **kwargs) + + monkeypatch.setattr(manager, "set_state", save) + + async def modify(): + """Mutate state and suspend until superseded.""" + async with manager.modify_state(token) as state: + state.count = 1 + started.set() + await asyncio.Event().wait() + + task = asyncio.create_task(modify()) + try: + await asyncio.wait_for(started.wait(), timeout=5) + task.cancel("superseded") + await asyncio.wait_for(saving.wait(), timeout=5) + for _ in range(2): + task.cancel("shutdown") + await asyncio.sleep(0) + assert not task.done() + assert await manager.redis.get(manager._lock_key(token)) is not None + release_save.set() + with pytest.raises(asyncio.CancelledError, match="superseded"): + await task + persisted = await manager.get_state(token) + assert isinstance(persisted, RedisTestState) + assert persisted.count == 1 + assert await manager.redis.get(manager._lock_key(token)) is None + finally: + release_save.set() + task.cancel() + with contextlib.suppress(asyncio.CancelledError): + await task + + async def test_modify_exception_does_not_persist_state( state_manager_redis: StateManagerRedis, ): From 6df001171e6a2f5863c44e7d2fd6b416c2f8caa4 Mon Sep 17 00:00:00 2001 From: Farhan Date: Thu, 29 Oct 2026 21:47:18 +0500 Subject: [PATCH 3/3] fix(state): keep redis writes of handlers that raise A handler exception dropped its writes on the non-cached Redis paths, while memory, disk and the Redis cached-lease path keep them. Save on any handler exception, not only on cancellation, so every path matches. Document that a write after the last yield reaches the frontend only with a later update on Redis. --- docs/events/yield_events.md | 2 +- news/+persist-cancelled-redis-state.bugfix.md | 2 +- reflex/istate/manager/redis.py | 4 ++-- tests/units/istate/manager/test_redis.py | 6 +++--- 4 files changed, 7 insertions(+), 7 deletions(-) diff --git a/docs/events/yield_events.md b/docs/events/yield_events.md index f26a3bedb4c..a238586b055 100644 --- a/docs/events/yield_events.md +++ b/docs/events/yield_events.md @@ -8,7 +8,7 @@ A regular event handler will send a `StateUpdate` when it has finished running. To do so, we can use the Python keyword `yield`. For every yield inside the function, a `StateUpdate` will be sent to the frontend with the changes up to this point in the execution of the event handler. -When using Redis, cancelling a foreground event handler preserves its state changes as long as it still holds the state lock. This includes updates already yielded to the frontend before another invocation of an `@rx.event(supersedes=True)` handler cancels it. +A foreground event handler that raises an exception or is cancelled keeps the state changes it made, as long as it still holds the state lock. This includes updates already yielded to the frontend before another invocation of an `@rx.event(supersedes=True)` handler cancels it. When using Redis, a change made after the handler's last `yield` is stored, but the frontend receives it only when a later update changes the same var, or after a page reload. This example below shows how to yield 100 updates to the UI. diff --git a/news/+persist-cancelled-redis-state.bugfix.md b/news/+persist-cancelled-redis-state.bugfix.md index a03edbf02be..6463ecbc822 100644 --- a/news/+persist-cancelled-redis-state.bugfix.md +++ b/news/+persist-cancelled-redis-state.bugfix.md @@ -1 +1 @@ -Preserve state changes made by cancelled foreground event handlers when using Redis, including updates already yielded to the client before a superseding event. +Preserve state changes made by foreground event handlers that raise or are cancelled when using Redis, including updates already yielded to the client before a superseding event, matching the other state managers. diff --git a/reflex/istate/manager/redis.py b/reflex/istate/manager/redis.py index ffbb07751dd..6e1fff728af 100644 --- a/reflex/istate/manager/redis.py +++ b/reflex/istate/manager/redis.py @@ -575,7 +575,7 @@ async def _modify_state_with_lock( lock_id: bytes, **context: Unpack[StateModificationContext], ) -> AsyncIterator[TOKEN_TYPE]: - """Load and save state under an existing lock, including on cancellation. + """Load and save state under an existing lock, also when the handler fails. Args: token: The token to modify the state for. @@ -588,7 +588,7 @@ async def _modify_state_with_lock( state = await self.get_state(token) try: yield state - except asyncio.CancelledError: + except (Exception, asyncio.CancelledError): # A handler may already have yielded these writes to the client. save = asyncio.create_task( self.set_state(token, state, lock_id=lock_id, **context) diff --git a/tests/units/istate/manager/test_redis.py b/tests/units/istate/manager/test_redis.py index 47f89b615fa..d380172b4ee 100644 --- a/tests/units/istate/manager/test_redis.py +++ b/tests/units/istate/manager/test_redis.py @@ -410,10 +410,10 @@ async def modify(): await task -async def test_modify_exception_does_not_persist_state( +async def test_modify_exception_persists_state( state_manager_redis: StateManagerRedis, ): - """Keep the existing write-discarding behavior for non-cancellation errors. + """Keep the writes of a handler that raises, like the other state managers. Args: state_manager_redis: The StateManagerRedis to test. @@ -428,7 +428,7 @@ async def test_modify_exception_does_not_persist_state( raise RuntimeError(msg) persisted = await manager.get_state(token) assert isinstance(persisted, RedisTestState) - assert persisted.count == 0 + assert persisted.count == 1 async def test_modify_cancelled_does_not_save_after_lock_expiry(