[ENG-3989] Ensure non-serialized states are present in StateManagerDi…

…sk (#4230)

If a non-root state was serialized, but its substates were not, then these
would not be populated when reloading the pickled state, because only substates
from the root were being populated with fresh versions.

Now, capture the substates from all fresh states and apply them to the
deserialized state for each substate to ensure that the entire state tree has
all substates instantiated after deserializing, even substates that were never
serialized originally.

reflex/state.py

@@ -2842,9 +2842,13 @@ async def populate_substates(
         for substate in state.get_substates():
             substate_token = _substate_key(client_token, substate)
 
+            fresh_instance = await root_state.get_state(substate)
             instance = await self.load_state(substate_token)
-            if instance is None:
-                instance = await root_state.get_state(substate)
+            if instance is not None:
+                # Ensure all substates exist, even if they weren't serialized previously.
+                instance.substates = fresh_instance.substates
+            else:
+                instance = fresh_instance
             state.substates[substate.get_name()] = instance
             instance.parent_state = state
 

tests/units/test_state.py

@@ -3313,3 +3313,36 @@ def handle_var(self):
 
     state.handle_supported_regular_vars()
     state.handle_non_var()
+
+
+@pytest.mark.asyncio
+async def test_deserialize_gc_state_disk(token):
+    """Test that a state can be deserialized from disk with a grandchild state.
+
+    Args:
+        token: A token.
+    """
+
+    class Root(BaseState):
+        pass
+
+    class State(Root):
+        num: int = 42
+
+    class Child(State):
+        foo: str = "bar"
+
+    dsm = StateManagerDisk(state=Root)
+    async with dsm.modify_state(token) as root:
+        s = await root.get_state(State)
+        s.num += 1
+        c = await root.get_state(Child)
+        assert s._get_was_touched()
+        assert not c._get_was_touched()
+
+    dsm2 = StateManagerDisk(state=Root)
+    root = await dsm2.get_state(token)
+    s = await root.get_state(State)
+    assert s.num == 43
+    c = await root.get_state(Child)
+    assert c.foo == "bar"