fix(fastapi)!: tear down after user lifespan shutdown

lite_bootstrap/bootstrappers/fastapi_bootstrapper.py

@@ -35,7 +35,6 @@
 if import_checker.is_fastapi_installed:
     import fastapi
     from fastapi.middleware.cors import CORSMiddleware
-    from fastapi.routing import _merge_lifespan_context
 
 if import_checker.is_fastapi_installed and import_checker.is_sentry_installed:
     from sentry_sdk.integrations.fastapi import FastApiIntegration
@@ -299,11 +298,14 @@ def __init__(self, bootstrap_config: FastAPIConfig) -> None:
         self._attach_teardown_once(application, lambda: self._wrap_lifespan(application))
 
     def _wrap_lifespan(self, application: "fastapi.FastAPI") -> None:
-        old_lifespan_manager = application.router.lifespan_context
-        application.router.lifespan_context = _merge_lifespan_context(
-            old_lifespan_manager,
-            self.lifespan_manager,
-        )
+        user_lifespan = application.router.lifespan_context
+
+        @contextlib.asynccontextmanager
+        async def merged_lifespan(app: "fastapi.FastAPI") -> AsyncGenerator[Mapping[str, typing.Any]]:
+            async with self.lifespan_manager(app) as bootstrap_state, user_lifespan(app) as user_state:
+                yield {**bootstrap_state, **(user_state or {})}
+
+        application.router.lifespan_context = merged_lifespan
 
     def is_ready(self) -> bool:
         return import_checker.is_fastapi_installed

pyproject.toml

@@ -82,8 +82,6 @@ free-all = [
     "lite-bootstrap[sentry,otl,logging,pyroscope]",
 ]
 fastapi = [
-    # >=0.112.2 for fastapi.routing._merge_lifespan_context, which FastAPIBootstrapper._wrap_lifespan
-    # uses to chain teardown onto the user's lifespan. Private, so it can move without a deprecation.
     "fastapi>=0.112.2",
     # fastapi 0.112.2 requires typing-extensions>=4.8.
     "typing-extensions>=4.8",

tests/test_fastapi_bootstrap.py

@@ -5,7 +5,7 @@
 import typing
 import uuid
 import warnings
-from collections.abc import Generator
+from collections.abc import AsyncGenerator, Generator
 from unittest.mock import patch
 
 import fastapi
@@ -64,6 +64,33 @@ def fastapi_config() -> FastAPIConfig:
     )
 
 
+def test_fastapi_exports_user_shutdown_span_and_preserves_lifespan_state() -> None:
+    @contextlib.asynccontextmanager
+    async def lifespan(_app: fastapi.FastAPI) -> AsyncGenerator[dict[str, str]]:
+        yield {"service": "ready"}
+        with get_tracer_provider().get_tracer(__name__).start_as_current_span("user-shutdown"):
+            pass
+
+    app = fastapi.FastAPI(lifespan=lifespan)
+
+    @app.get("/state")
+    async def read_state(request: fastapi.Request) -> str:
+        return request.state.service
+
+    bootstrapper = FastAPIBootstrapper(bootstrap_config=FastAPIConfig(application=app, opentelemetry_log_traces=True))
+    application = bootstrapper.bootstrap()
+    provider = get_tracer_provider()
+    assert isinstance(provider, SDKTracerProvider)
+    exporter = InMemorySpanExporter()
+    provider.add_span_processor(SimpleSpanProcessor(exporter))
+
+    with TestClient(application) as client:
+        assert client.get("/state").json() == "ready"
+
+    assert "user-shutdown" in [span.name for span in exporter.get_finished_spans()]
+    assert not bootstrapper.is_bootstrapped
+
+
 def _instrument_app_kwargs(config: FastAPIConfig) -> dict[str, typing.Any]:
     """Bootstrap with FastAPIInstrumentor stubbed, and return what it was handed."""
     with patch.object(fastapi_bootstrapper, "FastAPIInstrumentor") as mock_instrumentor: