[ENG-3717] [flexgen] Initialize app from refactored code (#3918)

reflex/constants/base.py

@@ -117,7 +117,9 @@ class Templates(SimpleNamespace):
     REFLEX_BUILD_POLL_URL = REFLEX_BUILD_BACKEND + "/api/init/{reflex_init_token}"
 
     # The URL to fetch the generation's reflex code
-    REFLEX_BUILD_CODE_URL = REFLEX_BUILD_BACKEND + "/api/gen/{generation_hash}"
+    REFLEX_BUILD_CODE_URL = (
+        REFLEX_BUILD_BACKEND + "/api/gen/{generation_hash}/refactored"
+    )
 
     class Dirs(SimpleNamespace):
         """Folders used by the template system of Reflex."""

reflex/utils/exceptions.py

@@ -91,3 +91,7 @@ class EventFnArgMismatch(ReflexError, TypeError):
 
 class DynamicRouteArgShadowsStateVar(ReflexError, NameError):
     """Raised when a dynamic route arg shadows a state var."""
+
+
+class GeneratedCodeHasNoFunctionDefs(ReflexError):
+    """Raised when refactored code generated with flexgen has no functions defined."""

reflex/utils/prerequisites.py

@@ -15,7 +15,7 @@
 import stat
 import sys
 import tempfile
-import textwrap
+import time
 import zipfile
 from datetime import datetime
 from fileinput import FileInput
@@ -36,6 +36,7 @@
 from reflex.compiler import templates
 from reflex.config import Config, get_config
 from reflex.utils import console, net, path_ops, processes
+from reflex.utils.exceptions import GeneratedCodeHasNoFunctionDefs
 from reflex.utils.format import format_library_name
 from reflex.utils.registry import _get_best_registry
 
@@ -1435,19 +1436,37 @@ def initialize_main_module_index_from_generation(app_name: str, generation_hash:
     Args:
         app_name: The name of the app.
         generation_hash: The generation hash from reflex.build.
+
+    Raises:
+        GeneratedCodeHasNoFunctionDefs: If the fetched code has no function definitions
+            (the refactored reflex code is expected to have at least one root function defined).
     """
     # Download the reflex code for the generation.
-    resp = net.get(
-        constants.Templates.REFLEX_BUILD_CODE_URL.format(
-            generation_hash=generation_hash
+    url = constants.Templates.REFLEX_BUILD_CODE_URL.format(
+        generation_hash=generation_hash
+    )
+    resp = net.get(url)
+    while resp.status_code == httpx.codes.SERVICE_UNAVAILABLE:
+        console.debug("Waiting for the code to be generated...")
+        time.sleep(1)
+        resp = net.get(url)
+    resp.raise_for_status()
+
+    # Determine the name of the last function, which renders the generated code.
+    defined_funcs = re.findall(r"def ([a-zA-Z_]+)\(", resp.text)
+    if not defined_funcs:
+        raise GeneratedCodeHasNoFunctionDefs(
+            f"No function definitions found in generated code from {url!r}."
         )
-    ).raise_for_status()
+    render_func_name = defined_funcs[-1]
 
     def replace_content(_match):
         return "\n".join(
             [
-                "def index() -> rx.Component:",
-                textwrap.indent("return " + resp.text, "    "),
+                resp.text,
+                "",
+                "" "def index() -> rx.Component:",
+                f"    return {render_func_name}()",
                 "",
                 "",
             ],