Save states with the schema hash 0.9 workers check

Since #7461, `_serialize` wrote the new `_to_schema` hash (field names and
types only) next to each pickled state. Workers of 0.9.x and 0.10.0a1 compare
that hash with their own, which includes the defaults, and replace the state
with a fresh one on any mismatch. A rolling deploy or a rollback sharing Redis
or the disk state manager with them therefore silently reset every session,
and the old worker wrote the fresh state back, so the data was lost for good.
The rest of the payload was already readable by those workers.

Save the previous format's hash (`_legacy_state_schema`, now cached like
`_to_schema`) next to the pickle, and record the new hash inside the pickled
dict under `__reflex_schema__`. Older workers keep that entry as a plain
instance attribute and write it back, which stays correct because they only
accept a state whose names, types and defaults match their own.
`_deserialize` checks the recorded hash when present, so defaults remain
outside the schema for this release, and falls back to the saved hash for
pickles from 0.10.0a2 and older releases.

Fixes #7470

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01SjCrbdguMkfe4bhbaREiwo

Author: claude

news/+state-schema-rollback.bugfix.md

@@ -0,0 +1 @@
+Workers of 0.9 and 0.10.0a1 no longer discard state saved by this release (they did since 0.10.0a2), so a rolling deploy or a rollback that shares Redis or the disk state manager with them keeps user sessions instead of silently resetting them. As before, those workers keep a saved state only while the declared defaults of its vars are unchanged; workers of this release ignore defaults.

reflex/state.py

@@ -97,6 +97,11 @@
 from reflex.utils import console, format, types
 from reflex.utils.exec import is_testing_env
 
+# Entry in each pickle holding the `_to_schema` hash, which `_deserialize`
+# checks and takes out of the restored state. The hash saved next to the pickle
+# is `_legacy_state_schema`, the only one 0.9 and 0.10.0a1 workers check.
+_SCHEMA_PICKLE_KEY = "__reflex_schema__"
+
 # Entries in each pickle for workers of the previous release, which kept the
 # dirty tracking and backend vars in the instance dict. Remove in 1.0.
 _PREVIOUS_RELEASE_PICKLE_KEYS: dict[str, Any] = {
@@ -1071,8 +1076,9 @@ def _init_var_dependency_dicts(cls):
                     parent_state.get_parent_state(),
                 )
 
-        # Reset cached schema value
+        # Reset cached schema values
         cls._to_schema.cache_clear()
+        _legacy_state_schema.cache_clear()
 
     @classmethod
     def _iter_functions(cls) -> Iterator[tuple[str, FunctionType]]:
@@ -2040,7 +2046,7 @@ def __getstate__(self):
         are stored on it, so inherited ones are not included either.
 
         Returns:
-            The state dict for serialization.
+            The state dict for serialization, with the schema hash of the class.
         """
         cls = type(self)
         fields = vars(self)
@@ -2049,8 +2055,13 @@ def __getstate__(self):
         for name, f in cls.__fields__.items():
             if f._owner is cls and name not in fields:
                 fields[name] = f.default_value()
-        # Empty entries that let workers of the previous release load it.
-        return {**fields, **_PREVIOUS_RELEASE_PICKLE_KEYS}
+        # The schema this release checks, and empty entries that let workers of
+        # the previous release load it.
+        return {
+            **fields,
+            _SCHEMA_PICKLE_KEY: cls._to_schema(),
+            **_PREVIOUS_RELEASE_PICKLE_KEYS,
+        }
 
     def __setstate__(self, state: builtins.dict[str, Any]):
         """Set the state from redis deserialization.
