Send ASGI lifespan events to api_transformer apps; drop the extra top-level Starlette app

docs/api-routes/overview.md

@@ -128,9 +128,10 @@ async def get_items():
 def add_logging_middleware(app):
     # This is a simple example middleware that logs requests
     async def middleware(scope, receive, send):
-        # Log the request path
-        path = scope["path"]
-        print("Request:", path)
+        # Log the request path; lifespan scopes have none
+        path = scope.get("path")
+        if path is not None:
+            print("Request:", path)
         await app(scope, receive, send)
 
     return middleware
@@ -140,6 +141,50 @@ def add_logging_middleware(app):
 app = rx.App(api_transformer=[fastapi_app, add_logging_middleware])
 ```
 
+## Lifespan
+
+Reflex sends [ASGI lifespan](https://asgi.readthedocs.io/en/latest/specs/lifespan.html)
+events to the app that `api_transformer` produces, as an ASGI server would, once
+the app's [lifespan tasks](/docs/utility-methods/lifespan-tasks/) have started.
+Middleware returned by a callable transformer should pass lifespan events on to
+the app it wraps. Starlette and FastAPI do not pass them to the apps they mount,
+so Reflex also sends them to the app each Starlette or FastAPI transformer
+mounts. Lifespans start from the innermost transformer out and stop in reverse,
+before the lifespan tasks stop. Any state they yield is available on
+`request.state`.
+
+```python
+from contextlib import asynccontextmanager
+
+import reflex as rx
+from fastapi import FastAPI, Request
+
+
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    app.state.greeting = "hello"
+    yield {"client": "connected"}
+
+
+fastapi_app = FastAPI(lifespan=lifespan)
+
+
+@fastapi_app.get("/api/greeting")
+async def greeting(request: Request):
+    return dict(
+        greeting=request.app.state.greeting,
+        client=request.state.client,
+    )
+
+
+app = rx.App(api_transformer=fastapi_app)
+```
+
+Apps you mount on a transformer yourself do not get lifespan events; register
+their startup and shutdown code with `app.register_lifespan_task` instead. An app
+that does not support lifespan events is skipped, and Reflex logs that it
+appears not to support them.
+
 ## Reserved Routes
 
 Some routes on the backend are reserved for the runtime of Reflex, and should not be overridden unless you know what you are doing.

docs/utility_methods/lifespan_tasks.md

@@ -62,7 +62,8 @@ If the task accepts the special argument, `app`, it will be passed the Reflex ap
 instance (`rx.App`/`LifespanMixin`).
 
 If the task accepts the special argument, `starlette_app`, it will be passed the
-underlying `Starlette` application instance.
+`Starlette` application that serves Reflex's own routes (the one a Starlette or
+FastAPI `api_transformer` mounts and a callable `api_transformer` wraps).
 
 ```python
 app = rx.App()

news/+api-transformer-lifespan.bugfix.md

@@ -0,0 +1 @@
+Apps passed to `api_transformer` now get ASGI lifespan events, so the `lifespan` (and `on_startup`/`on_shutdown` handlers) of a FastAPI or Starlette transformer run, and state the lifespan yields is available on `request.state`. Events are forwarded through callable transformers as the ASGI spec describes; apps you mount on a transformer yourself still get none. If you worked around this by also registering the transformer's lifespan with `register_lifespan_task`, or by finding the transformer through the `starlette_app` lifespan task argument, remove the workaround: the first now runs the lifespan twice, and `starlette_app` is now the `Starlette` app serving Reflex's own routes.

news/+lighter-asgi-entrypoint.performance.md

@@ -0,0 +1 @@
+The backend no longer wraps the app in an extra `Starlette` application to run lifespan tasks, which removes a routing layer from every request. As a result, calling the `rx.App` instance returns a plain ASGI app rather than a `Starlette` instance.

reflex/app.py

@@ -347,27 +347,66 @@ def _route_arg_label(arg_type: str) -> str:
     return "list" if arg_type == constants.RouteArgType.LIST else "single"
 
 
-class _ContextMiddleware:
-    """Ensure Reflex contexts are attached for each ASGI request.
+class _RootMiddleware:
+    """The outermost ASGI layer: attach Reflex contexts and serve the lifespan.
 
     Many ASGI servers start each request with a fresh contextvars scope, so this
     middleware re-applies the RegistrationContext and EventContext that are
-    needed for Reflex state and event processing.
+    needed for Reflex state and event processing. It answers the server's
+    lifespan events itself, entering the given lifespan context in between.
     """
 
-    def __init__(self, app: ASGIApp, reflex_app: App):
+    def __init__(
+        self,
+        app: ASGIApp,
+        reflex_app: App,
+        lifespan: Callable[[Scope], contextlib.AbstractAsyncContextManager[Any]],
+    ):
         """Wrap an ASGI app so that it runs with the Reflex contexts set.
 
         Args:
             app: The next ASGI app in the middleware stack.
             reflex_app: The Reflex app owning the contexts to attach.
