ENG-8664: MutableProxy.wrap_recursive fully rebinds existing proxy values

reflex/istate/proxy.py

@@ -509,11 +509,10 @@ def _wrap_recursive(self, value: Any) -> Any:
         # When called from dataclasses internal code, return the unwrapped value
         if self._is_called_from_dataclasses_internal():
             return value
-        # If we already have a proxy, make sure the state reference is up to date and return it.
+        # If we already have a proxy, unwrap and rewrap to make sure the state
+        # reference is up to date.
         if isinstance(value, MutableProxy):
-            if value._self_state is not self._self_state:
-                value._self_state = self._self_state
-            return value
+            value = value.__wrapped__
         # Recursively wrap mutable types.
         if is_mutable_type(type(value)):
             base_cls = globals()[self.__base_proxy__]

tests/units/test_state.py

@@ -35,6 +35,7 @@
 from reflex.istate.manager.redis import StateManagerRedis
 from reflex.state import (
     BaseState,
+    ImmutableMutableProxy,
     ImmutableStateError,
     MutableProxy,
     OnLoadInternalState,
@@ -4419,20 +4420,30 @@ async def test_rebind_mutable_proxy(mock_app: rx.App, token: str) -> None:
             "token": token,
             "sid": "test_sid",
         })
+        assert isinstance(state, MutableProxyState)
+        assert isinstance(state.data, MutableProxy)
+        assert not isinstance(state.data, ImmutableMutableProxy)
         state_proxy = StateProxy(state)
-        assert isinstance(state_proxy.data, MutableProxy)
+        assert isinstance(state_proxy.data, ImmutableMutableProxy)
     async with state_proxy:
+        # This assigns an ImmutableMutableProxy to data["a"].
         state_proxy.data["a"] = state_proxy.data["b"]
+    assert isinstance(state_proxy.data["a"], ImmutableMutableProxy)
     assert state_proxy.data["a"] is not state_proxy.data["b"]
     assert state_proxy.data["a"].__wrapped__ is state_proxy.data["b"].__wrapped__
 
+    # Rebinding with a non-proxy should return a MutableProxy object (not ImmutableMutableProxy).
+    assert isinstance(state_proxy.__wrapped__.data["a"], MutableProxy)
+    assert not isinstance(state_proxy.__wrapped__.data["a"], ImmutableMutableProxy)
+
     # Flush any oplock.
     await mock_app.state_manager.close()
 
     new_state_proxy = StateProxy(state)
     assert state_proxy is not new_state_proxy
     assert new_state_proxy.data["a"]._self_state is new_state_proxy
     assert state_proxy.data["a"]._self_state is state_proxy
+    assert state_proxy.__wrapped__.data["a"]._self_state is state_proxy.__wrapped__
 
     async with state_proxy:
         state_proxy.data["a"].append(3)