reregister memo in app harness

news/+harness-real-app-no-reload.bugfix.md

@@ -0,0 +1 @@
+`AppHarness` no longer destructively reloads a real (pre-imported) app on start: reloading only re-ran the top module, silently dropping every page, `rx.State` subclass, and `@rx.memo` defined in a submodule (flaky "Invalid path" compile errors and unresolved `$/app_components/...` memo imports). Such apps are now imported once and used as-is (`reload=False`); generated `app_source` apps are unaffected.

packages/reflex-base/news/+memo-harness-reregister.bugfix.md

@@ -0,0 +1 @@
+Re-register `@rx.memo` definitions on use so a cleared `MEMOS` registry (e.g. an `AppHarness` reload, which only re-imports the top app module) still emits every memo module the compiled tree references, instead of failing to resolve them at runtime.

packages/reflex-base/src/reflex_base/components/memo.py

@@ -492,6 +492,35 @@ def _register_memo_definition(definition: MemoDefinition) -> None:
     MEMOS[key] = definition
 
 
+def _reregister_used_memo(definition: MemoDefinition) -> None:
+    """Re-register a used ``@rx.memo`` so a cleared ``MEMOS`` is repopulated from
+    usage. Passthrough auto-memos (defined in this module) are tracked
+    separately and stay out of ``MEMOS``.
+    """
+    if definition.fn.__module__ != __name__:
+        _register_memo_definition(definition)
+
+
+def materialize_registered_memo_bodies() -> None:
+    """Evaluate every registered memo body, until ``MEMOS`` stops growing.
+
+    Reading a memo body re-registers any nested ``@rx.memo`` it references (see
+    :func:`_reregister_used_memo`), so the compiler must do this before it
+    snapshots ``MEMOS`` — else a dependency surfaced only while compiling a lazy
+    var-memo body is left out. Bodies cache, so re-reading one is a no-op.
+    """
+    evaluated: set[tuple[str, str | None]] = set()
+    while pending := [
+        (key, definition) for key, definition in MEMOS.items() if key not in evaluated
+    ]:
+        for key, definition in pending:
+            evaluated.add(key)
+            if isinstance(definition, MemoFunctionDefinition):
+                _ = definition.function
+            elif isinstance(definition, MemoComponentDefinition):
+                _ = definition.component
+
+
 def _annotation_inner_type(annotation: Any) -> Any:
     """Unwrap a Var-like annotation to its inner type.
 
@@ -1528,6 +1557,7 @@ def call(self, *args: Any, **kwargs: Any) -> Var:
         Returns:
             The function call var.
         """
+        _reregister_used_memo(self._definition)
         return self._imported_var.call(
             *_bind_function_runtime_args(self._definition, *args, **kwargs)
         )
@@ -1542,6 +1572,7 @@ def partial(self, *args: Any, **kwargs: Any) -> FunctionVar:
         Returns:
             The partially applied function var.
         """
+        _reregister_used_memo(self._definition)
         return self._imported_var.partial(
             *_bind_function_runtime_args(self._definition, *args, **kwargs)
         )
@@ -1552,6 +1583,7 @@ def _as_var(self) -> FunctionVar:
         Returns:
             The imported function var.
         """
+        _reregister_used_memo(self._definition)
         return self._imported_var
 
 
@@ -1590,6 +1622,7 @@ def __call__(self, *children: Any, **props: Any) -> MemoComponent:
             The rendered memo component.
         """
         definition = self._definition
+        _reregister_used_memo(definition)
         rest_param = self._rest_param
 
         # Validate positional children usage and reserved keywords.
@@ -1661,6 +1694,7 @@ def _as_var(self) -> Var:
         Returns:
             The imported component var.
         """
+        _reregister_used_memo(self._definition)
         return _component_import_var(
             self._definition.export_name, self._definition.source_module
         )
@@ -2046,5 +2080,6 @@ def memo(
     "MemoFunctionDefinition",
     "create_component_memo",
     "create_passthrough_component_memo",
+    "materialize_registered_memo_bodies",
     "memo",
 ]

pyi_hashes.json

@@ -120,5 +120,5 @@
   "packages/reflex-components-sonner/src/reflex_components_sonner/toast.pyi": "58521fcd1b514804f534d97624e82c9a",
   "reflex/__init__.pyi": "56385a4f0d9431eb0056dbc5553a58f9",
   "reflex/components/__init__.pyi": "9facd05a776d0641432696bbf8e34388",
-  "reflex/experimental/memo.pyi": "36e5d5f97eb64e94c0974e909e7e2952"
+  "reflex/experimental/memo.pyi": "63d6db54828066685da93102d94220d4"
 }

reflex/compiler/compiler.py

@@ -24,6 +24,7 @@
     MemoDefinition,
     MemoFunctionDefinition,
     create_component_memo,
+    materialize_registered_memo_bodies,
     reset_memo_component_classes,
 )
 from reflex_base.config import get_config