+            lifespan: Returns the context to enter for the server's lifespan scope.
         """
         self.app = app
         self.reflex_app = reflex_app
+        self.lifespan = lifespan
 
     async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
         self.reflex_app._set_contexts_internal()
-        await self.app(scope, receive, send)
+        if scope["type"] == "lifespan":
+            await self._serve_lifespan(scope, receive, send)
+        else:
+            await self.app(scope, receive, send)
+
+    async def _serve_lifespan(self, scope: Scope, receive: Receive, send: Send):
+        """Answer the server's lifespan events around the lifespan context.
+
+        Args:
+            scope: The lifespan scope.
+            receive: Receives the server's lifespan events.
+            send: Reports the outcome of startup and shutdown to the server.
+
+        Raises:
+            BaseException: Whatever made startup or shutdown fail, once reported.
+        """
+        started = False
+        await receive()
+        try:
+            async with self.lifespan(scope):
+                await send({"type": "lifespan.startup.complete"})
+                started = True
+                await receive()
+        except BaseException:
+            await send({
+                "type": "lifespan.shutdown.failed"
+                if started
+                else "lifespan.startup.failed",
+                "message": traceback.format_exc(),
+            })
+            raise
+        await send({"type": "lifespan.shutdown.complete"})
 
 
 class _DrainOversizedBodyMiddleware:
@@ -888,13 +927,16 @@ def __call__(self) -> ASGIApp:
             msg = "The app has not been initialized."
             raise ValueError(msg)
 
-        asgi_app = self._api
+        api = asgi_app = self._api
 
         if environment.REFLEX_MOUNT_FRONTEND_COMPILED_APP.get():
             from reflex.utils.exec import get_frontend_mount
 
             asgi_app.routes.append(get_frontend_mount(router=self.router))
 
+        # Lifespan events go to the outermost app, which may forward them, and to
+        # each app a Starlette transformer mounts, as Starlette does not forward them.
+        lifespan_apps: list[ASGIApp] = []
         if self.api_transformer is not None:
             api_transformers: Sequence[Starlette | Callable[[ASGIApp], ASGIApp]] = (
                 [self.api_transformer]
@@ -904,6 +946,15 @@ def __call__(self) -> ASGIApp:
 
             for api_transformer in api_transformers:
                 if isinstance(api_transformer, Starlette):
+                    if api_transformer.router.lifespan_context in self._lifespan_tasks:
+                        logger.warning(
+                            f"The lifespan of api_transformer {api_transformer!r} is "
+                            "also registered with register_lifespan_task, so it runs "
+                            "twice: Reflex now runs a transformer's own lifespan. "
+                            "Remove the register_lifespan_task call."
+                        )
+                    if asgi_app is not api:
+                        lifespan_apps.append(asgi_app)
                     # Mount the api to the starlette app.
                     App._add_cors(api_transformer)
                     api_transformer.mount("", asgi_app)
@@ -912,11 +963,16 @@ def __call__(self) -> ASGIApp:
                     # Transform the asgi app.
                     asgi_app = api_transformer(asgi_app)
 
-        top_asgi_app = Starlette(lifespan=self._run_lifespan_tasks)
-        # Make sure Reflex contexts are attached for each request.
-        top_asgi_app.add_middleware(_ContextMiddleware, reflex_app=self)
-        top_asgi_app.mount("", asgi_app)
-        App._add_cors(top_asgi_app)
+            if asgi_app is not api:
+                lifespan_apps.append(asgi_app)
+
+        top_asgi_app = _RootMiddleware(
+            # Covers the routes a callable transformer adds and the 500 responses
+            # that the outermost Starlette app's error handler sends.
+            cors.CORSMiddleware(asgi_app, **App._cors_options()),
+            reflex_app=self,
+            lifespan=functools.partial(self._run_lifespan, api, lifespan_apps),
+        )
         if otel.asgi_middleware is not None:
             return otel.asgi_middleware(top_asgi_app)
         return top_asgi_app
@@ -975,20 +1031,28 @@ def _add_optional_endpoints(self):
         if environment.REFLEX_ADD_ALL_ROUTES_ENDPOINT.get():
             self.add_all_routes_endpoint()
 
+    @staticmethod
+    def _cors_options() -> dict[str, Any]:
+        """Get the options of the CORS middleware.
+
+        Returns:
+            The keyword arguments for ``CORSMiddleware``.
+        """
+        return {
+            "allow_credentials": True,
+            "allow_methods": ["*"],
+            "allow_headers": ["*"],
+            "allow_origins": get_config().cors_allowed_origins,
+        }
+
     @staticmethod
     def _add_cors(api: Starlette):
         """Add CORS middleware to the app.
 
         Args:
             api: The Starlette app to add CORS middleware to.
         """
-        api.add_middleware(
-            cors.CORSMiddleware,
-            allow_credentials=True,
-            allow_methods=["*"],
-            allow_headers=["*"],
-            allow_origins=get_config().cors_allowed_origins,
-        )
+        api.add_middleware(cors.CORSMiddleware, **App._cors_options())
 
     @property
     def state_manager(self) -> StateManager:

reflex/app_mixins/lifespan.py

@@ -9,11 +9,12 @@
 import inspect
 import logging
 import time
-from collections.abc import Callable, Coroutine
+from collections.abc import Callable, Coroutine, Sequence
 from typing import TYPE_CHECKING, TypeVar, overload
 
 from reflex_base.utils import console
 from reflex_base.utils.exceptions import InvalidLifespanTaskTypeError
+from reflex_base.utils.types import ASGIApp, Message, Scope
 from starlette.applications import Starlette
 
 from reflex.utils import prerequisites
@@ -66,6 +67,83 @@ def _log_cleanup_error(task: asyncio.Task) -> None:
         )
 
 
+async def _receive_message(
+    app_task: asyncio.Task, sent: asyncio.Queue
+) -> Message | None:
+    """Wait for the next message an app sends on its lifespan scope.
+
+    Args:
+        app_task: The task running the app on the lifespan scope.
+        sent: The queue of messages the app sent.
+
+    Returns:
+        The message, or None if the app returned without sending one.
+    """
+    getter = asyncio.ensure_future(sent.get())
+    try:
+        await asyncio.wait((getter, app_task), return_when=asyncio.FIRST_COMPLETED)
+    finally:
+        if not getter.done():
+            getter.cancel()
+    if getter.done() and not getter.cancelled():
+        return getter.result()
+    # Re-raise the error that ended the app.
+    app_task.result()
+    return None
+
+
+@contextlib.asynccontextmanager
+async def _run_asgi_lifespan(app: ASGIApp, server_scope: Scope):
+    """Send ASGI lifespan events to an app around the body of the context.
+
+    Like an ASGI server's "auto" lifespan mode, an app that raises, returns or
+    sends another message instead of answering the startup event is taken not
+    to support the protocol.
+
+    Args:
+        app: The ASGI app.
+        server_scope: The server's lifespan scope, whose ``state`` the app shares.
+
+    Yields:
+        None once the app has started.
+
+    Raises:
+        RuntimeError: If the app reports that its startup or shutdown failed.
+    """
+    received: asyncio.Queue[Message] = asyncio.Queue()
+    sent: asyncio.Queue[Message] = asyncio.Queue()
+    app_task = asyncio.ensure_future(app({**server_scope}, received.get, sent.put))
+    # A failure is reported through the messages or re-raised from _receive_message.
+    app_task.add_done_callback(lambda task: task.cancelled() or task.exception())
+    try:
+        await received.put({"type": "lifespan.startup"})
+        try:
+            message = await _receive_message(app_task, sent) or {"type": "no answer"}
+        except Exception as err:
+            message = {"type": repr(err)}
+        if message["type"] == "lifespan.startup.failed":
+            msg = f"Lifespan startup of {app!r} failed: {message.get('message', '')}"
+            raise RuntimeError(msg)
+        if message["type"] != "lifespan.startup.complete":
+            logger.info(
+                f"ASGI lifespan appears unsupported by {app!r}: {message['type']}"
+            )
+            yield
+            return
+        try:
+            yield
+        finally:
+            await received.put({"type": "lifespan.shutdown"})
+            message = await _receive_message(app_task, sent)
+            if message is not None and message["type"] == "lifespan.shutdown.failed":
+                msg = (
+                    f"Lifespan shutdown of {app!r} failed: {message.get('message', '')}"
+                )
+                raise RuntimeError(msg)
+    finally:
+        app_task.cancel()
+
+
 @dataclasses.dataclass
 class LifespanMixin(AppMixin):
     """A Mixin that allow tasks to run during the whole app lifespan.
@@ -189,6 +267,28 @@ async def _run_lifespan_tasks(self, starlette_app: Starlette):
             await state_manager.close()
         await prerequisites.close_health_redis()
 
+    @contextlib.asynccontextmanager
+    async def _run_lifespan(
+        self, starlette_app: Starlette, asgi_apps: Sequence[ASGIApp], scope: Scope
+    ):
+        """Run the lifespan tasks, then the ASGI lifespan protocol of each app.
+
+        Shutdown runs in reverse order, so the apps stop before the tasks.
+
+        Args:
+            starlette_app: The Starlette app passed to tasks accepting ``starlette_app``.
+            asgi_apps: The apps to send lifespan events to, in startup order.
+            scope: The server's lifespan scope, whose ``state`` the apps share.
+
+        Yields:
+            None once every lifespan has started.
+        """
+        async with contextlib.AsyncExitStack() as stack:
+            await stack.enter_async_context(self._run_lifespan_tasks(starlette_app))
+            for asgi_app in asgi_apps:
+                await stack.enter_async_context(_run_asgi_lifespan(asgi_app, scope))
+            yield
+
     @overload
     def register_lifespan_task(
         self, task: _LifespanTaskT, **task_kwargs