support async context for mutable proxies

[harsh21234i]

Fixes #6689

Summary

• add async context-manager support to MutableProxy and
  ImmutableMutableProxy
• refresh the proxy from its bound state field after entering the state
  context
• cover ImmutableMutableProxy mutation through async with proxy

Testing

• uv run pytest tests/units/
  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
• uv run ruff check reflex/istate/proxy.py tests/units/test_state.py
• uv run towncrier check --config pyproject.toml --dir . --compare-with
  upstream/main

[chatgpt-codex-connector]

💡 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-apps]

Greptile Summary

This PR adds async with proxy support to MutableProxy and ImmutableMutableProxy, allowing background-task proxies that wrap nested mutable state fields to enter and exit the owning StateProxy context directly. It also introduces an AccessSpec-based path-tracking system so each proxy knows how to re-fetch itself from the freshly-locked state on entry.

• MutableProxy.__aenter__/__aexit__: enters the bound StateProxy context, traverses the recorded access path to re-bind __wrapped__ to the fresh object, and unconditionally clears _self_actx_state in a finally block to prevent permanent lock-out on failures.
• Path tracking (AccessSpec, _self_path, __iter_path_proxies__): records how each nested proxy was reached (("attr", name), ("item", key)) and marks iteration-sourced proxies as un-refreshable; __dataclass_proxies__ key is fixed to (proxy_cls, wrapped_cls) to avoid cross-class collisions.
• __wrap_mutable_attrs__/_wrap_recursive_decorator: wraps dict.get and dict.setdefault return values with the correct path segment so those proxies can also be used as async context managers.

Confidence Score: 5/5

Safe to merge — the changes are well-scoped to the proxy layer, no state-manager or serialization logic is touched, and the existing test suite is augmented with six targeted async tests including edge-case recovery scenarios.

Path tracking is logically consistent across all access methods and the exception-safety issues from earlier review rounds have been resolved. No new unsafe patterns were introduced.

No files require special attention.

Important Files Changed

Filename	Overview
reflex/istate/proxy.py	Adds AccessSpec path tracking and __aenter__/__aexit__ to MutableProxy; fixes __dataclass_proxies__ key to include proxy class; adds __iter_path_proxies__ for un-refreshable iter-sourced proxies; extends __wrap_mutable_attrs__ and _wrap_recursive_decorator for dict get/setdefault path tracking. Previously-raised exception-safety concerns are addressed in this revision.
tests/units/test_state.py	Adds six new async tests covering: basic proxy context entry/exit, dict method paths (get/setdefault), iter-sourced proxy rejection, recovery from a failed __aenter__, and recovery when __aexit__ cleanup itself fails. Coverage is thorough.
news/6689.feature.md	Changelog entry for the new async context manager support on mutable state proxies.

_{Reviews (11): Last reviewed commit: "Optimize dataclass mutable proxy cache" | Re-trigger Greptile}

[codspeed-hq]

Merging this PR will improve performance by 9.3%

⚠️ Different runtime environments detected

Some benchmarks with significant performance changes were compared across different runtime environments,
which may affect the accuracy of the results.

Open the report in CodSpeed to investigate

⚡ 1 improved benchmark
✅ 25 untouched benchmarks
⏩ 8 skipped benchmarks¹

Performance Changes

	Benchmark	BASE	HEAD	Efficiency
⚡	test_var_access[mutable_list]	70.9 ms	64.8 ms	+9.3%

Tip

Curious why this is faster? Comment @codspeedbot explain why this is faster on this PR, or directly use the CodSpeed MCP with your agent.

---

_{Comparing harsh21234i:fix/mutable-proxy-async-context-6689 (a3bfa51) with main (9a5c4d3)}

Footnotes

1. 8 benchmarks were skipped, so the baseline results were used instead. If they were deleted from the codebase, click here and archive them to remove them from the performance reports. ↩

[harsh21234i]

Addressed the review feedback in d8f12f3.

Changes made:

• Added typed AccessSpec path entries for attr, item, and iter.
• Updated _wrap_recursive to append a single new path segment internally.
• Added _self_actx_state to internal proxy attrs.
• Rejected nested async with proxy usage explicitly.
• Made proxy __aexit__ idempotent when already exited.
• Made async refresh path replay defensive and fail clearly for iteration-sourced proxies.
• Added coverage for nested context rejection, idempotent exit, and iteration-derived proxy rejection.

Checks passed:

• uv run pytest tests/units/test_state.py::test_immutable_mutable_proxy_async_context_manager tests/units/ test_state.py::test_immutable_mutable_proxy_async_context_rejects_iter_proxy tests/units/test_state.py::test_rebind_mutable_proxy tests/units/istate/ test_proxy.py -q
• uv run ruff check reflex/istate/proxy.py tests/units/test_state.py
• uv run ruff format --check reflex/istate/proxy.py tests/units/test_state.py
• uv run pyright reflex/istate/proxy.py tests/units/test_state.py tests/units/istate/test_proxy.py

[FarhanAliRaza]

@codspeedbot fix this regression

[harsh21234i]

hey can you go through this @masenf @FarhanAliRaza

[FarhanAliRaza]

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.

[harsh21234i]

hey @FarhanAliRaza
Thanks for the detailed review. I addressed the confirmed dict.get() / setdefault() bug in 040ae098.

Changes made:

• dict.get(key) now records an item path when the key exists, so returned mutable proxies refresh to the actual
  child value.
• dict.setdefault(key, value) now records the inserted/existing item path.
• dict.get(missing_key, mutable_default) now uses the unrefreshable marker, so async with fails with the existing
  Unable to refresh mutable proxy error instead of rebinding to the parent container.
• Added a regression test covering the nested dict[str, dict] corruption case.
• Added the missing Raises: docstring section and clarified the overlapping async context error message.