reregister memo in app harness

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

@@ -0,0 +1 @@
+`AppHarness` no longer reloads an existing app (one created without `app_source`), which dropped the pages, states and `@rx.memo` components defined in the app's submodules.

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

@@ -0,0 +1 @@
+Fix `@rx.memo` definitions from a module shared by several `AppHarness` apps in one test process going missing from later apps, which left their imports unresolved at runtime.

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

@@ -244,6 +244,10 @@ def is_ready(self) -> bool:
         """
         return self._ready
 
+    def reset(self) -> None:
+        """Drop the cached body so the next :meth:`get` runs the thunk again."""
+        self._ready = False
+
     def get(self, thunk: Callable[[], _BodyT] | None = None) -> _BodyT:
         """Return the body, running and caching ``thunk`` on first read.
 
@@ -653,6 +657,48 @@ def _register_memo_definition(definition: MemoDefinition) -> None:
     MEMOS[key] = definition
 
 
+def _reregister_used_memo(definition: MemoDefinition) -> None:
+    """Re-register a used ``@rx.memo`` that a cleared ``MEMOS`` dropped.
+
+    Its cached body is reset, so the next read re-runs it and re-registers the
+    nested memos it uses too. Passthrough auto-memos (defined in this module)
+    are tracked separately and stay out of ``MEMOS``.
+
+    Args:
+        definition: The memo definition being used.
+    """
+    if definition.fn.__module__ == __name__:
+        return
+    key = _memo_registry_key(definition)
+    if key in MEMOS:
+        return
+    MEMOS[key] = definition
+    if isinstance(definition, MemoComponentDefinition):
+        definition._component.reset()
+    elif isinstance(definition, MemoFunctionDefinition):
+        definition._function.reset()
+
+
+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.
 
@@ -1713,6 +1759,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)
         )
@@ -1727,6 +1774,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)
         )
@@ -1737,6 +1785,7 @@ def _as_var(self) -> FunctionVar:
         Returns:
             The imported function var.
         """
+        _reregister_used_memo(self._definition)
         return self._imported_var
 
 
@@ -1775,6 +1824,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.
@@ -1868,6 +1918,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
         )
@@ -2452,5 +2503,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": "f170ac685b6ba5892370166c80684db3",
   "reflex/__init__.pyi": "1c841ad90d789507cd56d966b474d11c",
   "reflex/components/__init__.pyi": "9facd05a776d0641432696bbf8e34388",
-  "reflex/experimental/memo.pyi": "87e485038dc79f66b2ab15144845ef1e"
+  "reflex/experimental/memo.pyi": "39df3eded29349502d91bac39ad223fd"
 }

reflex/compiler/compiler.py

@@ -25,6 +25,7 @@
     MemoDefinition,
     MemoFunctionDefinition,
     create_component_memo,
+    materialize_registered_memo_bodies,
     reset_memo_component_classes,
 )
 from reflex_base.config import get_config
@@ -1401,6 +1402,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

@@ -267,32 +267,44 @@ def _initialize_app(self):
         # prerendering, so the dev server serves prerendered page HTML whose
         # hydration failures break event delivery (notably under vite >= 8.2).
         environment.REFLEX_ENV_MODE.set(reflex.constants.Env.DEV)
-        # 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):
+                # Load this app's config before .web is initialized from it.
+                reload_config()
                 reflex.utils.prerequisites.initialize_frontend_dependencies()
+                os.environ.pop(reflex.constants.PYTEST_CURRENT_TEST, None)
+                os.environ[reflex.constants.APP_HARNESS_FLAG] = "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:
@@ -303,19 +315,12 @@ def _initialize_app(self):
             new_registration_context = AppHarness._base_registration_context.fork()
             self._registry_token = RegistrationContext.set(new_registration_context)
             # ensure config and app are reloaded when testing different app
-            config = reload_config()
+            reload_config()
             # 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

@@ -24,7 +24,10 @@
     _analyze_params,
     _LazyBody,
     _MemoCallBinding,
+    _reregister_used_memo,
     component_hash,
+    create_passthrough_component_memo,
+    materialize_registered_memo_bodies,
     memo_tag,
 )
 from reflex_base.event import EventChain, EventHandler, no_args_event_spec
@@ -2010,6 +2013,119 @@ def recursive_count(n: rx.vars.NumberVar[int]) -> rx.Var[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"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
+
+    # 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_cached_memo_bodies_reregister_nested_memos_after_registry_clear():
+    """Memos nested in a body an earlier compile cached are emitted again.
+
+    Regression: a later ``AppHarness`` clears ``MEMOS`` without re-importing a
+    shared memo module, so only re-running the cached outer bodies can
+    re-register the inner memos their compiled modules import.
+    """
+
+    @rx.memo
+    def nested_fmt(*, value: rx.Var[str]) -> rx.Var[str]:
+        return rx.Var.create(f"[{value}]")
+
+    @rx.memo
+    def wrapping_fmt(*, value: rx.Var[str]) -> rx.Var[str]:
+        return nested_fmt(value=value)
+
+    @rx.memo
+    def nested_card(*, title: rx.Var[str]) -> rx.Component:
+        return rx.text(title)
+
+    @rx.memo
+    def wrapping_card(*, title: rx.Var[str]) -> rx.Component:
+        return rx.box(nested_card(title=title))
+
+    # The first app's compile caches every body.
+    materialize_registered_memo_bodies()
+
+    MEMOS.clear()
+    wrapping_fmt(value=rx.Var.create("x"))
+    wrapping_card(title="Hello")
+    materialize_registered_memo_bodies()
+
+    assert ("nested_fmt", __name__) in MEMOS
+    assert ("NestedCard", __name__) in MEMOS
+    files, _ = compiler.compile_memo_components(tuple(MEMOS.values()))
+    for name in ("nested_fmt", "NestedCard"):
+        sym = memo_paths.mirrored_symbol(name, __name__)
+        assert any(f"export const {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 == {}
+
+
 class _AppWrapProbe(Component):
     """A component whose only per-instance artifact is an app-wrap component."""