regres: MutableProxy identity mismatch when nested in Dict/List during State serialization

[ShaikZaheer-developer]

If a MutableProxy (used for state reactivity) is nested inside a standard Python collection (such as a list or dict) inside the rx.State, the object identity of the MutableProxy is not maintained during the pickling/unpickling operation. As Reflex uses object identities to maintain "deltas" and push updates to the frontend, this results in the UI ceasing to respond to state changes once a session is reloaded from a Redis or disk-backed cache.

To Reproduce Steps to reproduce the behavior:

Code/Link to Repo:

import reflex as rx
import pickle
from reflex.state import MutableProxy

   # 1. Define a base object that Reflex would wrap in a proxy
   class SubState(rx.Base):
    value: int = 0

def reproduce_identity_loss():
    # 2. Simulate a proxy being put into a nested structure
    original_obj = SubState()
    proxy = MutableProxy(original_obj)
    state_dict = {"data": proxy}
    original_id = id(state_dict["data"])
    
    # 3. Simulate session serialization (Pickling)
    serialized = pickle.dumps(state_dict)
    reloaded_dict = pickle.loads(serialized)
    
    # 4. Compare IDs
    new_id = id(reloaded_dict["data"])
    
    print(f"Original ID: {original_id}")
    print(f"Reloaded ID: {new_id}")
    
    if original_id != new_id:
        raise AssertionError("Identity Mismatch: The state tracker will lose track of this object!")

if __name__ == "__main__":
    reproduce_identity_loss()

The id() of the MutableProxy should remain the same or be automatically re-registered with the State tracker when unpickling. This ensures that even after a session refresh or a server restart (where state is retrieved from a cache), the frontend stays responsive to changes in the nested data structures.

Specifics:

Python Version: 3.10 or 3.11

Reflex Version: 0.8.26 (Latest)

OS: (Insert your OS here, e.g., Windows 11 / macOS / Ubuntu)

Browser (Optional): Chrome/Firefox

Additional context :

This seems to be a regression/edge-case issue related to the recent fixes in #6097 and #6105. Although top-level state variables are properly addressed, the recursive tracking for nested proxies in collections appears to drop the memory reference in setstate. This is especially important for Enterprise users running Reflex in production with persistent session storage.

[masenf]

As Reflex uses object identities to maintain "deltas" and push updates to the frontend

Reflex doesn't use object identities to maintain deltas. The framework tracks State vars that are set during event processing, and constructs the delta based on those "dirty vars" and computed vars that depend on those dirty vars.

The id() of the MutableProxy should remain the same or be automatically re-registered with the State tracker when unpickling.

Anytime an object is unpickled, it is by definition creating a copy of the object. The unpickled object should NOT retain the identity of the original proxy and in fact, in the unpickled object the proxy is nowhere to be seen. This is by design. If the unpickled object retained the same MutableProxy reference, then the copy would be tracking changes as if it were the original object, which is not correct.

When reflex unpickles a state, the resulting state is populated with non-proxy copies of the previously serialized state. These objects get re-wrapped in new MutableProxy instances tied to the new copy of the state when subsequently accessed through the state. Generally speaking, it's not possible to preserve cross-state MutableProxy references across event boundaries.

this results in the UI ceasing to respond to state changes once a session is reloaded from a Redis or disk-backed cache.

Do you have a reproducible example of this effect? That sounds like an actual bug. The provided repro code seems to be working as I would expect

This seems to be a regression/edge-case issue related to the recent fixes in #6097 and #6105.

The intent of #6097 was to allow different MutableProxy instances wrapping the same object to have referential integrity within the pickle, so when they get unpickled, both objects remain the same, instead of being separate copies.

import reflex as rx
import pickle
from reflex.state import MutableProxy


# 1. Define a base object that Reflex would wrap in a proxy
class SubState(rx.Base):
    value: int = 0


class AState(rx.State):
    myobj: SubState = SubState()


def reproduce_identity_loss():
    # 2. Simulate a proxy being put into a nested structure
    a_state = AState(_reflex_internal_init=True)
    state_dict = {"data": {"a": a_state.myobj, "b": a_state.myobj}}
    # The proxies are different objects
    assert state_dict["data"]["a"] is not state_dict["data"]["b"]
    # The wrapped objects are the identical
    assert state_dict["data"]["a"].__wrapped__ is state_dict["data"]["b"].__wrapped__

    # 3. Simulate session serialization (Pickling)
    serialized = pickle.dumps(state_dict)
    reloaded_dict = pickle.loads(serialized)

    # Deserialized objects are no longer proxies
    assert type(reloaded_dict["data"]["a"]) is not MutableProxy
    # Internal references remain the same object after deserialization
    assert reloaded_dict["data"]["a"] is reloaded_dict["data"]["b"]


if __name__ == "__main__":
    reproduce_identity_loss()

On previous versions (<=0.8.26), this last assertion will fail.

I don't see how #6105 would be related to the issue at hand.

---

Short story long, I don't see a bug here. If you disagree, please reopen with more realistic repro code.