ENG-6769: Fix module app.app not found during on_load

Author: masenf

reflex/state.py

@@ -2463,16 +2463,20 @@ class OnLoadInternalState(State):
     This is a separate substate to avoid deserializing the entire state tree for every page navigation.
     """
 
+    # Cannot properly annotate this as `App` due to circular import issues.
+    _app_ref: ClassVar = None
+
     def on_load_internal(self) -> list[Event | EventSpec | event.EventCallback] | None:
         """Queue on_load handlers for the current page.
 
         Returns:
             The list of events to queue for on load handling.
         """
-        # Do not app._compile()!  It should be already compiled by now.
-        load_events = prerequisites.get_and_validate_app().app.get_load_events(
-            self.router._page.path
-        )
+        app = type(self)._app_ref or prerequisites.get_and_validate_app().app
+        # Cache the app reference for subsequent calls.
+        if type(self)._app_ref is None:
+            type(self)._app_ref = app
+        load_events = app.get_load_events(self.router._page.path)
         if not load_events:
             self.is_hydrated = True
             return None  # Fast path for navigation with no on_load events defined.
@@ -2646,6 +2650,9 @@ def reload_state_module(
         state: Recursive argument for the state class to reload.
 
     """
+    # Reset the _app_ref of OnLoadInternalState to avoid stale references.
+    if state is OnLoadInternalState:
+        state._app_ref = None
     # Clean out all potentially dirty states of reloaded modules.
     for pd_state in tuple(state._potentially_dirty_states):
         with contextlib.suppress(ValueError):

tests/units/test_app.py

@@ -7,6 +7,7 @@
 from collections.abc import Generator
 from contextlib import nullcontext as does_not_raise
 from pathlib import Path
+from typing import ClassVar
 from unittest.mock import AsyncMock
 
 import pytest
@@ -945,6 +946,7 @@ class DynamicState(BaseState):
     is_hydrated: bool = False
     loaded: int = 0
     counter: int = 0
+    _app_ref: ClassVar = None
 
     @rx.event
     def on_load(self):
@@ -982,6 +984,7 @@ def test_dynamic_arg_shadow(
         app_module_mock: Mocked app module.
         mocker: pytest mocker object.
     """
+    DynamicState._app_ref = None
     arg_name = "counter"
     route = f"/test/[{arg_name}]"
     app = app_module_mock.app = App(_state=DynamicState)
@@ -1030,6 +1033,7 @@ async def test_dynamic_route_var_route_change_completed_on_load(
         app_module_mock: Mocked app module.
         mocker: pytest mocker object.
     """
+    DynamicState._app_ref = None
     arg_name = "dynamic"
     route = f"test/[{arg_name}]"
     app = app_module_mock.app = App(_state=DynamicState)

tests/units/test_state.py

@@ -2948,6 +2948,7 @@ async def test_preprocess(
         expected: Expected delta.
         mocker: pytest mock object.
     """
+    OnLoadInternalState._app_ref = None
     mocker.patch(
         "reflex.state.State.class_subclasses", {test_state, OnLoadInternalState}
     )
@@ -3002,6 +3003,7 @@ async def test_preprocess_multiple_load_events(
         token: A token.
         mocker: pytest mock object.
     """
+    OnLoadInternalState._app_ref = None
     mocker.patch(
         "reflex.state.State.class_subclasses", {OnLoadState, OnLoadInternalState}
     )