refactor(api): rebuild main.py lifespan on AsyncExitStack (#59)

- Extract the 80-line inline lifespan out of create_app() into a module-level _lifespan built on AsyncExitStack: each resource registers its cleanup as it is acquired, so the LIFO teardown order (cancel loops -> drain replay tasks -> stop containers -> clear factory -> dispose engine) holds by construction, including when startup fails midway
- Extract named helpers (_reconcile_stale_sessions, _stop_running_containers, _drain_replay_tasks, _cancel_task) in place of the hand-ordered finally block
- Hoist function-level imports to module scope (no cycles) and drop narration comments, including a leftover story-spec reference
- Behavior preserved: same teardown order, same exception-swallowing semantics; 282/282 tests pass, ruff and mypy clean

Author: mariuspruvot

CLAUDE.md

@@ -44,7 +44,7 @@ infra/
 
 ## Key Patterns
 
-- **App factory**: `helprs.main:create_app()` — lifespan manages DB engine
+- **App factory**: `helprs.main:create_app()` — module-level `_lifespan` owns the engine and background loops via `AsyncExitStack`: each resource registers its cleanup at acquisition, teardown runs LIFO (cancel loops → drain replay tasks → stop containers → clear factory → dispose engine)
 - **Flat modules** (identity, installation, webhook, container): `router.py`, `service.py`, `models.py`, `schemas.py`
 - **Container orchestration**: `container` module manages ephemeral Docker lifecycle, credential injection, result relay
 - **Skills as agents**: each skill is a self-contained folder with workflow definitions, mounted into containers

apps/api/src/helprs/main.py

@@ -2,12 +2,16 @@
 
 import asyncio
 import contextlib
-from contextlib import asynccontextmanager
+from collections.abc import AsyncIterator
+from contextlib import AsyncExitStack, asynccontextmanager
 
 import structlog
 from fastapi import APIRouter, FastAPI, Request
 from fastapi.responses import JSONResponse
+from sqlalchemy import text
+from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker
 
+from helprs.admin.views import setup_admin
 from helprs.core.config import get_settings
 from helprs.core.database import (
     clear_session_factory,
@@ -17,35 +21,34 @@
 )
 from helprs.core.exceptions import DomainError, domain_exception_handler
 from helprs.core.middleware import configure_logging, configure_sentry, setup_middleware
+from helprs.modules.container import service as containers
+from helprs.modules.container.router import router as container_router
+from helprs.modules.identity.router import router as identity_router
+from helprs.modules.installation.router import router as installation_router
+from helprs.modules.webhook.repository import get_replayable_events
+from helprs.modules.webhook.router import router as webhook_router
+from helprs.modules.webhook.tasks import process_webhook_event
 
 logger = structlog.get_logger()
 
+SessionFactory = async_sessionmaker[AsyncSession]
+
 _REPLAY_DISCOVERY_TIMEOUT_SECONDS = 10.0
 _REPLAY_CONCURRENCY = 10
 _REAPER_INTERVAL_SECONDS = 300
 _CONTAINER_CLEANUP_INTERVAL_SECONDS = 300
 
 
 async def _replay_pending_webhook_events(app: FastAPI) -> None:
-    """Re-dispatch any webhook_events that survived a crash / are stuck.
-
-    Used by both the lifespan startup path and the periodic reaper loop.
-    Bounded by ``LIMIT`` inside ``get_replayable_events`` and a
-    ``Semaphore(_REPLAY_CONCURRENCY)`` here so a backlog cannot fan out into
-    thousands of concurrent tasks and exhaust the DB pool.
+    """Re-dispatch webhook_events that survived a crash or got stuck.
 
-    Failures in the discovery query never block startup — they are logged
-    and swallowed. A stuck DB connection is bounded by
-    ``_REPLAY_DISCOVERY_TIMEOUT_SECONDS`` via ``asyncio.wait_for``.
-
-    Spawned tasks are tracked on ``app.state.replay_tasks`` so they can be
-    awaited on lifespan shutdown (no orphaned ``create_task`` references,
-    no tasks cancelled mid-commit by ``engine.dispose()``).
+    Runs at startup and on every reaper tick. Discovery is bounded by
+    ``asyncio.wait_for`` and never blocks startup; dispatch is bounded by
+    ``Semaphore(_REPLAY_CONCURRENCY)`` so a backlog cannot exhaust the DB
+    pool. Spawned tasks are tracked on ``app.state.replay_tasks`` so
+    shutdown can await them before the engine is disposed.
     """
-    from helprs.modules.webhook.repository import get_replayable_events
-    from helprs.modules.webhook.tasks import process_webhook_event
-
-    session_factory = app.state.session_factory
+    session_factory: SessionFactory = app.state.session_factory
 
     async def _discover() -> list:
         async with session_factory() as session:
@@ -57,7 +60,6 @@ async def _discover() -> list:
         logger.warning("webhook_replay_discovery_timeout", timeout=_REPLAY_DISCOVERY_TIMEOUT_SECONDS)
         return
     except Exception:
-        # Never block startup on replay discovery failures.
         logger.exception("webhook_replay_discovery_failed")
         return
 
@@ -80,22 +82,17 @@ async def _bounded(event_id) -> None:
 
 
 async def _run_webhook_reaper(app: FastAPI, *, interval_seconds: int = _REAPER_INTERVAL_SECONDS) -> None:
-    """Periodic reaper that re-runs the replay query on a fixed interval.
+    """Re-run the replay query on a fixed interval.
 
-    Complements the boot-time replay: a blip mid-run (e.g., `mark_processed`
-    commit failure) that leaves a row in ``processing`` would otherwise only
-    be recovered on the next restart. With the reaper, recovery happens
-    within one ``interval_seconds`` window.
-
-    Cancelled cleanly by the lifespan teardown via ``task.cancel()``.
+    Complements the boot-time replay: a row left in ``processing`` by a
+    mid-run blip would otherwise only be recovered on the next restart.
     """
     try:
         while True:
             await asyncio.sleep(interval_seconds)
             try:
                 await _replay_pending_webhook_events(app)
             except Exception:
-                # Never crash the reaper loop — keep going on the next tick.
                 logger.exception("webhook_reaper_cycle_failed")
     except asyncio.CancelledError:
         logger.info("webhook_reaper_stopped")
@@ -107,21 +104,15 @@ async def _run_container_cleanup(
     *,
     interval_seconds: int = _CONTAINER_CLEANUP_INTERVAL_SECONDS,
 ) -> None:
-    """Periodic cleanup of expired container sessions.
-
-    Finds sessions past their TTL and destroys their Docker containers.
-    Cancelled cleanly by the lifespan teardown via ``task.cancel()``.
-    """
-    from helprs.modules.container.service import AioDockerClient, cleanup_expired
-
+    """Destroy Docker containers of sessions past their TTL, on a fixed interval."""
     ttl = get_settings().CONTAINER_TTL_SECONDS
     try:
         while True:
             await asyncio.sleep(interval_seconds)
             try:
-                docker = AioDockerClient()
+                docker = containers.AioDockerClient()
                 async with app.state.session_factory() as db:
-                    cleaned = await cleanup_expired(db, docker, ttl_seconds=ttl)
+                    cleaned = await containers.cleanup_expired(db, docker, ttl_seconds=ttl)
                     await db.commit()
                     if cleaned:
                         logger.info("container_cleanup_cycle", cleaned=cleaned)
@@ -132,144 +123,112 @@ async def _run_container_cleanup(
         raise
 
 
-def create_app() -> FastAPI:
-    settings = get_settings()
+async def _reconcile_stale_sessions(session_factory: SessionFactory) -> None:
+    """Mark sessions left RUNNING/PENDING by a previous run as FAILED."""
+    try:
+        async with session_factory() as db:
+            await containers.reconcile_stale_sessions(db)
+            await db.commit()
+    except Exception:
+        logger.exception("session_reconciliation_failed")
 
-    # Configure observability
-    configure_logging()
-    configure_sentry(settings)
 
-    @asynccontextmanager
-    async def lifespan(app: FastAPI):
-        """Manage database engine lifecycle, admin setup, and webhook reaper."""
+async def _stop_running_containers(session_factory: SessionFactory) -> None:
+    """Stop every container still running before the engine goes away."""
+    try:
+        docker = containers.AioDockerClient()
+        async with session_factory() as db:
+            stopped = await containers.cleanup_all_running(db, docker)
+            await db.commit()
+            if stopped:
+                logger.info("shutdown_containers_stopped", count=stopped)
+        await docker.close()
+    except Exception:
+        logger.exception("shutdown_container_cleanup_failed")
+
+
+async def _drain_replay_tasks(app: FastAPI) -> None:
+    """Await in-flight replay tasks so none writes to a disposed pool."""
+    if app.state.replay_tasks:
+        await asyncio.gather(*app.state.replay_tasks, return_exceptions=True)
+
+
+async def _cancel_task(task: asyncio.Task) -> None:
+    task.cancel()
+    with contextlib.suppress(asyncio.CancelledError, Exception):
+        await task
+
+
+@asynccontextmanager
+async def _lifespan(app: FastAPI) -> AsyncIterator[None]:
+    """Acquire long-lived resources; ``AsyncExitStack`` releases them in reverse order."""
+    settings = get_settings()
+    async with AsyncExitStack() as stack:
         engine = create_engine()
-        reaper_task: asyncio.Task | None = None
-        cleanup_task: asyncio.Task | None = None
-        try:
-            session_factory = create_session_factory(engine)
-            app.state.engine = engine
-            app.state.session_factory = session_factory
-            # Register the factory for ``get_db_context`` (used by
-            # background tasks that need DB access outside the request
-            # dependency graph).
-            set_session_factory(session_factory)
-            app.state.replay_semaphore = asyncio.Semaphore(_REPLAY_CONCURRENCY)
-            app.state.replay_tasks = set()
-
-            # Admin panel (needs engine)
-            from helprs.admin.views import setup_admin
-
-            setup_admin(app, engine, settings.SECRET_KEY)
-
-            # Crash-replay: any webhook_events still in pending/processing
-            # after a grace period are re-dispatched. Fire-and-forget so the
-            # server starts serving traffic immediately (AC #2).
-            await _replay_pending_webhook_events(app)
-
-            # Reconcile container sessions left in RUNNING/PENDING by a
-            # prior crash — mark them FAILED immediately rather than
-            # waiting for TTL expiry.
-            try:
-                from helprs.modules.container.service import reconcile_stale_sessions
+        stack.push_async_callback(engine.dispose)
 
-                async with session_factory() as db:
-                    await reconcile_stale_sessions(db)
-                    await db.commit()
-            except Exception:
-                logger.exception("session_reconciliation_failed")
-
-            # Periodic reaper: handles rows that get stuck mid-run (e.g.
-            # mark_processed commit failure) without waiting for the next
-            # restart.
-            reaper_task = asyncio.create_task(_run_webhook_reaper(app))
-            cleanup_task = asyncio.create_task(_run_container_cleanup(app))
-
-            yield
-        finally:
-            if cleanup_task is not None:
-                cleanup_task.cancel()
-                with contextlib.suppress(asyncio.CancelledError, Exception):
-                    await cleanup_task
-
-            if reaper_task is not None:
-                reaper_task.cancel()
-                with contextlib.suppress(asyncio.CancelledError, Exception):
-                    await reaper_task
-
-            # Wait for in-flight replay tasks before disposing the engine so
-            # they don't end up writing to a closed pool.
-            tracked: set[asyncio.Task] = getattr(app.state, "replay_tasks", set())
-            if tracked:
-                await asyncio.gather(*tracked, return_exceptions=True)
-
-            # Stop all running containers before shutting down.
-            try:
-                from helprs.modules.container.service import AioDockerClient, cleanup_all_running
+        session_factory = create_session_factory(engine)
+        app.state.engine = engine
+        app.state.session_factory = session_factory
+        app.state.replay_semaphore = asyncio.Semaphore(_REPLAY_CONCURRENCY)
+        app.state.replay_tasks = set()
 
-                docker = AioDockerClient()
-                async with session_factory() as db:
-                    stopped = await cleanup_all_running(db, docker)
-                    await db.commit()
-                    if stopped:
-                        logger.info("shutdown_containers_stopped", count=stopped)
-                await docker.close()
-            except Exception:
-                logger.exception("shutdown_container_cleanup_failed")
+        # get_db_context() (background tasks outside the request dependency
+        # graph) resolves sessions through this module-level registry.
+        set_session_factory(session_factory)
+        stack.callback(clear_session_factory)
+
+        setup_admin(app, engine, settings.SECRET_KEY)
+
+        await _replay_pending_webhook_events(app)
+        await _reconcile_stale_sessions(session_factory)
+
+        stack.push_async_callback(_stop_running_containers, session_factory)
+        stack.push_async_callback(_drain_replay_tasks, app)
+
+        for loop in (_run_webhook_reaper(app), _run_container_cleanup(app)):
+            stack.push_async_callback(_cancel_task, asyncio.create_task(loop))
+
+        yield
+
+
+async def _handle_unhandled_exception(request: Request, exc: Exception) -> JSONResponse:
+    logger.exception("unhandled_exception", path=request.url.path)
+    return JSONResponse(
+        status_code=500,
+        content={"error": "internal_error", "message": "An unexpected error occurred"},
+    )
 
-            clear_session_factory()
-            await engine.dispose()
+
+def create_app() -> FastAPI:
+    settings = get_settings()
+    configure_logging()
+    configure_sentry(settings)
 
     app = FastAPI(
         title="helPRs API",
         description="AI-powered pull request review assistant",
         version="0.1.0",
-        lifespan=lifespan,
+        lifespan=_lifespan,
     )
 
-    # Exception handlers
     app.add_exception_handler(DomainError, domain_exception_handler)
-
-    async def _unhandled_exception_handler(request: Request, exc: Exception) -> JSONResponse:
-        logger.exception("unhandled_exception", path=request.url.path)
-        return JSONResponse(
-            status_code=500,
-            content={"error": "internal_error", "message": "An unexpected error occurred"},
-        )
-
-    app.add_exception_handler(Exception, _unhandled_exception_handler)
-
-    # Middleware (CORS, logging, rate limiting)
+    app.add_exception_handler(Exception, _handle_unhandled_exception)
     setup_middleware(app, settings)
 
-    # API router
     api_router = APIRouter(prefix="/api/v1")
-
-    from helprs.modules.container.router import router as container_router
-    from helprs.modules.identity.router import router as identity_router
-    from helprs.modules.installation.router import router as installation_router
-    from helprs.modules.webhook.router import router as webhook_router
-
-    api_router.include_router(container_router)
-    api_router.include_router(identity_router)
-    api_router.include_router(installation_router)
-    api_router.include_router(webhook_router)
-
+    for router in (container_router, identity_router, installation_router, webhook_router):
+        api_router.include_router(router)
     app.include_router(api_router)
 
-    # Health check
     @app.get("/health")
     async def health_check():
-        from sqlalchemy import text
-
         try:
             async with app.state.session_factory() as db:
                 await db.execute(text("SELECT 1"))
             return {"status": "ok", "db": "ok"}
         except Exception:
-            return JSONResponse(
-                status_code=503,
-                content={"status": "degraded", "db": "unreachable"},
-            )
+            return JSONResponse(status_code=503, content={"status": "degraded", "db": "unreachable"})
 
     return app