add a config variable to add extra overlay components (#4763)

* add a config variable to add extra overlay components

* add integration test

* Apply suggestions from code review

---------

Co-authored-by: Masen Furer <[email protected]>

reflex/app.py

@@ -164,9 +164,26 @@ def default_overlay_component() -> Component:
     """
     config = get_config()
 
+    extra_config = config.extra_overlay_function
+    config_overlay = None
+    if extra_config:
+        module, _, function_name = extra_config.rpartition(".")
+        try:
+            module = __import__(module)
+            config_overlay = getattr(module, function_name)()
+        except Exception as e:
+            from reflex.compiler.utils import save_error
+
+            log_path = save_error(e)
+
+            console.error(
+                f"Error loading extra_overlay_function {extra_config}. Error saved to {log_path}"
+            )
+
     return Fragment.create(
         connection_pulser(),
         connection_toaster(),
+        *([config_overlay] if config_overlay else []),
         *([backend_disabled()] if config.is_reflex_cloud else []),
         *codespaces.codespaces_auto_redirect(),
     )

reflex/compiler/utils.py

@@ -158,6 +158,22 @@ def get_import_dict(lib: str, default: str = "", rest: list[str] | None = None)
     }
 
 
+def save_error(error: Exception) -> str:
+    """Save the error to a file.
+
+    Args:
+        error: The error to save.
+
+    Returns:
+        The path of the saved error.
+    """
+    timestamp = datetime.now().strftime("%Y-%m-%d__%H-%M-%S")
+    constants.Reflex.LOGS_DIR.mkdir(parents=True, exist_ok=True)
+    log_path = constants.Reflex.LOGS_DIR / f"error_{timestamp}.log"
+    traceback.TracebackException.from_exception(error).print(file=log_path.open("w+"))
+    return str(log_path)
+
+
 def compile_state(state: Type[BaseState]) -> dict:
     """Compile the state of the app.
 
@@ -170,10 +186,7 @@ def compile_state(state: Type[BaseState]) -> dict:
     try:
         initial_state = state(_reflex_internal_init=True).dict(initial=True)
     except Exception as e:
-        timestamp = datetime.now().strftime("%Y-%m-%d__%H-%M-%S")
-        constants.Reflex.LOGS_DIR.mkdir(parents=True, exist_ok=True)
-        log_path = constants.Reflex.LOGS_DIR / f"state_compile_error_{timestamp}.log"
-        traceback.TracebackException.from_exception(e).print(file=log_path.open("w+"))
+        log_path = save_error(e)
         console.warn(
             f"Failed to compile initial state with computed vars. Error log saved to {log_path}"
         )

reflex/config.py

@@ -709,6 +709,9 @@ class Config:  # pyright: ignore [reportIncompatibleVariableOverride]
     # Whether the app is running in the reflex cloud environment.
     is_reflex_cloud: bool = False
 
+    # Extra overlay function to run after the app is built. Formatted such that `from path_0.path_1... import path[-1]`, and calling it with no arguments would work. For example, "reflex.components.moment.momnet".
+    extra_overlay_function: Optional[str] = None
+
     def __init__(self, *args, **kwargs):
         """Initialize the config values.
 

tests/integration/test_extra_overlay_function.py

@@ -0,0 +1,87 @@
+"""Test case for adding an overlay component defined in the rxconfig."""
+
+from typing import Generator
+
+import pytest
+from selenium.webdriver.common.by import By
+
+from reflex.testing import AppHarness, WebDriver
+
+
+def ExtraOverlay():
+    import reflex as rx
+
+    rx.config.get_config().extra_overlay_function = "reflex.components.moment.moment"
+
+    def index():
+        return rx.vstack(
+            rx.el.input(
+                id="token",
+                value=rx.State.router.session.client_token,
+                is_read_only=True,
+            ),
+            rx.text(
+                "Hello World",
+            ),
+        )
+
+    app = rx.App(_state=rx.State)
+    app.add_page(index)
+
+
+@pytest.fixture(scope="module")
+def extra_overlay(tmp_path_factory) -> Generator[AppHarness, None, None]:
+    """Start ExtraOverlay app at tmp_path via AppHarness.
+
+    Args:
+        tmp_path_factory: pytest tmp_path_factory fixture
+
+    Yields:
+        running AppHarness instance
+    """
+    with AppHarness.create(
+        root=tmp_path_factory.mktemp("extra_overlay"),
+        app_source=ExtraOverlay,
+    ) as harness:
+        assert harness.app_instance is not None, "app is not running"
+        yield harness
+
+
+@pytest.fixture
+def driver(extra_overlay: AppHarness):
+    """Get an instance of the browser open to the extra overlay app.
+
+    Args:
+        extra_overlay: harness for the ExtraOverlay app.
+
+    Yields:
+        WebDriver instance.
+    """
+    driver = extra_overlay.frontend()
+    try:
+        token_input = driver.find_element(By.ID, "token")
+        assert token_input
+        # wait for the backend connection to send the token
+        token = extra_overlay.poll_for_value(token_input)
+        assert token is not None
+
+        yield driver
+    finally:
+        driver.quit()
+
+
+def test_extra_overlay(driver: WebDriver, extra_overlay: AppHarness):
+    """Test the ExtraOverlay app.
+
+    Args:
+        driver: WebDriver instance.
+        extra_overlay: harness for the ExtraOverlay app.
+    """
+    # Check that the text is displayed.
+    text = driver.find_element(By.XPATH, "//*[contains(text(), 'Hello World')]")
+    assert text
+    assert text.text == "Hello World"
+
+    time = driver.find_element(By.TAG_NAME, "time")
+    assert time
+    assert time.text