fix: Preserve object identity when pickling MutableProxy (#6097)

* fix: Preserve object identity when pickling MutableProxy

* thanks greptile, it wasn't a dataclass once

Author: benedikt-bartscher

reflex/istate/proxy.py

@@ -26,6 +26,7 @@
     from reflex.state import BaseState, StateUpdate
 
 T_STATE = TypeVar("T_STATE", bound="BaseState")
+T = TypeVar("T")
 
 
 class StateProxy(wrapt.ObjectProxy):
@@ -671,19 +672,23 @@ def __deepcopy__(self, memo: dict[int, Any] | None = None) -> Any:
         return copy.deepcopy(self.__wrapped__, memo=memo)
 
     def __reduce_ex__(self, protocol_version: SupportsIndex):
-        """Get the state for redis serialization.
+        """Serialize the wrapped object for pickle, stripping off the proxy.
 
-        This method is called by cloudpickle to serialize the object.
-
-        It explicitly serializes the wrapped object, stripping off the mutable proxy.
+        Returns a function that reconstructs to the wrapped object directly,
+        ensuring pickle's memo system correctly tracks object identity.
 
         Args:
             protocol_version: The protocol version.
 
         Returns:
-            Tuple of (wrapped class, empty args, class __getstate__)
+            Tuple that reconstructs to the wrapped object.
         """
-        return self.__wrapped__.__reduce_ex__(protocol_version)
+        return (_unwrap_for_pickle, (self.__wrapped__,))
+
+
+def _unwrap_for_pickle(obj: T) -> T:
+    """Return the object unchanged. Used by MutableProxy.__reduce_ex__."""
+    return obj
 
 
 @serializer

tests/units/istate/test_proxy.py

@@ -0,0 +1,37 @@
+"""Tests for MutableProxy pickle behavior."""
+
+import dataclasses
+import pickle
+
+import reflex as rx
+from reflex.istate.proxy import MutableProxy
+
+
+@dataclasses.dataclass
+class Item:
+    """Simple picklable object for testing."""
+
+    id: int
+
+
+class ProxyTestState(rx.State):
+    """Test state with a list field."""
+
+    items: list[Item] = []
+
+
+def test_mutable_proxy_pickle_preserves_object_identity():
+    """Test that same object referenced directly and via proxy maintains identity."""
+    state = ProxyTestState()
+    obj = Item(1)
+
+    data = {
+        "direct": [obj],
+        "proxied": [MutableProxy(obj, state, "items")],
+    }
+
+    unpickled = pickle.loads(pickle.dumps(data))
+
+    assert unpickled["direct"][0].id == 1
+    assert unpickled["proxied"][0].id == 1
+    assert unpickled["direct"][0] is unpickled["proxied"][0]

tests/units/test_state.py

@@ -4468,9 +4468,5 @@ async def test_rebind_mutable_proxy(mock_app: rx.App, token: str) -> None:
     ) as state:
         assert isinstance(state, MutableProxyState)
         assert state.data["a"] == [2, 3]
-        if isinstance(mock_app.state_manager, StateManagerRedis):
-            # In redis mode, the object identity does not persist across async with self calls.
-            assert state.data["b"] == [2]
-        else:
-            # In disk/memory mode, the fact that data["b"] was mutated via data["a"] persists.
-            assert state.data["b"] == [2, 3]
+        # Object identity persists across serialization, so data["b"] is also mutated.
+        assert state.data["b"] == [2, 3]