Reassociate sid and token when connecting websocket (#5794)

* Reassociate sid and token when connecting websocket

If there are active background tasks for a given client_token, reassociate the
token and sid in the state when reconnecting so that updates from the
background task go to to the new websocket.

* Do not queue websocket events until the websocket is connected

* fully drain the queue on socket connect event

Author: masenf

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

@@ -478,8 +478,8 @@ export const queueEvents = async (
  * @param params The params object from React Router
  */
 export const processEvent = async (socket, navigate, params) => {
-  // Only proceed if the socket is up and no event in the queue uses state, otherwise we throw the event into the void
-  if (!socket && isStateful()) {
+  // Only proceed if the socket is up or no event in the queue uses state, otherwise we throw the event into the void
+  if (isStateful() && !(socket && socket.connected)) {
     return;
   }
 
@@ -576,11 +576,15 @@ export const connect = async (
   };
 
   // Once the socket is open, hydrate the page.
-  socket.current.on("connect", () => {
+  socket.current.on("connect", async () => {
     setConnectErrors([]);
     window.addEventListener("pagehide", pagehideHandler);
     window.addEventListener("beforeunload", disconnectTrigger);
     window.addEventListener("unload", disconnectTrigger);
+    // Drain any initial events from the queue.
+    while (event_queue.length > 0 && !event_processing) {
+      await processEvent(socket.current, navigate, () => params.current);
+    }
   });
 
   socket.current.on("connect_error", (error) => {
@@ -893,7 +897,7 @@ export const useEventLoop = (
   // Main event loop.
   useEffect(() => {
     // Skip if the backend is disabled
-    if (isBackendDisabled()) {
+    if (isBackendDisabled() || !socket.current || !socket.current.connected) {
       return;
     }
     (async () => {

reflex/app.py

@@ -2189,3 +2189,10 @@ async def link_token_to_sid(self, sid: str, token: str):
         if new_token:
             # Duplicate detected, emit new token to client
             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(
+            _substate_key(new_token or token, self.app.state_manager.state)
+        ) as state:
+            state.router_data[constants.RouteVar.SESSION_ID] = sid
+            state.router = RouterData.from_router_data(state.router_data)