Send a "reconnect" hydrate whenever the socket reconnects

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

@@ -568,6 +568,7 @@ export const connect = async (
       !socket.current.wait_connect
     ) {
       socket.current.wait_connect = true;
+      socket.current.rehydrate = true;
       socket.current.io.opts.query = { token: getToken() }; // Update token for reconnect.
       socket.current.connect();
     }
@@ -615,6 +616,22 @@ export const connect = async (
     window.addEventListener("pagehide", pagehideHandler);
     window.addEventListener("beforeunload", disconnectTrigger);
     window.addEventListener("unload", disconnectTrigger);
+    if (socket.current.rehydrate) {
+      socket.current.rehydrate = false;
+      const events = initialEvents();
+      if (events.length > 0) {
+        // On reconnect, we only hydrate, do not re-run on_load events.
+        const hydrate_event = initialEvents()[0];
+        hydrate_event.payload.is_reconnect = true;
+        queueEvents(
+          [hydrate_event],
+          socket,
+          true,
+          navigate,
+          () => params.current,
+        );
+      }
+    }
     // Drain any initial events from the queue.
     while (event_queue.length > 0 && !event_processing) {
       await processEvent(socket.current, navigate, () => params.current);

reflex/app.py

@@ -2263,7 +2263,7 @@ async def link_token_to_sid(self, sid: str, token: str):
             await self.emit("new_token", new_token, to=sid)
 
         # Update client state to apply new sid/token for running background tasks.
-        async with self.app.modify_state(
+        async with self.app.state_manager.modify_state(
             _substate_key(new_token or token, self.app.state_manager.state)
         ) as state:
             state.router_data[constants.RouteVar.SESSION_ID] = sid

reflex/middleware/hydrate_middleware.py

@@ -35,11 +35,14 @@ async def preprocess(
         if event.name != get_hydrate_event(state):
             return None
 
-        # Clear client storage, to respect clearing cookies
-        state._reset_client_storage()
-
-        # Mark state as not hydrated (until on_loads are complete)
-        setattr(state, constants.CompileVars.IS_HYDRATED, False)
+        # In reconnect mode, don't reset client storage or call on_load.
+        is_reconnect = event.payload.get("is_reconnect", False)
+        if not is_reconnect:
+            # Clear client storage, to respect clearing cookies
+            state._reset_client_storage()
+
+            # Mark state as not hydrated (until on_loads are complete)
+            setattr(state, constants.CompileVars.IS_HYDRATED, False)
 
         # Get the initial state.
         delta = await _resolve_delta(state.dict())

tests/units/utils/test_token_manager.py

@@ -519,7 +519,7 @@ def new_event_namespace() -> EventNamespace:
         state.router_data = {}
 
         mock_app = Mock()
-        mock_app.modify_state = Mock(
+        mock_app.state_manager.modify_state = Mock(
             return_value=AsyncMock(__aenter__=AsyncMock(return_value=state))
         )