set env mode before importing modules that contain rx.State subclasses

reflex/reflex.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+# WARNING: do not import any modules that contain rx.State subclasses here
 import atexit
 import os
 from pathlib import Path
@@ -15,7 +16,6 @@ from reflex_cli.utils import dependency
 from reflex import constants
 from reflex.config import EnvironmentVariables, get_config
 from reflex.custom_components.custom_components import custom_components_cli
-from reflex.state import reset_disk_state_manager
 from reflex.utils import console, redir, telemetry
 
 # Disable typer+rich integration for help panels
@@ -154,14 +154,17 @@ def _run(
     loglevel: constants.LogLevel = config.loglevel,
 ):
     """Run the app in the given directory."""
+
+    # Set env mode in the environment
+    # This must be set before importing modules that contain rx.State subclasses
+    EnvironmentVariables.REFLEX_ENV_MODE.set(env)
+
+    from reflex.state import reset_disk_state_manager
     from reflex.utils import build, exec, prerequisites, processes
 
     # Set the log level.
     console.set_log_level(loglevel)
 
-    # Set env mode in the environment
-    EnvironmentVariables.REFLEX_ENV_MODE.set(env)
-
     # Show system info
     exec.output_system_info()
 
@@ -304,7 +307,6 @@ def run(
         raise typer.Exit(1)
     EnvironmentVariables.REFLEX_BACKEND_ONLY.set(backend)
     EnvironmentVariables.REFLEX_FRONTEND_ONLY.set(frontend)
-    EnvironmentVariables.REFLEX_ENV_MODE.set(env)
 
     _run(env, frontend, backend, frontend_port, backend_port, backend_host, loglevel)