Minor pickle improvement 2

reflex/state.py

@@ -2105,14 +2105,30 @@ def __getstate__(self):
             state["__dict__"].pop("router", None)
             state["__dict__"].pop("router_data", None)
         # Never serialize parent_state or substates.
-        state["__dict__"]["parent_state"] = None
-        state["__dict__"]["substates"] = {}
+        state["__dict__"].pop("parent_state", None)
+        state["__dict__"].pop("substates", None)
+        state["__dict__"].pop("dirty_vars", None)
+        state["__dict__"].pop("dirty_substates", None)
         state["__dict__"].pop("_was_touched", None)
         # Remove all inherited vars.
         for inherited_var_name in self.inherited_vars:
             state["__dict__"].pop(inherited_var_name, None)
         return state
 
+    def __setstate__(self, state: dict[str, Any]):
+        """Set the state from redis deserialization.
+
+        This method is called by pickle to deserialize the object.
+
+        Args:
+            state: The state dict for deserialization.
+        """
+        state["__dict__"]["parent_state"] = None
+        state["__dict__"]["substates"] = {}
+        state["__dict__"]["dirty_vars"] = set()
+        state["__dict__"]["dirty_substates"] = set()
+        super().__setstate__(state)
+
     def _check_state_size(
         self,
         pickle_state_size: int,

tests/units/test_state.py

@@ -2513,10 +2513,7 @@ def test_mutable_copy(mutable_state: MutableTestState, copy_func: Callable):
         assert getattr(ms_copy, attr) is not getattr(mutable_state, attr)
     ms_copy.custom.array.append(42)
     assert "custom" in ms_copy.dirty_vars
-    if copy_func is copy.copy:
-        assert "custom" in mutable_state.dirty_vars
-    else:
-        assert not mutable_state.dirty_vars
+    assert not mutable_state.dirty_vars
 
 
 @pytest.mark.parametrize(