[ENG-4005] Proxy backend requests on '/' to the frontend

pyproject.toml

@@ -57,6 +57,7 @@ setuptools = ">=75.0"
 httpx = ">=0.25.1,<1.0"
 twine = ">=4.0.0,<7.0"
 tomlkit = ">=0.12.4,<1.0"
+asgiproxy = { version = "==0.1.1", optional = true }
 lazy_loader = ">=0.4"
 reflex-chakra = ">=0.6.0"
 typing_extensions = ">=4.6.0"
@@ -80,10 +81,14 @@ selenium = ">=4.11.0,<5.0"
 pytest-benchmark = ">=4.0.0,<6.0"
 playwright = ">=1.46.0"
 pytest-playwright = ">=0.5.1"
+asgiproxy = "==0.1.1"
 
 [tool.poetry.scripts]
 reflex = "reflex.reflex:cli"
 
+[tool.poetry.extras]
+proxy = ["asgiproxy"]
+
 [build-system]
 requires = ["poetry-core>=1.5.1"]
 build-backend = "poetry.core.masonry.api"

reflex/app.py

@@ -331,6 +331,12 @@ def __post_init__(self):
 
             self.register_lifespan_task(windows_hot_reload_lifespan_hack)
 
+        # Enable proxying to frontend server.
+        if not environment.REFLEX_BACKEND_ONLY.get():
+            from reflex.proxy import proxy_middleware
+
+            self.register_lifespan_task(proxy_middleware)
+
     def _enable_state(self) -> None:
         """Enable state for the app."""
         if not self.state:

reflex/proxy.py

@@ -0,0 +1,76 @@
+"""Handle proxying frontend requests from the backend server."""
+
+from __future__ import annotations
+
+from contextlib import asynccontextmanager
+from typing import AsyncGenerator
+from urllib.parse import urlparse
+
+from fastapi import FastAPI
+from starlette.types import ASGIApp
+
+from .config import get_config
+from .utils import console
+
+try:
+    from asgiproxy.config import BaseURLProxyConfigMixin, ProxyConfig
+    from asgiproxy.context import ProxyContext
+    from asgiproxy.simple_proxy import make_simple_proxy_app
+except ImportError:
+
+    @asynccontextmanager
+    async def proxy_middleware(*args, **kwargs) -> AsyncGenerator[None, None]:
+        """A no-op proxy middleware for when asgiproxy is not installed.
+
+        Args:
+            *args: The positional arguments.
+            **kwargs: The keyword arguments.
+
+        Yields:
+            None
+        """
+        yield
+else:
+
+    def _get_proxy_app_with_context(frontend_host: str) -> tuple[ProxyContext, ASGIApp]:
+        """Get the proxy app with the given frontend host.
+
+        Args:
+            frontend_host: The frontend host to proxy requests to.
+
+        Returns:
+            The proxy context and app.
+        """
+
+        class LocalProxyConfig(BaseURLProxyConfigMixin, ProxyConfig):
+            upstream_base_url = frontend_host
+            rewrite_host_header = urlparse(upstream_base_url).netloc
+
+        proxy_context = ProxyContext(LocalProxyConfig())
+        proxy_app = make_simple_proxy_app(proxy_context)
+        return proxy_context, proxy_app
+
+    @asynccontextmanager
+    async def proxy_middleware(  # pyright: ignore[reportGeneralTypeIssues]
+        app: FastAPI,
+    ) -> AsyncGenerator[None, None]:
+        """A middleware to proxy requests to the separate frontend server.
+
+        The proxy is installed on the / endpoint of the FastAPI instance.
+
+        Args:
+            app: The FastAPI instance.
+
+        Yields:
+            None
+        """
+        config = get_config()
+        backend_port = config.backend_port
+        frontend_host = f"http://localhost:{config.frontend_port}"
+        proxy_context, proxy_app = _get_proxy_app_with_context(frontend_host)
+        app.mount("/", proxy_app)
+        console.debug(
+            f"Proxying '/' requests on port {backend_port} to {frontend_host}"
+        )
+        async with proxy_context:
+            yield

reflex/testing.py

@@ -44,6 +44,7 @@
 import reflex.utils.prerequisites
 import reflex.utils.processes
 from reflex.config import environment
+from reflex.proxy import proxy_middleware
 from reflex.state import (
     BaseState,
     StateManager,
@@ -297,6 +298,9 @@ def _initialize_app(self):
             self.state_manager = StateManagerRedis.create(self.app_instance.state)
         else:
             self.state_manager = self.app_instance._state_manager
+        # Disable proxy for app harness tests.
+        if proxy_middleware in self.app_instance.lifespan_tasks:
+            self.app_instance.lifespan_tasks.remove(proxy_middleware)
 
     def _reload_state_module(self):
         """Reload the rx.State module to avoid conflict when reloading."""
@@ -364,9 +368,12 @@ async def _reset_backend_state_manager(self):
     def _start_frontend(self):
         # Set up the frontend.
         with chdir(self.app_path):
+            backend_host, backend_port = self._poll_for_servers().getsockname()
             config = reflex.config.get_config()
+            config.backend_port = backend_port
             config.api_url = "http://{0}:{1}".format(
-                *self._poll_for_servers().getsockname(),
+                backend_host,
+                backend_port,
             )
             reflex.utils.build.setup_frontend(self.app_path)
 
@@ -391,6 +398,7 @@ def _wait_frontend(self):
                 self.frontend_url = m.group(1)
                 config = reflex.config.get_config()
                 config.deploy_url = self.frontend_url
+                config.frontend_port = int(self.frontend_url.rpartition(":")[2])
                 break
         if self.frontend_url is None:
             raise RuntimeError("Frontend did not start")
@@ -915,17 +923,20 @@ def _run_frontend(self):
             root=web_root,
             error_page_map=error_page_map,
         ) as self.frontend_server:
-            self.frontend_url = "http://localhost:{1}".format(
-                *self.frontend_server.socket.getsockname()
-            )
+            config = reflex.config.get_config()
+            config.frontend_port = self.frontend_server.server_address[1]
+            self.frontend_url = f"http://localhost:{config.frontend_port}"
             self.frontend_server.serve_forever()
 
     def _start_frontend(self):
         # Set up the frontend.
         with chdir(self.app_path):
+            backend_host, backend_port = self._poll_for_servers().getsockname()
             config = reflex.config.get_config()
+            config.backend_port = backend_port
             config.api_url = "http://{0}:{1}".format(
-                *self._poll_for_servers().getsockname(),
+                backend_host,
+                backend_port,
             )
             reflex.reflex.export(
                 zipping=False,

scripts/wait_for_listening_port.py

@@ -10,6 +10,8 @@
 from concurrent.futures import ThreadPoolExecutor, as_completed
 from typing import Tuple
 
+import httpx
+
 # psutil is already a dependency of Reflex itself - so it's OK to use
 import psutil
 
@@ -61,6 +63,12 @@ def main():
         else:
             print(f"FAIL: {msg}")
             exit(1)
+    # Make sure the HTTP response for both ports is the same (proxy frontend to backend).
+    responses = [(port, httpx.get(f"http://localhost:{port}")) for port in args.port]
+    n_port, n_resp = responses[0]
+    for port, resp in responses[1:]:
+        assert resp.content == n_resp.content, f"HTTP response on {port} is not equal."
+        print(f"OK: HTTP responses for :{n_port}/ and :{port}/ are equal.")
 
 
 if __name__ == "__main__":