Fix stale on_load during websocket connection

packages/reflex-base/news/+7487.bugfix.md

@@ -0,0 +1 @@
+Use the browser's current route when building the websocket connection's initial hydration event, so navigation during connection setup cannot run the previous page's `on_load` handler or redirect.

packages/reflex-base/src/reflex_base/.templates/web/utils/state.js

@@ -719,6 +719,12 @@ const withRouterData = (event, params) => {
   return event;
 };
 
+const currentRouteKey = () => {
+  const loc = locationRef.current ?? window.location;
+  const liveLoc = env.MOUNT_TARGET ? loc : window.location;
+  return loc.pathname + (liveLoc.search ?? "");
+};
+
 /**
  * Send an event to the server via REST.
  * @param event The current event.
@@ -916,19 +922,21 @@ export const connect = async (
 
   // Get backend URL object from the endpoint.
   const endpoint = getBackendURL(EVENTURL);
+  let bootRoute;
 
   // The hydrate event rides in the socket.io CONNECT packet, so the backend
   // starts loading state as soon as the namespace connects instead of after
   // an extra round trip for the connect acknowledgement. The keys are read by
   // the backend as CompileVars.CONNECT_AUTH_EVENT and CONNECT_AUTH_STATE_SCHEMA;
   // the state schema goes along on every (re)connect, since a reconnect can
   // land on another backend version.
-  const bootAuth = (first) => {
+  const bootAuth = (first) => (callback) => {
     const boot_event = withRouterData(app.initialEvents(first)[0], params);
+    bootRoute = currentRouteKey();
     recordSentStorageValues(boot_event);
     // Instrumentation hook (installed by reflex-otel): may add a traceparent.
     window.__reflex_otel?.onEventSend?.(boot_event);
-    return { event: boot_event, state_schema: app.stateSchema };
+    callback({ event: boot_event, state_schema: app.stateSchema });
   };
 
   // Create the socket. A new session's token is saved here, once the app has
@@ -949,6 +957,8 @@ export const connect = async (
     warmSocket ?? createSocket(endpoint, transports, session_token);
   warmSocket = null;
   cancelWarmup();
+  // Socket.IO calls this when the namespace connection is sent, so route data
+  // reflects any navigation that happened while the transport was connecting.
   socket.current.auth = bootAuth(true);
   socket.current.wait_connect = !socket.current.connected;
   // Ensure undefined fields in events are sent as null instead of removed
@@ -1063,6 +1073,19 @@ export const connect = async (
     window.__reflex_otel?.onSocketConnect?.();
     window.addEventListener("pagehide", pagehideHandler);
     window.addEventListener("beforeunload", disconnectTrigger);
+    // The connection auth event already loads its route, and every queued
+    // on_load loads the route current when it is sent. Keep the last queued
+    // on_load, after every navigation's storage update, only if the route
+    // changed after the auth event was built.
+    let keepOnLoad = bootRoute !== currentRouteKey();
+    for (let i = event_queue.length - 1; i >= 0; i--) {
+      if (event_queue[i].name !== app.on_load_internal_event) continue;
+      if (keepOnLoad) {
+        keepOnLoad = false;
+      } else {
+        event_queue.splice(i, 1);
+      }
+    }
     // Drain any initial events from the queue.
     while (event_queue.length > 0) {
       await processEvent(socket.current, navigate, params);

packages/reflex-base/src/reflex_base/compiler/templates.py

@@ -312,6 +312,7 @@ def context_template(
     state_str = (
         rf"""
 export const state_name = "{state_name}"
+export const on_load_internal_event = '{state_name}.{constants.CompileVars.ON_LOAD_INTERNAL}'
 
 export const exception_state_name = "{constants.CompileVars.FRONTEND_EXCEPTION_STATE_FULL}"
 
@@ -366,6 +367,7 @@ def context_template(
         if state_name
         else """
 export const state_name = undefined
+export const on_load_internal_event = undefined
 
 export const exception_state_name = undefined
 
@@ -426,6 +428,7 @@ def context_template(
   initialState,
   clientStorage,
   state_name,
+  on_load_internal_event,
   stateSchema,
   exception_state_name,
   onLoadInternalEvent,

tests/integration/tests_playwright/test_navigation_before_connect.py

@@ -0,0 +1,87 @@
+"""Regressions for navigation while the initial websocket is connecting."""
+
+from collections.abc import Generator
+
+import pytest
+from playwright.sync_api import Page, WebSocketRoute, expect
+
+from reflex.testing import AppHarness
+
+
+def NavigationBeforeConnectApp():
+    """App whose first route redirects if its stale on_load runs."""
+    import reflex as rx
+
+    class NavigationState(rx.State):
+        loads: list[str] = []
+
+        @rx.event
+        def slow_load(self):
+            self.loads.append("slow")
+            return rx.redirect("/hijack")
+
+        @rx.event
+        def other_load(self):
+            self.loads.append("other")
+
+    app = rx.App()
+
+    def page_content():
+        return rx.vstack(
+            rx.foreach(NavigationState.loads, lambda item: rx.text(item)),
+            rx.link("Other", href="/other", id="other-link"),
+        )
+
+    app.add_page(page_content, route="/slow", on_load=NavigationState.slow_load)
+    app.add_page(page_content, route="/other", on_load=NavigationState.other_load)
+    app.add_page(page_content, route="/hijack")
+
+
+@pytest.fixture(scope="module")
+def navigation_before_connect_app(
+    app_harness_env: type[AppHarness],
+    tmp_path_factory: pytest.TempPathFactory,
+) -> Generator[AppHarness, None, None]:
+    """Start the app in dev or production mode.
+
+    Yields:
+        The running app harness.
+    """
+    with app_harness_env.create(
+        root=tmp_path_factory.mktemp("navigation_before_connect"),
+        app_name=f"navigationbeforeconnect_{app_harness_env.__name__.lower()}",
+        app_source=NavigationBeforeConnectApp,
+    ) as harness:
+        yield harness
+
+
+def test_navigation_before_socket_connect_uses_current_route(
+    navigation_before_connect_app: AppHarness, page: Page
+):
+    """A delayed connect must not run the old page's on_load or redirect."""
+    assert navigation_before_connect_app.frontend_url is not None
+    frontend_url = navigation_before_connect_app.frontend_url.rstrip("/")
+    sockets: list[WebSocketRoute] = []
+    page_errors: list[str] = []
+    page.on("pageerror", lambda error: page_errors.append(str(error)))
+
+    def hold_socket(socket: WebSocketRoute) -> None:
+        sockets.append(socket)
+
+    page.route_web_socket("**/_event/**", hold_socket)
+
+    page.goto(f"{frontend_url}/slow")
+    expect(page.get_by_role("link", name="Other")).to_be_visible()
+    if not sockets:
+        page.wait_for_event("websocket", timeout=60000)
+    assert len(sockets) == 1, page_errors
+
+    page.get_by_role("link", name="Other").click()
+    expect(page).to_have_url(f"{frontend_url}/other")
+
+    sockets[0].connect_to_server()
+
+    expect(page.locator("body")).to_contain_text("other")
+    expect(page.get_by_text("other", exact=True)).to_have_count(1)
+    expect(page).to_have_url(f"{frontend_url}/other")
+    expect(page.locator("body")).not_to_contain_text("slow")