add partial hydrate event

reflex/.templates/web/utils/state.js

@@ -471,12 +471,7 @@ export const connect = async (
   const getSubstateFromUpdate = (update, substate_name) => {
     if (update.__patch) {
       if (last_substate_hash[substate_name] !== update.__previous_hash) {
-        throw new Error(
-          "Patch received out of order" +
-            update.__hash +
-            " " +
-            last_substate_hash[substate_name]
-        );
+        throw new Error("Patch received out of order");
       }
       last_substate_hash[substate_name] = update.__hash;
       return applyPatch(last_substate_info, update.__patch).newDocument;
@@ -488,14 +483,22 @@ export const connect = async (
 
   // On each received message, queue the updates and events.
   socket.current.on("event", async (update) => {
+    const failed_states = [];
     for (const substate in update.delta) {
-      const new_substate_info = getSubstateFromUpdate(
-        update.delta[substate],
-        substate
-      );
+      try {
+        const new_substate_info = getSubstateFromUpdate(
+          update.delta[substate],
+          substate
+        );
+      } catch (e) {
+        console.error("Received patch out of order", e);
+        states_failed.push(substate);
+        continue;
+      }
       last_substate_info[substate] = new_substate_info;
       dispatch[substate](new_substate_info);
     }
+    // TODO: Handle failed states
     applyClientStorageDelta(client_storage, update.delta);
     event_processing = !update.final;
     if (update.events) {

reflex/compiler/utils.py

@@ -202,10 +202,10 @@ def compile_state(state: Type[BaseState]) -> dict:
                 console.warn(
                     f"Had to get initial state in a thread 🤮 {resolved_initial_state}",
                 )
-                return resolved_initial_state.data
+                return dict(**resolved_initial_state.data)
 
     # Normally the compile runs before any event loop starts, we asyncio.run is available for calling.
-    return asyncio.run(_resolve_delta(initial_state)).data
+    return dict(**asyncio.run(_resolve_delta(initial_state)).data)
 
 
 def _compile_client_storage_field(

reflex/constants/compiler.py

@@ -55,6 +55,8 @@ class CompileVars(SimpleNamespace):
     EVENTS = "events"
     # The name of the initial hydrate event.
     HYDRATE = "hydrate"
+    # The name of the partial hydrate event.
+    PARTIAL_HYDRATE = "partial_hydrate"
     # The name of the is_hydrated variable.
     IS_HYDRATED = "is_hydrated"
     # The name of the function to add events to the queue.

reflex/event.py

@@ -1194,6 +1194,18 @@ def get_hydrate_event(state: BaseState) -> str:
     return get_event(state, constants.CompileVars.HYDRATE)
 
 
+def get_partial_hydrate_event(state: BaseState) -> str:
+    """Get the name of the partial hydrate event for the state.
+
+    Args:
+        state: The state.
+
+    Returns:
+        The name of the partial hydrate event.
+    """
+    return get_event(state, constants.CompileVars.PARTIAL_HYDRATE)
+
+
 def call_event_handler(
     event_callback: EventHandler | EventSpec,
     event_spec: ArgsSpec | Sequence[ArgsSpec],

reflex/middleware/hydrate_middleware.py

@@ -3,10 +3,10 @@
 from __future__ import annotations
 
 import dataclasses
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, ChainMap
 
 from reflex import constants
-from reflex.event import Event, get_hydrate_event
+from reflex.event import Event, get_hydrate_event, get_partial_hydrate_event
 from reflex.middleware.middleware import Middleware
 from reflex.state import BaseState, StateDelta, StateUpdate, _resolve_delta
 
@@ -54,3 +54,50 @@ class HydrateMiddleware(Middleware):
 
         # Return the state update.
         return StateUpdate(delta=delta, events=[])
+
+
+@dataclasses.dataclass(init=True)
+class PartialHyderateMiddleware(Middleware):
+    """Middleware to handle partial app hydration."""
+
+    async def preprocess(
+        self, app: App, state: BaseState, event: Event
+    ) -> StateUpdate | None:
+        """Preprocess the event.
+
+        Args:
+            app: The app to apply the middleware to."
+            state: The client state.""
+            event: The event to preprocess.""
+
+        Returns:
+            An optional delta or list of state updates to return.""
+        """
+        # If this is not the partial hydrate event, return None
+        if event.name != get_partial_hydrate_event(state):
+            return None
+
+        substates_names = event.payload.get("states", [])
+        if not substates_names:
+            return None
+
+        substates = [
+            substate
+            for substate_name in substates_names
+            if (substate := state.get_substate(substate_name)) is not None
+        ]
+
+        delta = await _resolve_delta(
+            StateDelta(
+                ChainMap(*[substate.dict() for substate in substates]),
+                client_token=state.router.session.client_token,
+                flush=True,
+            )
+        )
+
+        # since a full dict was captured, clean any dirtiness
+        for substate in substates:
+            substate._clean()
+
+        # Return the state update.
+        return StateUpdate(delta=delta, events=[])

reflex/state.py

@@ -26,6 +26,7 @@ from typing import (
     Callable,
     ClassVar,
     Dict,
+    Mapping,
     NamedTuple,
     Optional,
     Sequence,
@@ -117,7 +118,7 @@ var = computed_var
 class StateDelta:
     """A dictionary representing the state delta."""
 
-    data: dict[str, Any] = dataclasses.field(default_factory=dict)
+    data: Mapping[str, Any] = dataclasses.field(default_factory=dict)
     client_token: str | None = dataclasses.field(default=None)
     flush: bool = dataclasses.field(default=False)
 
@@ -132,23 +133,6 @@ class StateDelta:
         """
         return self.data[key]
 
-    def __setitem__(self, key: str, value: Any):
-        """Set the item in the delta.
-
-        Args:
-            key: The key to set.
-            value: The value to set.
-        """
-        self.data[key] = value
-
-    def __delitem__(self, key: str):
-        """Delete the item from the delta.
-
-        Args:
-            key: The key to delete.
-        """
-        del self.data[key]
-
     def __iter__(self) -> Any:
         """Iterate over the delta.
 
@@ -190,7 +174,7 @@ class StateDelta:
         Returns:
             The reversed delta.
         """
-        return reversed(self.data)
+        return reversed(dict(**self.data))
 
     def values(self):
         """Get the values of the delta.