@@ -2135,7 +2146,8 @@ def _serialize(self) -> bytes:
         """
         payload = b""
         error = ""
-        self_schema = self._to_schema()
+        # Workers of 0.9 and 0.10.0a1 drop a state saved with any other hash.
+        self_schema = _legacy_state_schema(type(self))
         pickle_function = pickle.dumps
         try:
             payload = pickle.dumps((self_schema, self))
@@ -2207,16 +2219,21 @@ def _deserialize(
         except Exception as err:
             msg = f"Stored state could not be unpickled: {err!r}"
             raise StateSchemaMismatchError(msg) from err
-        if (
-            substate_schema != state._to_schema()
-            and substate_schema != _legacy_state_schema(type(state))
-        ):
+        # A state saved by 0.10.0a2 or an older release has no recorded schema.
+        schema = vars(state).pop(_SCHEMA_PICKLE_KEY, substate_schema)
+        if schema != state._to_schema() and schema != _legacy_state_schema(type(state)):
             raise StateSchemaMismatchError
         return state
 
 
+@functools.cache
 def _legacy_state_schema(state_cls: type[BaseState]) -> str:
-    """Hash the previous schema format to restore compatible persisted states.
+    """Hash the schema format of 0.9 and 0.10.0a1, defaults included.
+
+    It is saved next to each pickled state, so that workers of those releases
+    load it, and restores the states they saved. Like `_to_schema`, it is
+    cached until a state class is created or gets a new var, so a default
+    assigned in between, which those releases ignore, does not change it.
 
     Args:
         state_cls: The class whose previous schema to compute.

tests/units/test_state.py

@@ -69,6 +69,7 @@
 from reflex.istate.manager.token import BaseStateToken
 from reflex.istate.proxy import MutableProxy, StateProxy
 from reflex.state import (
+    _SCHEMA_PICKLE_KEY,
     BaseState,
     Delta,
     ImmutableStateError,
@@ -5573,26 +5574,71 @@ class DifferentName(BaseState):
     assert DifferentName._to_schema() != schema
 
 
+def _previous_release_schema(state_cls: type[BaseState], **types: str) -> str:
+    """Hash a schema as workers before 0.10.0a2 did, defaults included.
+
+    Args:
+        state_cls: The state class whose schema to hash.
+        **types: Serialized types to hash in place of the declared ones, by field name.
+
+    Returns:
+        The hash those workers require a saved state to carry.
+    """
+    return md5(
+        pickle.dumps(
+            sorted(
+                (
+                    name,
+                    types.get(name)
+                    or f"{declared.type_.__module__}.{declared.type_.__qualname__}",
+                    declared.default if is_serializable(declared.default) else None,
+                )
+                for name, declared in state_cls.get_fields().items()
+                if name in state_cls.base_vars
+            )
+        )
+    ).hexdigest()
+
+
+class _SavedState:
+    """Pickles a state the way a given release saves it."""
+
+    def __init__(self, state: BaseState, recorded_schema: str | None):
+        """Wrap the state to pickle.
+
+        Args:
+            state: The state to pickle.
+            recorded_schema: The schema hash recorded inside the pickle, or
+                None for releases that did not record one.
+        """
+        self.state = state
+        self.recorded_schema = recorded_schema
+
+    def __reduce__(self) -> tuple[Any, ...]:
+        """Reduce to the wrapped state with the chosen schema entry.
+
+        Returns:
+            The arguments to recreate the state with, and its state dict.
+        """
+        state_dict = self.state.__getstate__()
+        del state_dict[_SCHEMA_PICKLE_KEY]
+        if self.recorded_schema is not None:
+            state_dict[_SCHEMA_PICKLE_KEY] = self.recorded_schema
+        return object.__new__, (type(self.state),), state_dict
+
+
 @pytest.mark.parametrize("type_name", ["builtins.int", "builtins.str"])
 def test_deserialize_previous_schema_format(type_name: str):
     """Accept compatible hashes from the previous schema format only.
 
     Args:
         type_name: The type stored in the previous schema.
     """
-    previous_fields = [
-        (
-            name,
-            type_name
-            if name == "value"
-            else f"{declared.type_.__module__}.{declared.type_.__qualname__}",
-            declared.default if is_serializable(declared.default) else None,
-        )
-        for name, declared in DefaultSchemaState.get_fields().items()
-        if name in DefaultSchemaState.base_vars
-    ]
-    legacy_schema = md5(pickle.dumps(sorted(previous_fields))).hexdigest()
-    data = pickle.dumps((legacy_schema, DefaultSchemaState(value=99)))
+    legacy_schema = _previous_release_schema(DefaultSchemaState, value=type_name)
+    data = pickle.dumps((
+        legacy_schema,
+        _SavedState(DefaultSchemaState(value=99), None),
+    ))
     if type_name == "builtins.int":
         restored = BaseState._deserialize(data)
         assert isinstance(restored, DefaultSchemaState)
@@ -5602,6 +5648,80 @@ def test_deserialize_previous_schema_format(type_name: str):
             BaseState._deserialize(data)
 
 
+def test_serialized_state_passes_the_previous_release_schema_check():
+    """Workers before 0.10.0a2 load a state saved by this release.
+
+    They drop a saved state unless the hash saved with it equals their own,
+    which includes the defaults, so a rollback or a rolling deploy sharing the
+    state would otherwise reset every session.
+    """
+    saved_schema, restored = pickle.loads(DefaultSchemaState(value=99)._serialize())
+    assert saved_schema == _previous_release_schema(DefaultSchemaState)
+    assert isinstance(restored, DefaultSchemaState)
+    assert restored.value == 99
+
+
+@pytest.mark.parametrize(
+    ("saved_schema", "recorded_schema", "loads"),
+    [
+        # Saved by this release, or re-saved by an older worker that loaded it.
+        ("previous", "current", True),
+        # Saved by this release before a default changed.
+        ("other", "current", True),
+        # Saved by this release before a field was added, renamed or retyped.
+        ("other", "other", False),
+        # Saved by 0.10.0a2.
+        ("current", None, True),
+        # Saved by an older release.
+        ("previous", None, True),
+        ("other", None, False),
+    ],
+)
+def test_deserialize_checks_the_recorded_schema(
+    saved_schema: str, recorded_schema: str | None, loads: bool
+):
+    """The schema recorded in the pickle decides; older pickles fall back to the saved hash.
+
+    Args:
+        saved_schema: Which hash is saved next to the pickled state.
+        recorded_schema: Which hash is recorded inside it, if any.
+        loads: Whether the state is restored.
+    """
+    hashes = {
+        "current": DefaultSchemaState._to_schema(),
+        "previous": _previous_release_schema(DefaultSchemaState),
+        "other": md5(b"other").hexdigest(),
+    }
+    data = pickle.dumps((
+        hashes[saved_schema],
+        _SavedState(
+            DefaultSchemaState(value=99),
+            None if recorded_schema is None else hashes[recorded_schema],
+        ),
+    ))
+    if not loads:
+        with pytest.raises(StateSchemaMismatchError):
+            BaseState._deserialize(data)
+        return
+    restored = BaseState._deserialize(data)
+    assert isinstance(restored, DefaultSchemaState)
+    assert restored.value == 99
+    assert _SCHEMA_PICKLE_KEY not in vars(restored)
+
+
+def test_serialize_rehashes_the_previous_schema_after_add_var():
+    """A var added at runtime is part of the hash saved for older workers."""
+
+    class AddVarSchemaState(BaseState):
+        value: int = 1
+
+    state = AddVarSchemaState(_reflex_internal_init=True)  # pyright: ignore [reportCallIssue]
+    before, _ = pickle.loads(state._serialize())
+    AddVarSchemaState.add_var("added", int, 2)
+    after, _ = pickle.loads(state._serialize())
+    assert before != after == _previous_release_schema(AddVarSchemaState)
+
+
 def test_reset_client_storage_uses_declared_factories():
     """Hydration evaluates declared factories for all browser storage types."""
     calls = []