From 241940efec641dac6e125af696d4d3372c05e541 Mon Sep 17 00:00:00 2001
From: Khaleel Al-Adhami <khaleel.aladhami@gmail.com>
Date: Fri, 25 Oct 2024 17:29:22 -0700
Subject: [PATCH] fix order of operations and add note

---
 reflex/app.py                  | 22 ++++++++++++----------
 reflex/components/component.py |  8 --------
 2 files changed, 12 insertions(+), 18 deletions(-)

diff --git a/reflex/app.py b/reflex/app.py
index 557c95358..5923e3389 100644
--- a/reflex/app.py
+++ b/reflex/app.py
@@ -900,6 +900,18 @@ class App(MiddlewareMixin, LifespanMixin, Base):
         all_imports = {}
         custom_components = set()
 
+        # This has to happen before compiling stateful components as that
+        # prevents recursive functions from reaching all components.
+        for component in self.pages.values():
+            # Add component._get_all_imports() to all_imports.
+            all_imports.update(component._get_all_imports())
+
+            # Add the app wrappers from this component.
+            app_wrappers.update(component._get_all_app_wrap_components())
+
+            # Add the custom components from the page to the set.
+            custom_components |= component._get_all_custom_components()
+
         # Perform auto-memoization of stateful components.
         (
             stateful_components_path,
@@ -985,16 +997,6 @@ class App(MiddlewareMixin, LifespanMixin, Base):
                 compile_results.append(future.result())
                 progress.advance(task)
 
-        for component in self.pages.values():
-            # Add component._get_all_imports() to all_imports.
-            all_imports.update(component._get_all_imports())
-
-            # Add the app wrappers from this component.
-            app_wrappers.update(component._get_all_app_wrap_components())
-
-            # Add the custom components from the page to the set.
-            custom_components |= component._get_all_custom_components()
-
         app_root = self._app_root(app_wrappers=app_wrappers)
 
         # Get imports from AppWrap components.
diff --git a/reflex/components/component.py b/reflex/components/component.py
index 504a01909..9fea2f05b 100644
--- a/reflex/components/component.py
+++ b/reflex/components/component.py
@@ -1968,14 +1968,6 @@ class StatefulComponent(BaseComponent):
         """
         from reflex.components.core.foreach import Foreach
 
-        if isinstance(component, CustomComponent):
-            # Do not memoize custom components.
-            return None
-
-        if component._get_app_wrap_components():
-            # Do not memoize app wrap components.
-            return None
-
         if component._memoization_mode.disposition == MemoizationDisposition.NEVER:
             # Never memoize this component.
             return None