From e0d0b19a9040eb69352e5e4233570d0725c18682 Mon Sep 17 00:00:00 2001 From: Harsh Thakare Date: Tue, 30 Jun 2026 13:33:20 +0530 Subject: [PATCH 1/7] support async context for mutable proxies --- news/6689.feature.md | 1 + reflex/istate/proxy.py | 26 ++++++++++++++++++++++++++ tests/units/test_state.py | 37 +++++++++++++++++++++++++++++++++++++ 3 files changed, 64 insertions(+) create mode 100644 news/6689.feature.md diff --git a/news/6689.feature.md b/news/6689.feature.md new file mode 100644 index 00000000000..b549066792a --- /dev/null +++ b/news/6689.feature.md @@ -0,0 +1 @@ +Support using mutable state proxies as async context managers. diff --git a/reflex/istate/proxy.py b/reflex/istate/proxy.py index 776a519fe2c..6f71a945dc0 100644 --- a/reflex/istate/proxy.py +++ b/reflex/istate/proxy.py @@ -467,6 +467,32 @@ def __repr__(self) -> str: """ return f"{type(self).__name__}({self.__wrapped__})" + async def __aenter__(self) -> Self: + """Enter the async context manager protocol through the bound state. + + Returns: + This proxy refreshed from the current state field. + """ + state = await self._self_state.__aenter__() + try: + refreshed_value = getattr(state, self._self_field_name) + if isinstance(refreshed_value, MutableProxy): + super().__setattr__("__wrapped__", refreshed_value.__wrapped__) + self._self_state = refreshed_value._self_state + self._self_field_name = refreshed_value._self_field_name + except (Exception, asyncio.CancelledError): + await self._self_state.__aexit__(*sys.exc_info()) + raise + return self + + async def __aexit__(self, *exc_info: Any) -> None: + """Exit the async context manager protocol through the bound state. + + Args: + exc_info: The exception info tuple. + """ + await self._self_state.__aexit__(*exc_info) + def _mark_dirty( self, wrapped: Callable | None = None, diff --git a/tests/units/test_state.py b/tests/units/test_state.py index b0aaf6e022f..a8e6d989450 100644 --- a/tests/units/test_state.py +++ b/tests/units/test_state.py @@ -5006,6 +5006,43 @@ async def test_rebind_mutable_proxy( assert state.data["b"] == [2, 3] +@pytest.mark.asyncio +async def test_immutable_mutable_proxy_async_context_manager( + token: str, attached_mock_event_context: EventContext +) -> None: + """Mutable state proxies can enter the owning StateProxy context.""" + state_manager = attached_mock_event_context.state_manager + + async with state_manager.modify_state( + BaseStateToken(ident=token, cls=MutableProxyState) + ) as state: + state.router = RouterData.from_router_data({ + "query": {}, + "token": token, + "sid": "test_sid", + }) + state_proxy = StateProxy(state) + data_proxy = state_proxy.data + + assert isinstance(data_proxy, ImmutableMutableProxy) + with pytest.raises(ImmutableStateError): + data_proxy["a"].append(3) + + async with data_proxy as mutable_data: + assert mutable_data is data_proxy + data_proxy["a"].append(3) + mutable_data["b"].append(4) + + with pytest.raises(ImmutableStateError): + data_proxy["a"].append(5) + + async with state_manager.modify_state( + BaseStateToken(ident=token, cls=MutableProxyState) + ) as state: + assert state.data["a"] == [1, 3] + assert state.data["b"] == [2, 4] + + def test_override_base_method_skips_event_handler_wrapping(): """A method marked with __override_base_method__ should not be wrapped as an EventHandler.""" from reflex.state import _override_base_method From 7063e53636b4aa0fb0b231fc82431d4265b4f6ce Mon Sep 17 00:00:00 2001 From: Harsh Thakare Date: Tue, 30 Jun 2026 13:50:49 +0530 Subject: [PATCH 2/7] preserve nested mutable proxy context targets --- reflex/istate/proxy.py | 54 +++++++++++++++++++++++++++++++++------ tests/units/test_state.py | 20 +++++++++++++-- 2 files changed, 64 insertions(+), 10 deletions(-) diff --git a/reflex/istate/proxy.py b/reflex/istate/proxy.py index 6f71a945dc0..4f293cf118b 100644 --- a/reflex/istate/proxy.py +++ b/reflex/istate/proxy.py @@ -12,7 +12,7 @@ from collections.abc import Callable, Sequence from importlib.util import find_spec from types import MethodType -from typing import TYPE_CHECKING, Any, SupportsIndex, TypeVar +from typing import TYPE_CHECKING, Any, NoReturn, SupportsIndex, TypeVar import wrapt from reflex_base.event import Event @@ -446,7 +446,13 @@ def __new__(cls, wrapped: Any, *args, **kwargs) -> MutableProxy: cls = cls.__dataclass_proxies__[wrapper_cls_name] return super().__new__(cls) # pyright: ignore[reportArgumentType] - def __init__(self, wrapped: Any, state: BaseState, field_name: str): + def __init__( + self, + wrapped: Any, + state: BaseState, + field_name: str, + path: tuple[tuple[str, Any], ...] | None = None, + ): """Create a proxy for a mutable object that tracks changes. Args: @@ -454,10 +460,13 @@ def __init__(self, wrapped: Any, state: BaseState, field_name: str): state: The state to mark dirty when the object is changed. field_name: The name of the field on the state associated with the wrapped object. + path: Access path from the state field to this wrapped object. """ super().__init__(wrapped) self._self_state = state self._self_field_name = field_name + self._self_path = path or () + self._self_actx_state = None def __repr__(self) -> str: """Get the representation of the wrapped object. @@ -473,25 +482,49 @@ async def __aenter__(self) -> Self: Returns: This proxy refreshed from the current state field. """ - state = await self._self_state.__aenter__() + context_state = self._self_state + self._self_actx_state = context_state + state = await context_state.__aenter__() try: refreshed_value = getattr(state, self._self_field_name) + for access_kind, access_key in self._self_path: + refreshed_value = ( + getattr(refreshed_value, access_key) + if access_kind == "attr" + else refreshed_value[access_key] + ) if isinstance(refreshed_value, MutableProxy): super().__setattr__("__wrapped__", refreshed_value.__wrapped__) self._self_state = refreshed_value._self_state self._self_field_name = refreshed_value._self_field_name + self._self_path = refreshed_value._self_path + else: + self._raise_refresh_error() except (Exception, asyncio.CancelledError): - await self._self_state.__aexit__(*sys.exc_info()) + await context_state.__aexit__(*sys.exc_info()) + self._self_actx_state = None raise return self + def _raise_refresh_error(self) -> NoReturn: + """Raise when this proxy cannot be refreshed from its state field.""" + msg = ( + "Unable to refresh mutable proxy from state field " + f"`{self._self_field_name}`." + ) + raise RuntimeError(msg) + async def __aexit__(self, *exc_info: Any) -> None: """Exit the async context manager protocol through the bound state. Args: exc_info: The exception info tuple. """ - await self._self_state.__aexit__(*exc_info) + context_state = self._self_actx_state or self._self_state + try: + await context_state.__aexit__(*exc_info) + finally: + self._self_actx_state = None def _mark_dirty( self, @@ -539,11 +572,14 @@ def _is_called_from_dataclasses_internal() -> bool: return True return False - def _wrap_recursive(self, value: Any) -> Any: + def _wrap_recursive( + self, value: Any, path: tuple[tuple[str, Any], ...] | None = None + ) -> Any: """Wrap a value recursively if it is mutable. Args: value: The value to wrap. + path: Access path from the state field to this wrapped object. Returns: The wrapped value. @@ -562,6 +598,7 @@ def _wrap_recursive(self, value: Any) -> Any: wrapped=value, state=self._self_state, field_name=self._self_field_name, + path=path or self._self_path, ) return value @@ -618,10 +655,11 @@ def __getattr__(self, __name: str) -> Any: if is_mutable_type(type(value)) and __name not in ( "__wrapped__", "_self_state", + "_self_path", "__dict__", ): # Recursively wrap mutable attribute values retrieved through this proxy. - return self._wrap_recursive(value) + return self._wrap_recursive(value, (*self._self_path, ("attr", __name))) return value @@ -638,7 +676,7 @@ def __getitem__(self, key: Any) -> Any: if isinstance(key, slice) and isinstance(value, list): return [self._wrap_recursive(item) for item in value] # Recursively wrap mutable items retrieved through this proxy. - return self._wrap_recursive(value) + return self._wrap_recursive(value, (*self._self_path, ("item", key))) def __iter__(self) -> Any: """Iterate over the proxied object and return a proxy if mutable. diff --git a/tests/units/test_state.py b/tests/units/test_state.py index a8e6d989450..4f17cdc7d1b 100644 --- a/tests/units/test_state.py +++ b/tests/units/test_state.py @@ -5023,23 +5023,39 @@ async def test_immutable_mutable_proxy_async_context_manager( }) state_proxy = StateProxy(state) data_proxy = state_proxy.data + items_proxy = data_proxy["a"] assert isinstance(data_proxy, ImmutableMutableProxy) + assert isinstance(items_proxy, ImmutableMutableProxy) with pytest.raises(ImmutableStateError): data_proxy["a"].append(3) + with pytest.raises(ImmutableStateError): + items_proxy.append(3) + + async with state_manager.modify_state( + BaseStateToken(ident=token, cls=MutableProxyState) + ) as state: + state.data["a"].append(2) async with data_proxy as mutable_data: assert mutable_data is data_proxy data_proxy["a"].append(3) mutable_data["b"].append(4) + async with items_proxy as mutable_items: + assert mutable_items is items_proxy + mutable_items.append(5) + assert items_proxy.__wrapped__ == [1, 2, 3, 5] + + with pytest.raises(ImmutableStateError): + data_proxy["a"].append(6) with pytest.raises(ImmutableStateError): - data_proxy["a"].append(5) + items_proxy.append(6) async with state_manager.modify_state( BaseStateToken(ident=token, cls=MutableProxyState) ) as state: - assert state.data["a"] == [1, 3] + assert state.data["a"] == [1, 2, 3, 5] assert state.data["b"] == [2, 4] From d7894a50f0e1b0286b1a28f5bd04ebaf0ce75ab3 Mon Sep 17 00:00:00 2001 From: Harsh Thakare Date: Wed, 1 Jul 2026 14:45:33 +0530 Subject: [PATCH 3/7] Fix mutable proxy async context test typing --- tests/units/test_state.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/units/test_state.py b/tests/units/test_state.py index 4f17cdc7d1b..5972cbbfa34 100644 --- a/tests/units/test_state.py +++ b/tests/units/test_state.py @@ -5035,6 +5035,7 @@ async def test_immutable_mutable_proxy_async_context_manager( async with state_manager.modify_state( BaseStateToken(ident=token, cls=MutableProxyState) ) as state: + assert isinstance(state, MutableProxyState) state.data["a"].append(2) async with data_proxy as mutable_data: @@ -5055,6 +5056,7 @@ async def test_immutable_mutable_proxy_async_context_manager( async with state_manager.modify_state( BaseStateToken(ident=token, cls=MutableProxyState) ) as state: + assert isinstance(state, MutableProxyState) assert state.data["a"] == [1, 2, 3, 5] assert state.data["b"] == [2, 4] From d8f12f3ab23492adde8966a402eebed382e1d008 Mon Sep 17 00:00:00 2001 From: Harsh Thakare Date: Thu, 2 Jul 2026 09:31:36 +0530 Subject: [PATCH 4/7] Address mutable proxy async context review --- reflex/istate/proxy.py | 62 +++++++++++++++++++++++++++------------ tests/units/test_state.py | 43 +++++++++++++++++++++++++++ 2 files changed, 87 insertions(+), 18 deletions(-) diff --git a/reflex/istate/proxy.py b/reflex/istate/proxy.py index 4f293cf118b..00fcf75a179 100644 --- a/reflex/istate/proxy.py +++ b/reflex/istate/proxy.py @@ -12,7 +12,7 @@ from collections.abc import Callable, Sequence from importlib.util import find_spec from types import MethodType -from typing import TYPE_CHECKING, Any, NoReturn, SupportsIndex, TypeVar +from typing import TYPE_CHECKING, Any, Literal, NoReturn, SupportsIndex, TypeVar import wrapt from reflex_base.event import Event @@ -29,6 +29,11 @@ T_STATE = TypeVar("T_STATE", bound="BaseState") T = TypeVar("T") +AccessSpec = ( + tuple[Literal["attr"], str] + | tuple[Literal["item"], Any] + | tuple[Literal["iter"], None] +) # Cached filename of the dataclasses module, used to detect reads originating # from `dataclasses.asdict`/`astuple` internals on the proxy read hot-path. @@ -451,7 +456,7 @@ def __init__( wrapped: Any, state: BaseState, field_name: str, - path: tuple[tuple[str, Any], ...] | None = None, + path: tuple[AccessSpec, ...] | None = None, ): """Create a proxy for a mutable object that tracks changes. @@ -482,21 +487,31 @@ async def __aenter__(self) -> Self: Returns: This proxy refreshed from the current state field. """ + if self._self_actx_state is not None: + msg = ( + "Mutable proxy is already mutable. Do not nest `async with proxy` " + "blocks." + ) + raise RuntimeError(msg) context_state = self._self_state self._self_actx_state = context_state state = await context_state.__aenter__() try: refreshed_value = getattr(state, self._self_field_name) - for access_kind, access_key in self._self_path: - refreshed_value = ( - getattr(refreshed_value, access_key) - if access_kind == "attr" - else refreshed_value[access_key] - ) - if isinstance(refreshed_value, MutableProxy): + for access_spec in self._self_path: + match access_spec: + case ("attr", access_key): + refreshed_value = getattr(refreshed_value, access_key) + case ("item", access_key): + refreshed_value = refreshed_value[access_key] + case _: + self._raise_refresh_error() + if ( + isinstance(refreshed_value, MutableProxy) + and self._self_field_name == refreshed_value._self_field_name + ): super().__setattr__("__wrapped__", refreshed_value.__wrapped__) self._self_state = refreshed_value._self_state - self._self_field_name = refreshed_value._self_field_name self._self_path = refreshed_value._self_path else: self._raise_refresh_error() @@ -520,7 +535,9 @@ async def __aexit__(self, *exc_info: Any) -> None: Args: exc_info: The exception info tuple. """ - context_state = self._self_actx_state or self._self_state + context_state = self._self_actx_state + if context_state is None: + return try: await context_state.__aexit__(*exc_info) finally: @@ -573,13 +590,13 @@ def _is_called_from_dataclasses_internal() -> bool: return False def _wrap_recursive( - self, value: Any, path: tuple[tuple[str, Any], ...] | None = None + self, value: Any, new_path_segment: AccessSpec | None = None ) -> Any: """Wrap a value recursively if it is mutable. Args: value: The value to wrap. - path: Access path from the state field to this wrapped object. + new_path_segment: Access path segment from this proxy to the value. Returns: The wrapped value. @@ -598,7 +615,11 @@ def _wrap_recursive( wrapped=value, state=self._self_state, field_name=self._self_field_name, - path=path or self._self_path, + path=( + self._self_path + if new_path_segment is None + else (*self._self_path, new_path_segment) + ), ) return value @@ -656,10 +677,11 @@ def __getattr__(self, __name: str) -> Any: "__wrapped__", "_self_state", "_self_path", + "_self_actx_state", "__dict__", ): # Recursively wrap mutable attribute values retrieved through this proxy. - return self._wrap_recursive(value, (*self._self_path, ("attr", __name))) + return self._wrap_recursive(value, ("attr", __name)) return value @@ -674,9 +696,13 @@ def __getitem__(self, key: Any) -> Any: """ value = super().__getitem__(key) # pyright: ignore[reportAttributeAccessIssue] if isinstance(key, slice) and isinstance(value, list): - return [self._wrap_recursive(item) for item in value] + indices = range(len(self.__wrapped__))[key] + return [ + self._wrap_recursive(item, ("item", index)) + for item, index in zip(value, indices, strict=False) + ] # Recursively wrap mutable items retrieved through this proxy. - return self._wrap_recursive(value, (*self._self_path, ("item", key))) + return self._wrap_recursive(value, ("item", key)) def __iter__(self) -> Any: """Iterate over the proxied object and return a proxy if mutable. @@ -686,7 +712,7 @@ def __iter__(self) -> Any: """ for value in super().__iter__(): # pyright: ignore[reportAttributeAccessIssue] # Recursively wrap mutable items retrieved through this proxy. - yield self._wrap_recursive(value) + yield self._wrap_recursive(value, ("iter", None)) def __delattr__(self, name: str): """Delete the attribute on the proxied object and mark state dirty. diff --git a/tests/units/test_state.py b/tests/units/test_state.py index 5972cbbfa34..efb60c3a59f 100644 --- a/tests/units/test_state.py +++ b/tests/units/test_state.py @@ -4954,6 +4954,12 @@ class MutableProxyState(BaseState): data: dict[str, list[int]] = {"a": [1], "b": [2]} +class IterableMutableProxyState(BaseState): + """A test state with mutable values that can be accessed by iteration.""" + + data: list[list[int]] = [[1]] + + @pytest.mark.asyncio async def test_rebind_mutable_proxy( token: str, attached_mock_event_context: EventContext @@ -5027,6 +5033,7 @@ async def test_immutable_mutable_proxy_async_context_manager( assert isinstance(data_proxy, ImmutableMutableProxy) assert isinstance(items_proxy, ImmutableMutableProxy) + await data_proxy.__aexit__(None, None, None) with pytest.raises(ImmutableStateError): data_proxy["a"].append(3) with pytest.raises(ImmutableStateError): @@ -5040,6 +5047,9 @@ async def test_immutable_mutable_proxy_async_context_manager( async with data_proxy as mutable_data: assert mutable_data is data_proxy + with pytest.raises(RuntimeError, match="already mutable"): + async with data_proxy: + pass data_proxy["a"].append(3) mutable_data["b"].append(4) @@ -5061,6 +5071,39 @@ async def test_immutable_mutable_proxy_async_context_manager( assert state.data["b"] == [2, 4] +@pytest.mark.asyncio +async def test_immutable_mutable_proxy_async_context_rejects_iter_proxy( + token: str, attached_mock_event_context: EventContext +) -> None: + """Iteration-sourced mutable proxies fail clearly as async context managers.""" + state_manager = attached_mock_event_context.state_manager + + async with state_manager.modify_state( + BaseStateToken(ident=token, cls=IterableMutableProxyState) + ) as state: + state.router = RouterData.from_router_data({ + "query": {}, + "token": token, + "sid": "test_sid", + }) + state_proxy = StateProxy(state) + [items_proxy] = state_proxy.data + + assert isinstance(items_proxy, ImmutableMutableProxy) + with pytest.raises(RuntimeError, match="Unable to refresh mutable proxy"): + async with items_proxy: + pass + + async with state_proxy: + state_proxy.data[0].append(2) + + async with state_manager.modify_state( + BaseStateToken(ident=token, cls=IterableMutableProxyState) + ) as state: + assert isinstance(state, IterableMutableProxyState) + assert state.data == [[1, 2]] + + def test_override_base_method_skips_event_handler_wrapping(): """A method marked with __override_base_method__ should not be wrapped as an EventHandler.""" from reflex.state import _override_base_method From 78dec3632dd77e776680da1985ab28a408061d13 Mon Sep 17 00:00:00 2001 From: Harsh Thakare Date: Thu, 2 Jul 2026 10:25:30 +0530 Subject: [PATCH 5/7] Handle mutable proxy enter failures --- reflex/istate/proxy.py | 7 ++++-- tests/units/test_state.py | 47 +++++++++++++++++++++++++++++++++++++++ 2 files changed, 52 insertions(+), 2 deletions(-) diff --git a/reflex/istate/proxy.py b/reflex/istate/proxy.py index 00fcf75a179..368b49e92e4 100644 --- a/reflex/istate/proxy.py +++ b/reflex/istate/proxy.py @@ -495,8 +495,10 @@ async def __aenter__(self) -> Self: raise RuntimeError(msg) context_state = self._self_state self._self_actx_state = context_state - state = await context_state.__aenter__() + aenter_ok = False try: + state = await context_state.__aenter__() + aenter_ok = True refreshed_value = getattr(state, self._self_field_name) for access_spec in self._self_path: match access_spec: @@ -516,7 +518,8 @@ async def __aenter__(self) -> Self: else: self._raise_refresh_error() except (Exception, asyncio.CancelledError): - await context_state.__aexit__(*sys.exc_info()) + if aenter_ok: + await context_state.__aexit__(*sys.exc_info()) self._self_actx_state = None raise return self diff --git a/tests/units/test_state.py b/tests/units/test_state.py index efb60c3a59f..be3a6249cd3 100644 --- a/tests/units/test_state.py +++ b/tests/units/test_state.py @@ -5104,6 +5104,53 @@ async def test_immutable_mutable_proxy_async_context_rejects_iter_proxy( assert state.data == [[1, 2]] +@pytest.mark.asyncio +async def test_immutable_mutable_proxy_async_context_recovers_from_enter_failure( + token: str, + attached_mock_event_context: EventContext, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A failed owning StateProxy enter does not permanently block the proxy.""" + state_manager = attached_mock_event_context.state_manager + + async with state_manager.modify_state( + BaseStateToken(ident=token, cls=MutableProxyState) + ) as state: + state.router = RouterData.from_router_data({ + "query": {}, + "token": token, + "sid": "test_sid", + }) + state_proxy = StateProxy(state) + data_proxy = state_proxy.data + + original_aenter = StateProxy.__aenter__ + fail_once = True + + async def fail_first_enter(self: StateProxy) -> StateProxy: + nonlocal fail_once + if fail_once: + fail_once = False + raise asyncio.CancelledError + return await original_aenter(self) + + monkeypatch.setattr(StateProxy, "__aenter__", fail_first_enter) + + with pytest.raises(asyncio.CancelledError): + async with data_proxy: + pass + + assert data_proxy._self_actx_state is None + async with data_proxy as mutable_data: + mutable_data["a"].append(2) + + async with state_manager.modify_state( + BaseStateToken(ident=token, cls=MutableProxyState) + ) as state: + assert isinstance(state, MutableProxyState) + assert state.data["a"] == [1, 2] + + def test_override_base_method_skips_event_handler_wrapping(): """A method marked with __override_base_method__ should not be wrapped as an EventHandler.""" from reflex.state import _override_base_method From 2aa4c8c975eb581de37c84e32b11104d9e0f4a62 Mon Sep 17 00:00:00 2001 From: Harsh Thakare Date: Sat, 4 Jul 2026 10:27:13 +0530 Subject: [PATCH 6/7] Optimize mutable proxy iteration path tracking --- reflex/istate/proxy.py | 17 +++++++++++------ 1 file changed, 11 insertions(+), 6 deletions(-) diff --git a/reflex/istate/proxy.py b/reflex/istate/proxy.py index 368b49e92e4..97cdbe3936f 100644 --- a/reflex/istate/proxy.py +++ b/reflex/istate/proxy.py @@ -34,6 +34,8 @@ | tuple[Literal["item"], Any] | tuple[Literal["iter"], None] ) +_ITER_ACCESS_SPEC: AccessSpec = ("iter", None) +_ROOT_ITER_PATH: tuple[AccessSpec, ...] = (_ITER_ACCESS_SPEC,) # Cached filename of the dataclasses module, used to detect reads originating # from `dataclasses.asdict`/`astuple` internals on the proxy read hot-path. @@ -614,15 +616,18 @@ def _wrap_recursive( # Recursively wrap mutable types. if is_mutable_type(type(value)): base_cls = globals()[self.__base_proxy__] + path = self._self_path + if new_path_segment is not None: + path = ( + _ROOT_ITER_PATH + if new_path_segment is _ITER_ACCESS_SPEC and not path + else (*path, new_path_segment) + ) return base_cls( wrapped=value, state=self._self_state, field_name=self._self_field_name, - path=( - self._self_path - if new_path_segment is None - else (*self._self_path, new_path_segment) - ), + path=path, ) return value @@ -715,7 +720,7 @@ def __iter__(self) -> Any: """ for value in super().__iter__(): # pyright: ignore[reportAttributeAccessIssue] # Recursively wrap mutable items retrieved through this proxy. - yield self._wrap_recursive(value, ("iter", None)) + yield self._wrap_recursive(value, _ITER_ACCESS_SPEC) def __delattr__(self, name: str): """Delete the attribute on the proxied object and mark state dirty. From 0f899c6723fa0bc60b45503a846c4b9b53be9eaf Mon Sep 17 00:00:00 2001 From: Harsh Thakare Date: Sat, 4 Jul 2026 10:37:41 +0530 Subject: [PATCH 7/7] Clear mutable proxy context after cleanup failure --- reflex/istate/proxy.py | 8 +++++--- tests/units/test_state.py | 34 ++++++++++++++++++++++++++++++++++ 2 files changed, 39 insertions(+), 3 deletions(-) diff --git a/reflex/istate/proxy.py b/reflex/istate/proxy.py index 97cdbe3936f..67cbd7a94fb 100644 --- a/reflex/istate/proxy.py +++ b/reflex/istate/proxy.py @@ -520,9 +520,11 @@ async def __aenter__(self) -> Self: else: self._raise_refresh_error() except (Exception, asyncio.CancelledError): - if aenter_ok: - await context_state.__aexit__(*sys.exc_info()) - self._self_actx_state = None + try: + if aenter_ok: + await context_state.__aexit__(*sys.exc_info()) + finally: + self._self_actx_state = None raise return self diff --git a/tests/units/test_state.py b/tests/units/test_state.py index be3a6249cd3..181d4e95c7b 100644 --- a/tests/units/test_state.py +++ b/tests/units/test_state.py @@ -5151,6 +5151,40 @@ async def fail_first_enter(self: StateProxy) -> StateProxy: assert state.data["a"] == [1, 2] +@pytest.mark.asyncio +async def test_immutable_mutable_proxy_async_context_clears_state_when_cleanup_fails( + token: str, + attached_mock_event_context: EventContext, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A failed cleanup while entering does not leave the proxy marked mutable.""" + state_manager = attached_mock_event_context.state_manager + + async with state_manager.modify_state( + BaseStateToken(ident=token, cls=IterableMutableProxyState) + ) as state: + state.router = RouterData.from_router_data({ + "query": {}, + "token": token, + "sid": "test_sid", + }) + state_proxy = StateProxy(state) + [items_proxy] = state_proxy.data + + async def fail_exit(self: StateProxy, *exc_info: Any) -> None: + await asyncio.sleep(0) + msg = "cleanup failed" + raise RuntimeError(msg) + + monkeypatch.setattr(StateProxy, "__aexit__", fail_exit) + + with pytest.raises(RuntimeError, match="cleanup failed"): + async with items_proxy: + pass + + assert items_proxy._self_actx_state is None + + def test_override_base_method_skips_event_handler_wrapping(): """A method marked with __override_base_method__ should not be wrapped as an EventHandler.""" from reflex.state import _override_base_method