fix(state): persist cancelled redis handler updates

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.
 
+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.
 
 ```python demo exec

news/+persist-cancelled-redis-state.bugfix.md

@@ -0,0 +1 @@
+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.

reflex/istate/manager/redis.py

@@ -568,6 +568,39 @@ 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, also when the handler fails.
+
+        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 (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)
+            )
+            # 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)
+
     @contextlib.asynccontextmanager
     async def _try_modify_state(
         self, token: StateToken[TOKEN_TYPE], **context: Unpack[StateModificationContext]
@@ -584,10 +617,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.
@@ -660,10 +694,11 @@ async def _try_modify_state(
             raise
 
         if 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
         # Nothing left to release once a lease took the lock over.
         await lock_held_ctx.aclose()

tests/units/istate/manager/test_redis.py

@@ -39,6 +39,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."""
@@ -279,6 +281,180 @@ 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_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_persists_state(
+    state_manager_redis: StateManagerRedis,
+):
+    """Keep the writes of a handler that raises, like the other state managers.
+
+    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 == 1
+
+
+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,

tests/units/reflex_base/event/processor/test_base_state_processor.py

@@ -1418,6 +1418,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,