@@ -1279,6 +1280,11 @@ def compile_app(
         ] = hydrate_fallback_definition
         hydrate_fallback_export = hydrate_fallback_definition.export_name
 
+    # Settle memo bodies before snapshotting: a var-returning memo's lazy body
+    # may reference another `@rx.memo` that an AppHarness reload cleared from
+    # `MEMOS`, re-registering it only when read. Doing so after the snapshot
+    # would omit the dependency's own module.
+    materialize_registered_memo_bodies()
     memo_component_files, memo_components_imports = compile_memo_components(
         (
             *MEMOS.values(),

reflex/testing.py

@@ -238,32 +238,43 @@ def _get_source_from_app_source(self, app_source: Any) -> str:
     def _initialize_app(self):
         # disable telemetry reporting for tests
         os.environ["REFLEX_TELEMETRY_ENABLED"] = "false"
-        # Reset the global memo registry so previous AppHarness apps do not
-        # leak compiled component definitions into the next test app.
-        MEMOS.clear()
         self.app_path.mkdir(parents=True, exist_ok=True)
-        if self.app_source is not None:
-            app_globals = self._get_globals_from_signature(self.app_source)
-            if isinstance(self.app_source, functools.partial):
-                self.app_source = self.app_source.func
-            # get the source from a function or module object
-            source_code = "\n".join([
-                "\n".join([
-                    self.get_app_global_source(k, v) for k, v in app_globals.items()
-                ]),
-                self._get_source_from_app_source(self.app_source),
-            ])
-            get_config().loglevel = reflex.constants.LogLevel.INFO
-            with chdir(self.app_path):
-                reflex.reflex._init(
-                    name=self.app_name,
-                    template=reflex.constants.Templates.DEFAULT,
-                )
-                self.app_module_path.write_text(source_code)
-        else:
-            # Just initialize the web folder.
+
+        if self.app_source is None:
+            # Real, pre-existing app: import once and use it as-is. Reloading an
+            # already-imported app re-runs only the top module, dropping every
+            # submodule page/state/@rx.memo, so use reload=False (no reset).
             with chdir(self.app_path):
                 reflex.utils.prerequisites.initialize_frontend_dependencies()
+                os.environ.pop(reflex.constants.PYTEST_CURRENT_TEST, None)
+                os.environ[reflex.constants.APP_HARNESS_FLAG] = "true"
+                get_config(reload=True)
+                self.app_instance, self.app_module = (
+                    reflex.utils.prerequisites.get_and_validate_app(reload=False)
+                )
+                self.app_asgi = self.app_instance()
+            return
+
+        # Generated app: reset the global memo registry so a previous AppHarness
+        # app does not leak compiled component definitions into the next test app.
+        MEMOS.clear()
+        app_globals = self._get_globals_from_signature(self.app_source)
+        if isinstance(self.app_source, functools.partial):
+            self.app_source = self.app_source.func
+        # get the source from a function or module object
+        source_code = "\n".join([
+            "\n".join([
+                self.get_app_global_source(k, v) for k, v in app_globals.items()
+            ]),
+            self._get_source_from_app_source(self.app_source),
+        ])
+        get_config().loglevel = reflex.constants.LogLevel.INFO
+        with chdir(self.app_path):
+            reflex.reflex._init(
+                name=self.app_name,
+                template=reflex.constants.Templates.DEFAULT,
+            )
+            self.app_module_path.write_text(source_code)
         with chdir(self.app_path):
             # Use a new registration context for a new app.
             if AppHarness._base_registration_context is None:
@@ -274,19 +285,12 @@ def _initialize_app(self):
             new_registration_context = deepcopy(AppHarness._base_registration_context)
             self._registry_token = RegistrationContext.set(new_registration_context)
             # ensure config and app are reloaded when testing different app
-            config = get_config(reload=True)
+            get_config(reload=True)
             # Ensure the AppHarness test does not skip State assignment due to running via pytest
             os.environ.pop(reflex.constants.PYTEST_CURRENT_TEST, None)
             os.environ[reflex.constants.APP_HARNESS_FLAG] = "true"
-            # Ensure we compile generated apps, and reload pre-existing app modules
-            # that were already imported so they can re-register memo definitions.
-            should_reload_app = (
-                self.app_source is not None or config.module in sys.modules
-            )
             self.app_instance, self.app_module = (
-                reflex.utils.prerequisites.get_and_validate_app(
-                    reload=should_reload_app
-                )
+                reflex.utils.prerequisites.get_and_validate_app(reload=True)
             )
             self.app_asgi = self.app_instance()
 

tests/units/components/test_memo.py

@@ -22,7 +22,10 @@
     _analyze_params,
     _LazyBody,
     _MemoCallBinding,
+    _reregister_used_memo,
     _strip_optional,
+    create_passthrough_component_memo,
+    materialize_registered_memo_bodies,
 )
 from reflex_base.event import EventChain, EventHandler, no_args_event_spec
 from reflex_base.style import Style
@@ -1812,3 +1815,83 @@ def recursive_count(n: rx.vars.NumberVar[int]) -> rx.Var[int]:
 
     invoked = recursive_count(n=Var(_js_expr="three", _var_type=int))
     assert "recursive_count" in str(invoked)
+
+
+def test_memo_reregisters_on_use_after_registry_clear():
+    """Using a ``@rx.memo`` repopulates ``MEMOS`` after the registry is cleared,
+    so the compiler still emits its module.
+    """
+
+    @rx.memo
+    def sample_card(*, title: rx.Var[str]) -> rx.Component:
+        return rx.box(rx.heading(title))
+
+    key = ("SampleCard", __name__)
+    assert key in MEMOS  # registered at decoration time
+
+    MEMOS.clear()
+    assert key not in MEMOS
+
+    # Using (rendering) the memo must re-register it.
+    sample_card(title="Hello")
+    assert key in MEMOS
+
+    # ...so the compiler still emits its module.
+    sym = memo_paths.mirrored_symbol("SampleCard", __name__)
+    files, _ = compiler.compile_memo_components(tuple(MEMOS.values()))
+    assert any(f"export const {sym} = memo(" in code for _, code in files)
+
+
+def test_nested_var_memo_reregistered_before_compile_snapshot():
+    """A var memo referenced only inside another memo's body is still emitted.
+
+    Regression: rendering ``outer`` after a registry clear leaves its lazy body
+    unread, so ``inner`` is missing from the snapshot the compiler takes.
+    ``materialize_registered_memo_bodies`` re-registers it first.
+    """
+
+    @rx.memo
+    def inner_fmt(*, value: rx.Var[str]) -> rx.Var[str]:
+        return rx.Var.create(f"[{value}]")
+
+    @rx.memo
+    def outer_fmt(*, value: rx.Var[str]) -> rx.Var[str]:
+        return inner_fmt(value=value)
+
+    inner_key = ("inner_fmt", __name__)
+    outer_key = ("outer_fmt", __name__)
+    inner_sym = memo_paths.mirrored_symbol("inner_fmt", __name__)
+
+    MEMOS.clear()
+
+    # Rendering the outer memo re-registers only itself; its lazy body (which
+    # references inner) is not read yet.
+    outer_fmt(value=rx.Var.create("x"))
+    assert outer_key in MEMOS
+    assert inner_key not in MEMOS
+
+    # A snapshot taken now (as the compiler used to) omits inner: compiling it
+    # reads outer's body — which re-registers inner too late for this snapshot —
+    # so inner's module is never emitted even though outer.jsx imports it.
+    stale_snapshot = tuple(MEMOS.values())
+    stale_files, _ = compiler.compile_memo_components(stale_snapshot)
+    assert not any(f"export const {inner_sym} =" in code for _, code in stale_files)
+
+    # The compiler pre-pass settles the registry before it snapshots MEMOS.
+    materialize_registered_memo_bodies()
+    assert inner_key in MEMOS
+
+    files, _ = compiler.compile_memo_components(tuple(MEMOS.values()))
+    assert any(f"export const {inner_sym} =" in code for _, code in files)
+
+
+def test_reregister_used_memo_skips_passthrough_auto_memos():
+    """Passthrough auto-memos (tracked separately by the compiler) must never
+    enter ``MEMOS`` on use.
+    """
+    _, definition = create_passthrough_component_memo(
+        rx.box(rx.text("x")), source_module=__name__
+    )
+    MEMOS.clear()
+    _reregister_used_memo(definition)
+    assert MEMOS == {}

tests/units/test_testing.py

@@ -111,10 +111,11 @@ def test_app_harness_initialize_clears_memo_registries(
     harness_mocks.get_and_validate_app.assert_called_once_with(reload=True)
 
 
-def test_app_harness_initialize_reloads_existing_imported_app(
+def test_app_harness_initialize_preserves_existing_imported_app(
     tmp_path, preserve_memo_registries, harness_mocks, monkeypatch
 ):
-    """Ensure pre-existing imported apps are reloaded after memo registry reset.
+    """A real (``app_source is None``) app is imported with ``reload=False`` and
+    its registrations (incl. memos) are kept, not reset.
 
     Args:
         tmp_path: pytest tmp_path fixture
@@ -132,8 +133,11 @@ def test_app_harness_initialize_reloads_existing_imported_app(
         harness_mocks.config.module,
         ModuleType(harness_mocks.config.module),
     )
+    MEMOS["existing_memo", None] = mock.sentinel.memo
 
     harness = AppHarness.create(root=tmp_path / "plain_app")
     harness._initialize_app()
 
-    harness_mocks.get_and_validate_app.assert_called_once_with(reload=True)
+    harness_mocks.get_and_validate_app.assert_called_once_with(reload=False)
+    # The real app's registrations (incl. memos) are preserved, not reset.
+    assert ("existing_memo", None) in MEMOS