diff --git a/docs/api-routes/overview.md b/docs/api-routes/overview.md index 51a160b9e1a..53b0ff9e1b2 100644 --- a/docs/api-routes/overview.md +++ b/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. diff --git a/docs/utility_methods/lifespan_tasks.md b/docs/utility_methods/lifespan_tasks.md index bf9ac97515b..0b107bded70 100644 --- a/docs/utility_methods/lifespan_tasks.md +++ b/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() diff --git a/news/+api-transformer-lifespan.bugfix.md b/news/+api-transformer-lifespan.bugfix.md new file mode 100644 index 00000000000..99bdc98b82b --- /dev/null +++ b/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. diff --git a/news/+lighter-asgi-entrypoint.performance.md b/news/+lighter-asgi-entrypoint.performance.md new file mode 100644 index 00000000000..85fe8443ca4 --- /dev/null +++ b/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. diff --git a/reflex/app.py b/reflex/app.py index de1cf8f71b0..635405df1e7 100644 --- a/reflex/app.py +++ b/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,6 +1031,20 @@ 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. @@ -982,13 +1052,7 @@ def _add_cors(api: Starlette): 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: diff --git a/reflex/app_mixins/lifespan.py b/reflex/app_mixins/lifespan.py index 4dd9964cc1a..d6adba0b2e4 100644 --- a/reflex/app_mixins/lifespan.py +++ b/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 diff --git a/tests/integration/test_lifespan.py b/tests/integration/test_lifespan.py index 6adac44175c..b9ee87f3e1d 100644 --- a/tests/integration/test_lifespan.py +++ b/tests/integration/test_lifespan.py @@ -32,6 +32,7 @@ def LifespanApp( lifespan_task_global = 0 lifespan_context_global = 0 + transformer_lifespan_global = 0 # noqa: F841 (read from the module by the test) raw_asyncio_task_global = 0 connected_tokens: set[str] = set() @@ -146,7 +147,20 @@ def index(): from fastapi import FastAPI - app = rx.App(api_transformer=FastAPI() if mount_api_transformer else None) + @asynccontextmanager + async def transformer_lifespan(app: FastAPI): + global transformer_lifespan_global + transformer_lifespan_global += 1 # pyright: ignore[reportUnboundVariable] + try: + yield + finally: + transformer_lifespan_global += 1 + + app = rx.App( + api_transformer=FastAPI(lifespan=transformer_lifespan) + if mount_api_transformer + else None + ) if mount_cached_fastapi: assert app._api is not None @@ -296,6 +310,8 @@ def test_lifespan(lifespan_app: AppHarness): assert lifespan_app.poll_for_content(context_global, exp_not_equal="0") == "2" assert lifespan_app.app_module.lifespan_context_global == 2 + transformer_starts = 1 if lifespan_app.app_module.app.api_transformer else 0 + assert lifespan_app.app_module.transformer_lifespan_global == transformer_starts original_task_global_text = task_global.text original_task_global_value = int(original_task_global_text) @@ -314,6 +330,7 @@ def test_lifespan(lifespan_app: AppHarness): assert lifespan_app.app_module.lifespan_task_global == 0 assert lifespan_app.app_module.lifespan_context_global == 4 assert lifespan_app.app_module.raw_asyncio_task_global == 0 + assert lifespan_app.app_module.transformer_lifespan_global == 2 * transformer_starts # --- Do NOT add new test cases below this line. --- diff --git a/tests/units/app_mixins/test_lifespan.py b/tests/units/app_mixins/test_lifespan.py index 16b740ea1d9..4d6b9758d79 100644 --- a/tests/units/app_mixins/test_lifespan.py +++ b/tests/units/app_mixins/test_lifespan.py @@ -4,9 +4,12 @@ import asyncio import contextlib +import logging +from typing import Any import pytest from reflex_base.utils.exceptions import InvalidLifespanTaskTypeError +from reflex_base.utils.types import Receive, Scope, Send from starlette.applications import Starlette from reflex.app_mixins.lifespan import LifespanMixin @@ -360,3 +363,138 @@ async def run(): finally: release.set() await shutdown + + +def _lifespan_scope() -> dict[str, Any]: + """Create a server's lifespan scope. + + Returns: + The lifespan scope, with an empty state. + """ + return {"type": "lifespan", "asgi": {"version": "3.0"}, "state": {}} + + +def _raw_lifespan_app(events: list[str], name: str, fail: str | None = None): + """Create a plain ASGI app that answers lifespan events itself. + + Args: + events: The list to record lifespan events in. + name: The name of the app in the recorded events. + fail: "startup" or "shutdown" to report that phase as failed. + + Returns: + The ASGI app. + """ + + async def app(scope: Scope, receive: Receive, send: Send) -> None: + assert scope["type"] == "lifespan" + for phase in ("startup", "shutdown"): + await receive() + events.append(f"{name} {phase}") + outcome = "failed" if fail == phase else "complete" + await send({"type": f"lifespan.{phase}.{outcome}", "message": "boom"}) + if fail == phase: + return + scope["state"][name] = phase + + return app + + +@pytest.mark.asyncio +async def test_run_lifespan_sends_lifespan_events_after_the_tasks(): + """Apps start after the tasks, in order, stop in reverse, and share the state.""" + mixin = LifespanMixin() + events: list[str] = [] + + @contextlib.asynccontextmanager + async def task(): + events.append("task start") + yield + events.append("task stop") + + @contextlib.asynccontextmanager + async def starlette_lifespan(app: Starlette): + events.append("starlette startup") + yield {"db": "conn"} + events.append("starlette shutdown") + + mixin.register_lifespan_task(task) + scope = _lifespan_scope() + apps = [_raw_lifespan_app(events, "raw"), Starlette(lifespan=starlette_lifespan)] + + async with mixin._run_lifespan(Starlette(), apps, scope): + assert scope["state"] == {"raw": "startup", "db": "conn"} + assert events == [ + "task start", + "raw startup", + "starlette startup", + "starlette shutdown", + "raw shutdown", + "task stop", + ] + + +@pytest.mark.asyncio +async def test_run_lifespan_skips_apps_without_lifespan_support(caplog): + """An app that does not answer the startup event is skipped.""" + mixin = LifespanMixin() + events: list[str] = [] + + async def http_only_app(scope: Scope, receive: Receive, send: Send) -> None: # noqa: RUF029 + assert scope["type"] == "http" + + async def silent_app(scope: Scope, receive: Receive, send: Send) -> None: + pass + + async def http_response_app(scope: Scope, receive: Receive, send: Send) -> None: + await send({"type": "http.response.start", "status": 200, "headers": []}) + + apps = [ + http_only_app, + silent_app, + http_response_app, + _raw_lifespan_app(events, "raw"), + ] + with caplog.at_level(logging.INFO): + async with mixin._run_lifespan(Starlette(), apps, _lifespan_scope()): + assert events == ["raw startup"] + assert events == ["raw startup", "raw shutdown"] + assert "appears unsupported" in caplog.text + + +@pytest.mark.asyncio +async def test_run_lifespan_startup_failure_stops_started_lifespans(): + """An app failing to start stops the lifespans already started.""" + mixin = LifespanMixin() + events: list[str] = [] + apps = [ + _raw_lifespan_app(events, "first"), + _raw_lifespan_app(events, "failing", fail="startup"), + _raw_lifespan_app(events, "never"), + ] + + with pytest.raises(RuntimeError, match=r"startup of .* failed: boom"): + async with mixin._run_lifespan(Starlette(), apps, _lifespan_scope()): + pass + assert events == ["first startup", "failing startup", "first shutdown"] + + +@pytest.mark.asyncio +async def test_run_lifespan_shutdown_failure_still_stops_the_rest(): + """An app failing to stop is reported after the other lifespans stop.""" + mixin = LifespanMixin() + events: list[str] = [] + apps = [ + _raw_lifespan_app(events, "first"), + _raw_lifespan_app(events, "failing", fail="shutdown"), + ] + + with pytest.raises(RuntimeError, match=r"shutdown of .* failed: boom"): + async with mixin._run_lifespan(Starlette(), apps, _lifespan_scope()): + pass + assert events == [ + "first startup", + "failing startup", + "failing shutdown", + "first shutdown", + ] diff --git a/tests/units/test_app.py b/tests/units/test_app.py index 270bc15db03..e00eb30c639 100644 --- a/tests/units/test_app.py +++ b/tests/units/test_app.py @@ -43,7 +43,7 @@ from reflex_base.style import Style from reflex_base.utils import exceptions, format, memo_paths from reflex_base.utils.imports import ImportVar -from reflex_base.utils.types import Receive, Scope, Send +from reflex_base.utils.types import ASGIApp, Receive, Scope, Send from reflex_base.vars.base import computed_var from reflex_components_core.base.bare import Bare from reflex_components_core.base.fragment import Fragment @@ -55,8 +55,9 @@ from starlette.applications import Starlette from starlette.datastructures import FormData, Headers, UploadFile from starlette.exceptions import HTTPException -from starlette.requests import ClientDisconnect -from starlette.responses import StreamingResponse +from starlette.middleware.cors import CORSMiddleware +from starlette.requests import ClientDisconnect, Request +from starlette.responses import JSONResponse, StreamingResponse from starlette.testclient import TestClient from starlette_admin.auth import AuthProvider @@ -67,9 +68,9 @@ App, ComponentCallable, EventNamespace, - _ContextMiddleware, _decode_asgi_headers, _DrainOversizedBodyMiddleware, + _RootMiddleware, _sio_dumps, _sio_loads, default_overlay_component, @@ -589,14 +590,230 @@ def test_admin_dashboard_routes_remain_reversible( ) mocker.patch.object(app, "_compile") api = app() - assert isinstance(api, Starlette) - assert str(api.url_path_for("admin:index")) == "/admin/" + router = app.api_transformer or app._api + assert isinstance(router, Starlette) + assert str(router.url_path_for("admin:index")) == "/admin/" with TestClient(api) as client: response = client.get("/admin/") assert response.status_code == 200 assert "Reflex Admin Dashboard" in response.text +def test_api_transformer_lifespans_run(mocker: MockerFixture): + """Send ASGI lifespan events through the api_transformer apps. + + The outermost app and each app a Starlette transformer mounts get them, after + the Reflex lifespan tasks start; apps mounted on a transformer by the user + do not. + + Args: + mocker: The mock fixture for skipping frontend compilation. + """ + fastapi = pytest.importorskip("fastapi") + events: list[str] = [] + + def recording_lifespan(name: str, state: dict[str, str] | None = None) -> Any: + @contextlib.asynccontextmanager + async def lifespan(app: Any): + events.append(f"{name} start") + yield state + events.append(f"{name} stop") + + return lifespan + + async def read_state(request: Request) -> JSONResponse: # noqa: RUF029 + return JSONResponse({"db": request.state.db, "own_app": request.app is outer}) + + def record_scope_types(asgi_app: ASGIApp) -> ASGIApp: + async def recorder(scope: Scope, receive: Receive, send: Send) -> None: + events.append(f"callable {scope['type']}") + await asgi_app(scope, receive, send) + + return recorder + + inner = fastapi.FastAPI(lifespan=recording_lifespan("inner")) + inner.mount("/sub", Starlette(lifespan=recording_lifespan("mounted"))) + outer = Starlette(lifespan=recording_lifespan("outer", {"db": "conn"})) + outer.add_route("/db", read_state) + app = App(api_transformer=[inner, record_scope_types, outer]) + app.register_lifespan_task(recording_lifespan("reflex")) + mocker.patch.object(app, "_compile") + + with TestClient(app()) as client: + assert events == [ + "reflex start", + "callable lifespan", + "inner start", + "outer start", + ] + assert client.get("/db").json() == {"db": "conn", "own_app": True} + assert client.get("/ping").text == '"pong"' + assert events[4:] == ["callable http", "outer stop", "inner stop", "reflex stop"] + + +@pytest.mark.parametrize("also_registered", [False, True]) +def test_api_transformer_lifespan_also_registered_warns( + mocker: MockerFixture, caplog: pytest.LogCaptureFixture, also_registered: bool +): + """Warn when a transformer's lifespan is also registered as a lifespan task. + + Args: + mocker: The mock fixture for skipping frontend compilation. + caplog: The log capture fixture. + also_registered: Whether the lifespan is also registered as a task. + """ + fastapi = pytest.importorskip("fastapi") + + @contextlib.asynccontextmanager + async def lifespan(app: Any): + yield + + app = App(api_transformer=fastapi.FastAPI(lifespan=lifespan)) + if also_registered: + app.register_lifespan_task(lifespan) + mocker.patch.object(app, "_compile") + + with caplog.at_level(logging.WARNING, logger="reflex.app"): + app() + assert ("runs twice" in caplog.text) is also_registered + + +def test_api_transformer_lifespan_startup_failure_stops_the_app( + mocker: MockerFixture, +): + """A transformer failing to start fails the server's startup. + + Args: + mocker: The mock fixture for skipping frontend compilation. + """ + stopped: list[str] = [] + + @contextlib.asynccontextmanager + async def reflex_task(): + try: + yield + finally: + stopped.append("reflex") + + @contextlib.asynccontextmanager + async def failing_lifespan(app: Starlette): + msg = "boom" + raise ValueError(msg) + yield + + app = App(api_transformer=Starlette(lifespan=failing_lifespan)) + app.register_lifespan_task(reflex_task) + mocker.patch.object(app, "_compile") + + with pytest.raises(RuntimeError, match="boom"), TestClient(app()): + pass + assert stopped == ["reflex"] + + +def test_api_transformer_lifespan_shutdown_failure_is_reported( + mocker: MockerFixture, +): + """A transformer failing to stop fails the server's shutdown after Reflex stops. + + Args: + mocker: The mock fixture for skipping frontend compilation. + """ + stopped: list[str] = [] + + @contextlib.asynccontextmanager + async def reflex_task(): + try: + yield + finally: + stopped.append("reflex") + + @contextlib.asynccontextmanager + async def failing_lifespan(app: Starlette): + yield + msg = "boom" + raise ValueError(msg) + + app = App(api_transformer=Starlette(lifespan=failing_lifespan)) + app.register_lifespan_task(reflex_task) + mocker.patch.object(app, "_compile") + + status_codes: list[int] = [] + with pytest.raises(RuntimeError, match="boom"), TestClient(app()) as client: + status_codes.append(client.get("/ping").status_code) + assert status_codes == [200] + assert stopped == ["reflex"] + + +@pytest.mark.parametrize("callable_last", [False, True]) +def test_api_transformer_routes_get_cors_headers( + mocker: MockerFixture, callable_last: bool +): + """Routes of the outermost transformer answer with CORS headers exactly once. + + Args: + mocker: The mock fixture for skipping frontend compilation. + callable_last: Whether the outermost transformer is a callable. + """ + + def add_route(asgi_app: ASGIApp) -> ASGIApp: + async def router(scope: Scope, receive: Receive, send: Send) -> None: + if scope["type"] == "http" and scope["path"] == "/extra": + await JSONResponse("extra")(scope, receive, send) + else: + await asgi_app(scope, receive, send) + + return router + + starlette_app = Starlette() + starlette_app.add_route("/extra", lambda _: JSONResponse("extra")) + app = App(api_transformer=add_route if callable_last else starlette_app) + mocker.patch.object(app, "_compile") + + with TestClient(app()) as client: + for path in ("/extra", "/ping"): + response = client.get(path, headers={"Origin": "http://example.com"}) + assert response.status_code == 200 + assert response.headers.get_list("access-control-allow-origin") == [ + "http://example.com" + ] + + +@pytest.mark.parametrize("transformer", [None, "starlette", "callable"]) +def test_unhandled_route_error_gets_cors_headers( + mocker: MockerFixture, transformer: str | None +): + """A 500 from an unhandled route error answers with CORS headers exactly once. + + Args: + mocker: The mock fixture for skipping frontend compilation. + transformer: The kind of api_transformer to use, if any. + """ + + async def boom(request: Request) -> JSONResponse: # noqa: RUF029 + msg = "boom" + raise ValueError(msg) + + starlette_app = Starlette() + api_transformers = { + None: None, + "starlette": starlette_app, + "callable": lambda asgi_app: asgi_app, + } + app = App(api_transformer=api_transformers[transformer]) + assert app._api is not None + # The route sits on the outermost Starlette app. + router = starlette_app if transformer == "starlette" else app._api + router.add_route("/boom", boom) + mocker.patch.object(app, "_compile") + + with TestClient(app(), raise_server_exceptions=False) as client: + response = client.get("/boom", headers={"Origin": "http://example.com"}) + assert response.status_code == 500 + assert response.headers.get_list("access-control-allow-origin") == [ + "http://example.com" + ] + + @pytest.mark.skipif( not find_spec("starlette_admin") or not find_spec("sqlmodel") @@ -3536,7 +3753,9 @@ def test_call_app(): app = App() app._compile = unittest.mock.Mock() api = app() - assert isinstance(api, Starlette) + assert isinstance(api, _RootMiddleware) + assert isinstance(api.app, CORSMiddleware) + assert api.app.app is app._api def _probe_worker_token_identity(app: App, redis: AsyncMock, sender: Connection): @@ -4164,11 +4383,11 @@ def _test(): isolated_context.run(_test) -def test_context_middleware_is_registered_as_a_class( +def test_registered_middleware_are_classes( compilable_app: tuple[App, Path], mocker: MockerFixture, ): - """The context middleware must reach Starlette as a class, not a bound method. + """Every middleware registered on a Starlette app must be a class, not a bound method. ASGI instrumentation libraries (sentry-sdk's Starlette integration among them) wrap every registered middleware by assigning to ``cls.__call__``. A bound @@ -4180,24 +4399,21 @@ def test_context_middleware_is_registered_as_a_class( asgi_app = app() - assert isinstance(asgi_app, Starlette) - registered = [m.cls for m in asgi_app.user_middleware] - assert _ContextMiddleware in registered, ( - f"context middleware not registered as a class, got {registered}" - ) + assert isinstance(asgi_app, _RootMiddleware) + assert app._api is not None # Mimic the instrumentation's patch on every middleware in the stack. - for cls in registered: - cls.__call__ = cls.__call__ + for middleware in app._api.user_middleware: + middleware.cls.__call__ = middleware.cls.__call__ -async def test_context_middleware_sets_contexts(app_with_processor: App): - """The context middleware attaches Reflex contexts before calling the app.""" +async def test_root_middleware_sets_contexts(app_with_processor: App): + """The root middleware attaches Reflex contexts before calling the app.""" seen: list[tuple[RegistrationContext, EventContext]] = [] async def inner_app(scope: Scope, receive: Receive, send: Send) -> None: # noqa: RUF029 seen.append((RegistrationContext.get(), EventContext.get())) - await _ContextMiddleware(inner_app, app_with_processor)( + await _RootMiddleware(inner_app, app_with_processor, contextlib.nullcontext)( {"type": "http"}, AsyncMock(), AsyncMock() )