diff --git a/reflex/istate/proxy.py b/reflex/istate/proxy.py index d937d5ff23..80c3f83e1f 100644 --- a/reflex/istate/proxy.py +++ b/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 diff --git a/tests/units/istate/test_proxy.py b/tests/units/istate/test_proxy.py new file mode 100644 index 0000000000..b71d91e619 --- /dev/null +++ b/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] diff --git a/tests/units/test_state.py b/tests/units/test_state.py index ca41ac37ab..62cc193803 100644 --- a/tests/units/test_state.py +++ b/tests/units/test_state.py @@ -4453,9 +4453,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]