diff --git a/reflex/.templates/web/utils/state.js b/reflex/.templates/web/utils/state.js index 00224a7c56e..6d1c0d9f307 100644 --- a/reflex/.templates/web/utils/state.js +++ b/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); diff --git a/reflex/app.py b/reflex/app.py index dc91180c075..262ab204462 100644 --- a/reflex/app.py +++ b/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 diff --git a/reflex/middleware/hydrate_middleware.py b/reflex/middleware/hydrate_middleware.py index ec18939dee3..268dcf6c45f 100644 --- a/reflex/middleware/hydrate_middleware.py +++ b/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()) diff --git a/tests/units/utils/test_token_manager.py b/tests/units/utils/test_token_manager.py index d9d891e3253..f061574a7c9 100644 --- a/tests/units/utils/test_token_manager.py +++ b/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)) )