diff --git a/news/+7478-refresh-expired-background-state.bugfix.md b/news/+7478-refresh-expired-background-state.bugfix.md new file mode 100644 index 00000000000..30f0d3a1d08 --- /dev/null +++ b/news/+7478-refresh-expired-background-state.bugfix.md @@ -0,0 +1 @@ +Mark recreated state trees dirty so background-task completion deltas refresh expired substate values in the client. diff --git a/reflex/istate/manager/__init__.py b/reflex/istate/manager/__init__.py index 156c5af687e..237703a1f8d 100644 --- a/reflex/istate/manager/__init__.py +++ b/reflex/istate/manager/__init__.py @@ -5,7 +5,8 @@ import logging import sys from abc import ABC, abstractmethod -from collections.abc import AsyncIterator +from collections.abc import AsyncIterator, Iterator +from contextvars import ContextVar from types import TracebackType from typing import TYPE_CHECKING, Generic, TypedDict, overload @@ -19,6 +20,9 @@ from reflex.utils import console, prerequisites logger = logging.getLogger(__name__) +_STATE_RECREATION_TRACKER: ContextVar[list["BaseState"] | None] = ContextVar( + "state_recreation_tracker", default=None +) if TYPE_CHECKING: from reflex.state import BaseState @@ -333,6 +337,56 @@ async def __aexit__( return await self._state_context.__aexit__(exc_type, exc, tb) or suppressed +@contextlib.contextmanager +def _track_state_recreation() -> Iterator[list["BaseState"]]: + """Track state instances recreated inside the current async context. + + Yields: + The state instances reported as recreated before the context exits. + """ + recreated_states: list[BaseState] = [] + tracker_token = _STATE_RECREATION_TRACKER.set(recreated_states) + try: + yield recreated_states + finally: + _STATE_RECREATION_TRACKER.reset(tracker_token) + + +@contextlib.contextmanager +def _no_state_recreation_tracking() -> Iterator[None]: + """Temporarily prevent child tasks from inheriting state tracking.""" + tracker_token = _STATE_RECREATION_TRACKER.set(None) + try: + yield + finally: + _STATE_RECREATION_TRACKER.reset(tracker_token) + + +def _report_state_recreation(state: "BaseState") -> None: + """Report a state recreated during a tracked modification context. + + Args: + state: The newly recreated state. + """ + if (recreated_states := _STATE_RECREATION_TRACKER.get()) is not None: + recreated_states.append(state) + + +def _report_state_tree_recreation(state: "BaseState") -> None: + """Report every state in a newly recreated tree. + + Args: + state: The root of the recreated state tree. + """ + if (recreated_states := _STATE_RECREATION_TRACKER.get()) is None: + return + pending = [state] + while pending: + current = pending.pop() + recreated_states.append(current) + pending.extend(current.substates.values()) + + def _release_state_tree(state: "BaseState"): """Clear the substate links of a purged state tree. diff --git a/reflex/istate/manager/disk.py b/reflex/istate/manager/disk.py index 920e19867fb..197ed64287a 100644 --- a/reflex/istate/manager/disk.py +++ b/reflex/istate/manager/disk.py @@ -19,6 +19,7 @@ StateModificationContext, _default_token_expiration, _release_state_tree, + _report_state_recreation, ) from reflex.istate.manager.token import TOKEN_TYPE, BaseStateToken, StateToken from reflex.state import BaseState @@ -186,6 +187,8 @@ async def populate_substates( instance.parent_state = state await self.populate_substates(token, instance, root_state) + if instance is fresh_instance: + _report_state_recreation(instance) @override async def get_state( @@ -216,6 +219,7 @@ async def get_state( fresh_root_state = root_state_cls(_reflex_internal_init=True) if root_state is None: root_state = fresh_root_state + _report_state_recreation(root_state) elif not isinstance(root_state, BaseState): msg = "Deserialized state is not an instance of BaseState, cannot populate substates." raise TypeError(msg) diff --git a/reflex/istate/manager/memory.py b/reflex/istate/manager/memory.py index bbbb1608026..56abad9493b 100644 --- a/reflex/istate/manager/memory.py +++ b/reflex/istate/manager/memory.py @@ -12,7 +12,9 @@ StateManager, StateModificationContext, _default_token_expiration, + _no_state_recreation_tracking, _release_state_tree, + _report_state_tree_recreation, ) from reflex.istate.manager.token import TOKEN_TYPE, BaseStateToken, StateToken @@ -56,9 +58,9 @@ def _get_or_create_state(self, token: StateToken[TOKEN_TYPE]) -> TOKEN_TYPE: key = token.cache_key if key not in self.states: if isinstance(token, BaseStateToken): - self.states[key] = token.cls.get_root_state()( - _reflex_internal_init=True - ) + state = token.cls.get_root_state()(_reflex_internal_init=True) + _report_state_tree_recreation(state) + self.states[key] = state else: self.states[key] = token.cls() return cast(TOKEN_TYPE, self.states[key]) @@ -139,10 +141,11 @@ def _ensure_expiration_task(self): """Ensure the expiration background task is running.""" if self._expiration_task is None or self._expiration_task.done(): asyncio.get_running_loop() # Ensure we're in an event loop. - self._expiration_task = asyncio.create_task( - self._expire_states(), - name="StateManagerMemory|Expiration", - ) + with _no_state_recreation_tracking(): + self._expiration_task = asyncio.create_task( + self._expire_states(), + name="StateManagerMemory|Expiration", + ) @override async def get_state(self, token: StateToken[TOKEN_TYPE]) -> TOKEN_TYPE: diff --git a/reflex/istate/manager/redis.py b/reflex/istate/manager/redis.py index 2d7d25ad244..8adad04e671 100644 --- a/reflex/istate/manager/redis.py +++ b/reflex/istate/manager/redis.py @@ -28,6 +28,8 @@ StateManager, StateModificationContext, _default_token_expiration, + _no_state_recreation_tracking, + _report_state_recreation, ) from reflex.istate.manager.token import TOKEN_TYPE, BaseStateToken, StateToken from reflex.state import BaseState @@ -127,16 +129,19 @@ class OplockFound(Exception): # noqa: N818 """Indicates that an opportunistic lock was found.""" -# KEYS: the lock key, then the state keys. ARGV: the lock id, the state -# expiration in seconds, then the serialized states in KEYS order. Writes only -# while the lock is held by ARGV[1], and returns the lock's remaining PTTL, or -# nil when the lock changed hands or expired and nothing was written. +# KEYS: the lock key, the state manifest key, then the state keys. +# ARGV: the lock id, the state expiration in seconds, then serialized states +# in KEYS order. The manifest tracks which per-state keys have been persisted. _FENCED_SAVE_SCRIPT = """ if redis.call('GET', KEYS[1]) ~= ARGV[1] then return nil end -for i = 2, #KEYS do - redis.call('SET', KEYS[i], ARGV[i + 1], 'EX', ARGV[2]) +for i = 3, #KEYS do + redis.call('SET', KEYS[i], ARGV[i], 'EX', ARGV[2]) + redis.call('SADD', KEYS[2], KEYS[i]) +end +if #KEYS > 2 then + redis.call('PEXPIRE', KEYS[2], tonumber(ARGV[2]) * 1000) end return redis.call('PTTL', KEYS[1]) """ @@ -357,18 +362,39 @@ async def get_state( key=lambda x: x.get_full_name(), ) - # Read the tree atomically with one command instead of a transaction - # containing a GET for each state. An already populated tree needs no IO. - redis_states = ( - await self.redis.mget([ - token._state_key(state_cls) for state_cls in required_state_classes - ]) - if required_state_classes - else [] - ) + # Read state values and their persisted-key manifest in one round trip. + state_keys = [ + token._state_key(state_cls) for state_cls in required_state_classes + ] + if required_state_classes: + pipeline = self.redis.pipeline(transaction=False) + pipeline.mget(state_keys) + pipeline.smembers(token._state_manifest_key()) + redis_states, manifest_keys = await pipeline.execute() + persisted_state_keys = { + key.decode() if isinstance(key, bytes) else key for key in manifest_keys + } + untracked_keys = [ + key + for key, redis_state in zip(state_keys, redis_states, strict=True) + if redis_state is not None and key not in persisted_state_keys + ] + if untracked_keys: + # Learn existing Redis keys on first access after introducing the manifest. + manifest_pipeline = self.redis.pipeline(transaction=False) + for key in untracked_keys: + manifest_pipeline.sadd(token._state_manifest_key(), key) + manifest_pipeline.pexpire( + token._state_manifest_key(), self.token_expiration * 1000 + ) + await manifest_pipeline.execute() + else: + redis_states = [] + persisted_state_keys = set() - for state_cls, redis_state in zip( + for state_cls, state_key, redis_state in zip( required_state_classes, + state_keys, redis_states, strict=True, ): @@ -379,11 +405,13 @@ async def get_state( with contextlib.suppress(StateSchemaMismatchError): state = BaseState._deserialize(data=redis_state) if state is None: - # Key didn't exist or schema mismatch so create a new instance for this token. + # A missing never-persisted state has no stale client value to reset. state = state_cls( init_substates=False, _reflex_internal_init=True, ) + if redis_state is not None or state_key in persisted_state_keys: + _report_state_recreation(state) state_full_name = state.get_full_name() flat_state_tree[state_full_name] = state if state.get_parent_state() is not None: @@ -438,11 +466,18 @@ async def set_state( pickle_state = token.serialize(state) writes = [(str(token), pickle_state)] if pickle_state else [] + manifest_key = ( + token._state_manifest_key() + if isinstance(token, BaseStateToken) + else f"{token.lock_key}_reflex_state_manifest" + ) if lock_id is None: if writes: pipeline = self.redis.pipeline(transaction=False) for key, pickle_state in writes: pipeline.set(key, pickle_state, ex=self.token_expiration) + pipeline.sadd(manifest_key, key) + pipeline.pexpire(manifest_key, self.token_expiration * 1000) await pipeline.execute() return @@ -460,8 +495,9 @@ async def set_state( fenced_save = cast("Callable[..., Awaitable[int | None]]", self.redis.eval) pttl = await fenced_save( _FENCED_SAVE_SCRIPT, - 1 + len(writes), + 2 + len(writes), lock_key, + manifest_key, *(key for key, _ in writes), lock_id, self.token_expiration, @@ -798,10 +834,11 @@ async def lease_breaker(): lock_key not in self._local_leases and await self._n_lock_contenders(self._lock_key(token)) == 0 ): - self._local_leases[lock_key] = task = asyncio.create_task( - lease_breaker(), - name=f"reflex_lease_breaker|{lock_key}|{lock_id.decode()}", - ) + with _no_state_recreation_tracking(): + self._local_leases[lock_key] = task = asyncio.create_task( + lease_breaker(), + name=f"reflex_lease_breaker|{lock_key}|{lock_id.decode()}", + ) # Fetch the requested state into the cache. self._cached_states[lock_key] = await self.get_state(token) return task @@ -908,12 +945,13 @@ async def _subscribe_lock_updates(self): def _ensure_lock_task(self) -> None: """Ensure the lock updates subscriber task is running.""" - ensure_task( - owner=self, - task_attribute="_lock_task", - coro_function=self._subscribe_lock_updates, - suppress_exceptions=[Exception], - ) + with _no_state_recreation_tracking(): + ensure_task( + owner=self, + task_attribute="_lock_task", + coro_function=self._subscribe_lock_updates, + suppress_exceptions=[Exception], + ) async def _ensure_lock_task_subscribed(self, timeout: float | None = None) -> None: """Ensure the lock updates subscriber task is running and subscribed to avoid missing notifications. diff --git a/reflex/istate/manager/token.py b/reflex/istate/manager/token.py index f6838b7abc6..c380bd59cd5 100644 --- a/reflex/istate/manager/token.py +++ b/reflex/istate/manager/token.py @@ -162,6 +162,14 @@ def _state_key(self, cls: type[BaseState]) -> str: """ return f"{self.ident}_{cls.get_full_name()}" + def _state_manifest_key(self) -> str: + """Get the Redis key recording state keys previously persisted for this token. + + Returns: + The Redis key for this token's persisted-state manifest. + """ + return f"{self.ident}_reflex_state_manifest" + def __str__(self) -> str: """The key used in the underlying StateManager store. diff --git a/reflex/istate/proxy.py b/reflex/istate/proxy.py index f15d4d1d3cc..79ae76753f8 100644 --- a/reflex/istate/proxy.py +++ b/reflex/istate/proxy.py @@ -211,6 +211,7 @@ def __init__( self._self_actx_lock = asyncio.Lock() self._self_actx_lock_holder = None self._self_parent_state_proxy = parent_state_proxy + self._self_recreated_states: list[BaseState] = [] # Whether `async with self` was ever entered; a background handler that # never did emitted no delta, so the processor flushes once for it. self._self_entered_context = False @@ -266,12 +267,29 @@ async def __aenter__(self) -> Self: self._self_actx = ctx.state_manager.modify_state_with_links( token=self._self_substate_token, event=self._self_event ) - mutable_state = await self._self_actx.__aenter__() + from reflex.istate.manager import _track_state_recreation + + with _track_state_recreation() as recreated_states: + mutable_state = await self._self_actx.__aenter__() + self._self_recreated_states = recreated_states + from reflex.state import State + + for state in recreated_states: + state_cls = type(state) + if state_cls is not State: + state.dirty_vars.update( + name + for name, field in state_cls.get_fields().items() + if field._owner is state_cls + ) + state.dirty_vars.update(state.computed_vars) + for computed_var in state.computed_vars.values(): + computed_var.mark_dirty(instance=state) + state._mark_dirty() + refreshed_state = mutable_state.get_substate(self._self_substate_path) self._self_mutable = True self._self_entered_context = True - super().__setattr__( - "__wrapped__", mutable_state.get_substate(self._self_substate_path) - ) + super().__setattr__("__wrapped__", refreshed_state) except (Exception, asyncio.CancelledError): # Restore the proxy to a consistent state since __aexit__ will not be called when __aenter__ raises. await self.__aexit__(*sys.exc_info()) @@ -292,7 +310,15 @@ async def __aexit__(self, *exc_info: Any) -> None: try: if self._self_mutable and self._self_actx is not None: root_state = self.__wrapped__._get_root_state() - delta = await root_state._get_resolved_delta() + touched_states = [ + (state, state._was_touched) for state in self._self_recreated_states + ] + try: + delta = await root_state._get_resolved_delta() + finally: + for state, was_touched in touched_states: + state._was_touched = was_touched + self._self_recreated_states.clear() root_state._clean() # When the frontend vars are modified emit the delta to the frontend. if delta: diff --git a/tests/integration/test_memory_state_manager_expiration.py b/tests/integration/test_memory_state_manager_expiration.py index f4d3e88d7dc..a1fb2a8b048 100644 --- a/tests/integration/test_memory_state_manager_expiration.py +++ b/tests/integration/test_memory_state_manager_expiration.py @@ -13,6 +13,8 @@ def MemoryExpirationApp(): """Reflex app that exposes state expiration through a simple counter UI.""" + import asyncio + import reflex as rx class State(rx.State): @@ -22,9 +24,36 @@ class State(rx.State): def increment(self): self.counter += 1 + @rx.event(background=True) + async def delayed_update(self): + await asyncio.sleep(2) + async with self: + self.counter = 100 + + class ChildState(State): + child_counter: int = 0 + + @rx.event + def set_values(self): + self.counter = 11 + self.child_counter = 13 + + class CounterComponent(rx.ComponentState): + value: int = 0 + + @rx.event + def set_value(self): + self.value = 17 + + @classmethod + def get_component(cls): + return rx.vstack( + rx.text(cls.value, id="component-counter"), + rx.button("Set component", id="set-component", on_click=cls.set_value), + ) + app = rx.App() - @app.add_page def index(): return rx.vstack( rx.input( @@ -33,9 +62,28 @@ def index(): is_read_only=True, ), rx.text(State.counter, id="counter"), + rx.text(ChildState.child_counter, id="child-counter"), rx.button("Increment", id="increment", on_click=State.increment), + rx.button("Set values", id="set-values", on_click=ChildState.set_values), + rx.button( + "Delayed update", + id="delayed-update", + on_click=State.delayed_update, + ), ) + app.add_page(index) + + def background_page(): + """Render the background expiry case with a component substate. + + Returns: + The page with state and component counters. + """ + return rx.vstack(index(), CounterComponent.create()) + + app.add_page(background_page, route="/background") + @pytest.fixture def memory_expiration_app( @@ -157,3 +205,38 @@ def test_memory_state_manager_delays_expiration_after_use_end_to_end( increment.click() AppHarness.expect(lambda: counter.text == "1") assert token_input.get_attribute("value") == token + + +def test_background_task_refreshes_recreated_state_tree( + memory_expiration_app: AppHarness, + driver: WebDriver, +): + """A delayed background update must reset expired child and component state.""" + frontend_url = memory_expiration_app.frontend_url + assert frontend_url is not None + driver.get(frontend_url.removesuffix("/") + "/background") + token_input = AppHarness.poll_for_or_raise_timeout( + lambda: driver.find_element(By.ID, "token") + ) + token = memory_expiration_app.poll_for_value(token_input) + assert token is not None + counter = AppHarness.poll_for_or_raise_timeout( + lambda: driver.find_element(By.ID, "counter") + ) + child_counter = AppHarness.poll_for_or_raise_timeout( + lambda: driver.find_element(By.ID, "child-counter") + ) + component_counter = AppHarness.poll_for_or_raise_timeout( + lambda: driver.find_element(By.ID, "component-counter") + ) + driver.find_element(By.ID, "set-values").click() + AppHarness.expect(lambda: counter.text == "11") + AppHarness.expect(lambda: child_counter.text == "13") + driver.find_element(By.ID, "set-component").click() + AppHarness.expect(lambda: component_counter.text == "17") + + driver.find_element(By.ID, "delayed-update").click() + AppHarness.expect(lambda: counter.text == "100", timeout=10) + AppHarness.expect(lambda: child_counter.text == "0", timeout=10) + AppHarness.expect(lambda: component_counter.text == "0", timeout=10) + AppHarness.expect(lambda: token_input.get_attribute("value") == token) diff --git a/tests/units/istate/manager/test_expiration.py b/tests/units/istate/manager/test_expiration.py index 93907e03569..566f851a1bf 100644 --- a/tests/units/istate/manager/test_expiration.py +++ b/tests/units/istate/manager/test_expiration.py @@ -9,6 +9,7 @@ import pytest import pytest_asyncio +from reflex.istate.manager import _STATE_RECREATION_TRACKER, _track_state_recreation from reflex.istate.manager.disk import StateManagerDisk from reflex.istate.manager.memory import StateManagerMemory from reflex.istate.manager.token import BaseStateToken @@ -61,6 +62,30 @@ async def state_manager_memory() -> AsyncGenerator[StateManagerMemory]: await state_manager.close() +@pytest.mark.asyncio +async def test_expiration_task_does_not_inherit_recreation_tracker( + monkeypatch: pytest.MonkeyPatch, +): + """The long-lived expiry task must not retain tracked state instances.""" + state_manager = StateManagerMemory(token_expiration=60) + tracker_values: list[object] = [] + + async def capture_tracker() -> None: + await asyncio.sleep(0) + tracker_values.append(_STATE_RECREATION_TRACKER.get()) + + monkeypatch.setattr(state_manager, "_expire_states", capture_tracker) + with _track_state_recreation(): + state_manager._ensure_expiration_task() + + expiration_task = state_manager._expiration_task + assert expiration_task is not None + await expiration_task + await state_manager.close() + + assert tracker_values == [None] + + @pytest.mark.asyncio async def test_memory_state_manager_evicts_expired_state( state_manager_memory: StateManagerMemory, @@ -183,6 +208,34 @@ async def test_disk_state_manager_set_state_refreshes_expiration( assert state_manager_disk._token_last_touched[token][0] > touched_at +@pytest.mark.asyncio +async def test_disk_state_manager_reports_recreated_states( + state_manager_disk: StateManagerDisk, + token: str, +): + """Missing persisted states should be reported for a background refresh.""" + state_token = BaseStateToken(ident=token, cls=ExpiringChildState) + + with _track_state_recreation() as recreated_states: + root = await state_manager_disk.get_state(state_token) + child = await root.get_state(ExpiringChildState) + + assert recreated_states == [root, child] + assert not root.dirty_vars + assert not child.dirty_vars + + root.value = 7 + await state_manager_disk.set_state(state_token, root) + state_manager_disk._purge_token(state_token) + + with _track_state_recreation() as recreated_states: + reloaded = await state_manager_disk.get_state(state_token) + reloaded_child = await reloaded.get_state(ExpiringChildState) + + assert getattr(reloaded, "value", None) == 7 + assert recreated_states == [reloaded_child] + + @pytest.mark.asyncio async def test_disk_state_manager_writes_expired_pending_state_tree( state_manager_disk: StateManagerDisk, diff --git a/tests/units/istate/manager/test_redis.py b/tests/units/istate/manager/test_redis.py index 5ed218e1ef4..0b4ce7d3b58 100644 --- a/tests/units/istate/manager/test_redis.py +++ b/tests/units/istate/manager/test_redis.py @@ -7,15 +7,16 @@ import time import uuid from collections import Counter -from collections.abc import AsyncGenerator +from collections.abc import AsyncGenerator, Awaitable from types import ModuleType -from typing import Any +from typing import Any, cast from unittest.mock import AsyncMock, Mock import pytest import pytest_asyncio from reflex_base.utils.exceptions import EnvironmentVarValueError +from reflex.istate.manager import _track_state_recreation from reflex.istate.manager.redis import ( StateManagerRedis, _default_lock_expiration, @@ -66,8 +67,8 @@ class RedisAppObjectState(BaseState): @pytest.mark.asyncio -async def test_get_state_reads_tree_in_one_command(): - """Read persisted and missing states together, preserving their tree positions.""" +async def test_get_state_reads_tree_in_one_round_trip(): + """Read persisted and missing states together in one round trip.""" redis = mock_redis() manager = StateManagerRedis(redis=redis) token = BaseStateToken(ident="batched-read", cls=RedisTestState) @@ -94,6 +95,65 @@ async def test_get_state_reads_tree_in_one_command(): await manager.close() +@pytest.mark.asyncio +async def test_get_state_does_not_dirty_persisted_states_on_reload(): + """Only recreated Redis states need defaults sent to the frontend.""" + redis = mock_redis() + manager = StateManagerRedis(redis=redis) + token = BaseStateToken(ident="persisted-read", cls=RedisTestState) + + state = await manager.get_state(token) + state.foo = "persisted" + await manager.set_state(token, state) + with _track_state_recreation() as recreated_states: + reloaded_state = await manager.get_state(token) + + assert isinstance(reloaded_state, RedisTestState) + assert reloaded_state.foo == "persisted" + assert not recreated_states + assert not reloaded_state.dirty_vars + assert all(not child.dirty_vars for child in reloaded_state.substates.values()) + await manager.close() + + +@pytest.mark.asyncio +async def test_get_state_reports_missing_redis_substates(): + """Missing Redis state keys are reported without reporting persisted states.""" + redis = mock_redis() + manager = StateManagerRedis(redis=redis) + token = BaseStateToken(ident="missing-substate", cls=RedisTestState) + state_classes = sorted( + manager._get_required_state_classes(RedisTestState, subclasses=True), + key=lambda cls: cls.get_full_name(), + ) + persisted_state = RedisTestState(foo="persisted") + for state_cls in state_classes[1:]: + await cast( + Awaitable[int], + redis.sadd( + token._state_manifest_key(), + token._state_key(state_cls), + ), + ) + redis.mget = AsyncMock( + return_value=[ + persisted_state._serialize() if cls is RedisTestState else None + for cls in state_classes + ] + ) + + with _track_state_recreation() as recreated_states: + state = await manager.get_state(token) + + assert isinstance(state, RedisTestState) + assert state.foo == "persisted" + assert {type(recreated).__name__ for recreated in recreated_states} == { + "SubState1", + "SubState2", + } + await manager.close() + + @pytest.mark.asyncio async def test_get_state_reuses_populated_tree_without_reading(): """Fetching an already attached state must not contact Redis or replace it.""" @@ -1179,7 +1239,9 @@ def recording_eval(script: str, numkeys: int, *keys_and_args: Any) -> Any: await state_manager_redis.set_state(token, state, lock_id=lock_id) assert not calls - assert saved_keys == [(lock_key, token._state_key(TreeFirst))] + assert saved_keys == [ + (lock_key, token._state_manifest_key(), token._state_key(TreeFirst)) + ] assert await _stored_tree_states(state_manager_redis, token) == {TreeFirst} await redis.delete(lock_key) diff --git a/tests/units/istate/test_proxy.py b/tests/units/istate/test_proxy.py index 2138b4d0e8a..f53e140753f 100644 --- a/tests/units/istate/test_proxy.py +++ b/tests/units/istate/test_proxy.py @@ -11,7 +11,7 @@ from collections.abc import Mapping from contextlib import asynccontextmanager from operator import attrgetter -from typing import Any, ClassVar +from typing import Any, ClassVar, cast import pytest from reflex_base.constants.state import FIELD_MARKER @@ -21,6 +21,7 @@ import reflex as rx from reflex.istate.data import HeaderData, PageData, RouterData from reflex.istate.manager import StateManager +from reflex.istate.manager.memory import StateManagerMemory from reflex.istate.manager.token import BaseStateToken from reflex.istate.proxy import ( ImmutableMutableProxy, @@ -132,8 +133,8 @@ class Item: id: int -class ProxyTestState(rx.State): - """Test state with a list field.""" +class ProxyTestState(BaseState): + """Isolated state tree with a list field for proxy tests.""" items: list[Item] = [] @@ -186,6 +187,65 @@ async def mock_modify_state_context(*args, **kwargs): pass +@pytest.mark.asyncio +async def test_recreated_state_invalidates_computed_vars_from_sibling_state( + attached_mock_event_context: EventContext, + emitted_deltas: list[tuple[str, Mapping[str, Mapping[str, Any]]]], +): + """StateProxy recreation refreshes a cross-state computed var cache.""" + + class RootState(BaseState): + pass + + class SourceState(RootState): + value: int = 7 + + class TargetState(RootState): + @rx.var + async def source_value(self) -> int: + source = await self.get_state(SourceState) + return source.value + + state_manager = StateManagerMemory(token_expiration=60) + original_state_manager = attached_mock_event_context.state_manager + object.__setattr__(attached_mock_event_context, "state_manager", state_manager) + state_token = BaseStateToken( + ident=attached_mock_event_context.token, + cls=TargetState, + ) + try: + async with state_manager.modify_state(state_token) as root: + source = cast( + SourceState, root.get_substate(SourceState.get_full_name().split(".")) + ) + target = cast( + TargetState, root.get_substate(TargetState.get_full_name().split(".")) + ) + source.value = 1 + computed_var = TargetState.computed_vars["source_value"] + cast(dict[str, Any], target.__dict__)[computed_var._cache_attr] = 1 + + state_proxy = StateProxy(target) + state_manager._purge_token(state_token) + async with state_proxy: + pass + + recreated_target = state_proxy.__wrapped__ + assert recreated_target.__dict__[computed_var._cache_attr] == 7 + assert ( + emitted_deltas[0][1][recreated_target.get_full_name()][ + "source_value" + FIELD_MARKER + ] + == 7 + ) + assert not recreated_target._was_touched + finally: + await state_manager.close() + object.__setattr__( + attached_mock_event_context, "state_manager", original_state_manager + ) + + @dataclasses.dataclass class TaggedModel: """A dataclass with a mutable attribute for attr-path refresh tests.""" @@ -609,6 +669,36 @@ async def test_immutable_mutable_proxy_async_context_manager( assert state.data["b"] == [2, 4] +@pytest.mark.asyncio +async def test_state_proxy_does_not_emit_delta_for_reloaded_unchanged_state( + token: str, + attached_mock_event_context: EventContext, + emitted_deltas: list[tuple[str, Mapping[str, Mapping[str, Any]]]], + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A freshly deserialized but unchanged state must not send a full delta.""" + state_manager = attached_mock_event_context.state_manager + state_token = BaseStateToken(ident=token, cls=ProxyTestState) + + async with state_manager.modify_state(state_token) as state: + state._clean() + state_proxy = StateProxy(state) + + original_modify_state = state_manager.modify_state + + @asynccontextmanager + async def reload_state(token: BaseStateToken, **context: Any): + async with original_modify_state(token, **context) as state: + yield token.deserialize(data=token.serialize(state)) + + monkeypatch.setattr(state_manager, "modify_state", reload_state) + + async with state_proxy: + pass + + assert not emitted_deltas + + @pytest.mark.asyncio async def test_immutable_mutable_proxy_async_context_dict_method_paths( token: str, attached_mock_event_context: EventContext diff --git a/tests/units/mock_redis.py b/tests/units/mock_redis.py index 1cbbd65ed95..66584f19889 100644 --- a/tests/units/mock_redis.py +++ b/tests/units/mock_redis.py @@ -123,6 +123,15 @@ async def mock_srem(key: KeyT, value: EncodableT) -> int: return 1 return 0 + async def mock_smembers(key: KeyT) -> set[EncodableT]: # noqa: RUF029 + _expire_keys() + keyset = keys.get(_key_bytes(key)) + if keyset is None: + return set() + if not isinstance(keyset, set): + raise TypeError(WRONGTYPE_MESSAGE) + return set(keyset) + async def mock_scard(key: KeyT) -> int: # noqa: RUF029 _expire_keys() keyset = keys.get(_key_bytes(key)) @@ -176,6 +185,12 @@ def set( def get(self, key: KeyT): self.results.append(redis_mock.get(key=key)) + def mget(self, keys: list[KeyT]): + self.results.append(redis_mock.mget(keys)) + + def smembers(self, key: KeyT): + self.results.append(redis_mock.smembers(key=key)) + def sadd(self, key: KeyT, value: EncodableT): self.results.append(redis_mock.sadd(key=key, value=value)) @@ -183,7 +198,7 @@ def pexpire(self, key: KeyT, px: int, xx: bool = False): self.results.append(redis_mock.pexpire(key=key, px=px, xx=xx)) async def execute(self): - results = await asyncio.gather(*self.results) + results = [await result for result in self.results] self.results = [] return results @@ -211,12 +226,15 @@ async def mock_eval(script: str, numkeys: int, *keys_and_args: Any) -> Any: if script != _FENCED_SAVE_SCRIPT: msg = "mock_redis only emulates the state manager's fenced save script." raise NotImplementedError(msg) - lock_key, *state_keys = keys_and_args[:numkeys] + lock_key, manifest_key, *state_keys = keys_and_args[:numkeys] lock_id, expiration, *payloads = keys_and_args[numkeys:] if await redis_mock.get(lock_key) != lock_id: return None for key, payload in zip(state_keys, payloads, strict=True): await redis_mock.set(key, payload, ex=int(expiration)) + await redis_mock.sadd(manifest_key, key) + if state_keys: + await redis_mock.pexpire(manifest_key, int(expiration) * 1000) return await redis_mock.pttl(lock_key) async def pttl(key: KeyT) -> int: # noqa: RUF029 @@ -297,6 +315,7 @@ async def listen() -> AsyncGenerator[dict[str, Any], None]: redis_mock.sadd = mock_sadd redis_mock.srem = mock_srem redis_mock.scard = mock_scard + redis_mock.smembers = mock_smembers redis_mock.pexpire = mock_pexpire redis_mock.pipeline = pipeline redis_mock.eval = mock_eval