fix: Preserve object identity when pickling MutableProxy

[benedikt-bartscher]

No description provided.

[codspeed-hq]

Merging this PR will not alter performance

✅ 8 untouched benchmarks

---

_{Comparing benedikt-bartscher:preserve-object-identity-mutable-proxy-pickle (38f0745) with main (be43052)}

[greptile-apps]

Greptile Overview

Greptile Summary

This PR fixes a bug where MutableProxy instances didn't preserve object identity during pickle serialization. The previous implementation delegated to the wrapped object's __reduce_ex__ method, which caused pickle's memo system to create separate instances for the same object when it appeared both directly and wrapped in a MutableProxy. The new implementation returns a custom reducer function _unwrap_for_pickle that ensures the memo system correctly tracks object identity.

The fix affects state serialization when using Redis state manager, ensuring that when the same object is referenced multiple times (e.g., data["a"] and data["b"] pointing to the same list), mutations to one reference correctly affect the other after deserialization.

Confidence Score: 4/5

• This PR is safe to merge with low risk
• The fix correctly addresses a real bug with pickle object identity preservation. The implementation is clean and well-tested. Minor style issue exists in the new test file (redundant __init__ in dataclass), but this doesn't affect functionality. The change is focused and doesn't introduce new risks.
• tests/units/istate/test_proxy.py has a minor redundancy that could be cleaned up

Important Files Changed

Filename	Overview
reflex/istate/proxy.py	Improved pickle serialization to preserve object identity using a helper function instead of delegating to wrapped object's __reduce_ex__
tests/units/istate/test_proxy.py	New test verifying that pickle preserves object identity when same object is referenced both directly and via MutableProxy; has redundant __init__ in dataclass
tests/units/test_state.py	Updated test expectations to verify that object identity now correctly persists across serialization in both Redis and non-Redis modes

Sequence Diagram

sequenceDiagram
    participant App as Application
    participant SM as StateManager
    participant Pickle as pickle.dumps
    participant MP as MutableProxy
    participant WO as Wrapped Object
    participant Memo as Pickle Memo

    Note over App,Memo: Before: Object identity lost
    App->>SM: Store state with MutableProxy
    SM->>Pickle: pickle.dumps(data)
    Pickle->>MP: __reduce_ex__(protocol)
    MP->>WO: delegate to __reduce_ex__
    WO-->>Pickle: (class, args, state)
    Note over Memo: New memo entry for wrapped object
    Pickle->>Memo: Store object reference
    Note over Memo: Memo doesn't recognize proxy wraps same object
    Pickle-->>SM: Serialized bytes (identity lost)

    Note over App,Memo: After: Object identity preserved
    App->>SM: Store state with MutableProxy
    SM->>Pickle: pickle.dumps(data)
    Pickle->>MP: __reduce_ex__(protocol)
    MP-->>Pickle: (_unwrap_for_pickle, (wrapped_obj,))
    Note over Memo: Memo sees actual wrapped object
    Pickle->>Memo: Check if wrapped_obj already in memo
    alt Object already in memo
        Memo-->>Pickle: Return existing reference
    else New object
        Memo->>Pickle: Store new reference
    end
    Pickle-->>SM: Serialized bytes (identity preserved)
    
    Note over App,Memo: Deserialization
    SM->>Pickle: pickle.loads(bytes)
    Pickle->>Memo: Reconstruct objects using memo
    Note over Memo: Same object referenced multiple times
    Pickle-->>SM: Reconstructed data (same object identity)

Loading

[greptile-apps]

_{1 file reviewed, 1 comment}

_{Edit Code Review Agent Settings | Greptile}