support async context for mutable proxies#6691
Conversation
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: e0d0b19a90
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
Greptile SummaryThis PR adds
Confidence Score: 5/5The core async context manager logic is correct and all critical edge cases (failed enter, failed cleanup, double-entry guard) are covered by tests. Path tracking works correctly for attr/item access patterns. The _wrap_recursive_decorator gap for dict.get/setdefault results is an unlikely usage pattern and does not affect normal operation of the feature. reflex/istate/proxy.py — the _wrap_recursive_decorator method and whether it should forward the item key as a path segment for wrap_mutable_attrs methods. Important Files Changed
Reviews (7): Last reviewed commit: "Clear mutable proxy context after cleanu..." | Re-trigger Greptile |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
Merging this PR will degrade performance by 18.34%
Warning Please fix the performance issues or acknowledge them on CodSpeed. Performance Changes
Tip Investigate this regression by commenting Comparing 
Footnotes
|
|
Addressed the review feedback in d8f12f3. Changes made:
Checks passed:
|
|
@codspeedbot fix this regression |
|
hey can you go through this @masenf @FarhanAliRaza |
FarhanAliRaza
left a comment
FarhanAliRaza
left a comment
There was a problem hiding this comment.
Thanks — the overall design here is solid: delegating to the bound StateProxy, replaying a recorded access path on the refreshed state, and poisoning iteration-sourced proxies so they fail loudly is the right shape, and the exception paths are carefully handled and tested. Tests, ruff, and pyright all pass locally.
Requesting changes for one confirmed correctness bug (reproduced locally): proxies returned by dict.get()/setdefault() carry their parent's access path, so async with on them silently rebinds the proxy to the parent container — in the worst case writes land in the wrong container with no error. Details inline, along with a few smaller comments.
| 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_path = refreshed_value._self_path | ||
| else: | ||
| self._raise_refresh_error() |
There was a problem hiding this comment.
Confirmed bug (reproduced): proxies returned by dict.get() / setdefault() silently rebind to the parent container here.
Those methods are in __wrap_mutable_attrs__ and route through _wrap_recursive_decorator, which calls self._wrap_recursive(wrapped(*args, **kwargs)) with no new_path_segment — so the child proxy inherits the parent's _self_path. This guard can't catch the mismatch because _self_field_name is the top-level field name for every descendant proxy:
data_proxy = state_proxy.data # dict[str, list[int]] = {"a": [1], "b": [2]}
child = data_proxy.get("a") # wraps [1], but path=() like its parent
async with child:
# child.__wrapped__ is now {"a": [1], "b": [2]} — the whole dict!
child.append(3) # AttributeErrorchild.append(...) at least raises AttributeError, but with homogeneous nesting (dict[str, dict]) writes like child["x"] = 5 land silently in the wrong container — data corruption with no error. Note data_proxy["a"] (via __getitem__) works correctly; only the two decorator-routed methods are affected.
Suggested fix: either pass ("item", args[0]) as the segment in _wrap_recursive_decorator (always correct for setdefault, correct for get whenever the value actually came from the dict), or — safer and minimal — tag decorator-returned proxies with the unresolvable sentinel (like the iter marker) so __aenter__ raises the existing "Unable to refresh" error instead of silently rebinding. The sentinel handles get(key, default) returning the default, which can't be represented as a path at all.
| """ | ||
| if self._self_actx_state is not None: | ||
| msg = ( | ||
| "Mutable proxy is already mutable. Do not nest `async with proxy` " |
There was a problem hiding this comment.
Minor: if the same proxy object is shared between two asyncio tasks, task B hits this immediately with "Do not nest" — which is misleading for the cross-task case — whereas async with state_proxy in task B would have blocked on _self_actx_lock and proceeded. Uncommon in practice (each StateProxy.__getattr__ access mints a fresh proxy), but worth either clarifying the message or keying the guard to the current task.
| """Enter the async context manager protocol through the bound state. | ||
|
|
||
| Returns: | ||
| This proxy refreshed from the current state field. |
There was a problem hiding this comment.
Nit: this raises RuntimeError (nesting guard, refresh failure) but the docstring has no Raises: section.
| self._self_actx_state = context_state | ||
| aenter_ok = False | ||
| try: | ||
| state = await context_state.__aenter__() |
There was a problem hiding this comment.
Worth noting: for a MutableProxy from a regular (non-background) handler, _self_state is a raw BaseState, whose __aenter__ is a no-op returning self — so async with proxy succeeds but does nothing meaningful (no lock, no refresh). Functionally harmless since mutations persist through the normal handler flow, but it reads as if it locks something. Might deserve at least a comment here, or a docs note.
| path = ( | ||
| _ROOT_ITER_PATH | ||
| if new_path_segment is _ITER_ACCESS_SPEC and not path | ||
| else (*path, new_path_segment) |
There was a problem hiding this comment.
Perf note: every nested mutable attribute/item access now allocates a path tuple (here plus the ("attr", name) / ("item", key) tuples at the call sites), paid by all event handlers even though only background-task async with proxy ever consumes the path. It's marginal next to the proxy allocation that already happens per access, but flagging since this is the state read hot path — fine if you've considered it acceptable.
| 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_ACCESS_SPEC) |
There was a problem hiding this comment.
Not a blocker — the poison marker is a sound choice — but for lists specifically you could record ("item", i) via enumerate instead, which would make iteration-sourced list proxies enterable (for row in self.data: async with row: ...). Could be a follow-up.
3 participants
Fixes #6689
Summary
ImmutableMutableProxy
context
async with proxyTesting
test_state.py::test_immutable_mutable_proxy_async_context_manager tests/
units/test_state.py::test_rebind_mutable_proxy tests/units/istate/
test_proxy.py -q
upstream/main