refactor(app): move task queue startup to lifespan

Author: haasonsaas

app.py

@@ -8,6 +8,8 @@
 from datetime import datetime, timedelta
 from typing import Any
 
+from contextlib import asynccontextmanager
+
 from fastapi import Depends, FastAPI, HTTPException, WebSocket, WebSocketDisconnect
 from fastapi.responses import PlainTextResponse
 from pydantic import BaseModel, model_validator
@@ -68,36 +70,39 @@
 from agent_pm.storage.tasks import TaskStatus, get_task_queue
 from agent_pm.tools import registry
 
+@asynccontextmanager
+async def lifespan(_app: FastAPI):
+    global _task_queue
+    _task_queue = await get_task_queue()
+    await _task_queue.start()
+    logger.info("Agent PM service started")
+    try:
+        yield
+    finally:
+        if _task_queue:
+            await _task_queue.stop()
+        logger.info("Agent PM service stopped")
+
+
+lifespan_app = FastAPI(title="Agent PM", version="0.1.0", lifespan=lifespan)
+
 if settings.log_format == "json":
     configure_structured_logging()
 else:
     configure_logging(settings.trace_dir)
 
 logger = logging.getLogger(__name__)
-app = FastAPI(title="Agent PM", version="0.1.0")
 _jira_lock = asyncio.Lock()
 _task_queue = None
 
-plugin_registry.attach_app(app)
+plugin_registry.attach_app(lifespan_app)
 
 
 class FollowupUpdate(BaseModel):
     status: str
 
 
-@app.on_event("startup")
-async def startup_event():
-    global _task_queue
-    _task_queue = await get_task_queue()
-    await _task_queue.start()
-    logger.info("Agent PM service started")
-
-
-@app.on_event("shutdown")
-async def shutdown_event():
-    if _task_queue:
-        await _task_queue.stop()
-    logger.info("Agent PM service stopped")
+app = lifespan_app
 
 
 async def ensure_project_allowed(plan: TicketPlan) -> TicketPlan:
@@ -435,10 +440,9 @@ async def metrics() -> PlainTextResponse:
 @app.get("/tasks")
 async def list_tasks(status: str | None = None, limit: int = 50, _admin_key: AdminKeyDep = None) -> dict[str, Any]:
     """List all tasks with optional status filter."""
-    if not _task_queue:
-        raise HTTPException(status_code=503, detail="Task queue not initialized")
+    task_queue = await get_task_queue()
     task_status = TaskStatus(status) if status else None
-    tasks = await _task_queue.list_tasks(status=task_status, limit=limit)
+    tasks = await task_queue.list_tasks(status=task_status, limit=limit)
     return {
         "tasks": [
             {
@@ -463,9 +467,8 @@ async def list_dead_letter(
     error_type: str | None = None,
     _admin_key: AdminKeyDep = None,
 ) -> dict[str, Any]:
-    if not _task_queue:
-        raise HTTPException(status_code=503, detail="Task queue not initialized")
-    items, total = await _task_queue.list_dead_letters(
+    task_queue = await get_task_queue()
+    items, total = await task_queue.list_dead_letters(
         limit=limit, offset=offset, workflow_id=workflow_id, error_type=error_type
     )
     return {
@@ -487,48 +490,43 @@ async def list_dead_letter(
 
 @app.get("/tasks/dead-letter/{task_id}")
 async def get_dead_letter(task_id: str, _admin_key: AdminKeyDep = None) -> dict[str, Any]:
-    if not _task_queue:
-        raise HTTPException(status_code=503, detail="Task queue not initialized")
-    item = await _task_queue.get_dead_letter(task_id)
+    task_queue = await get_task_queue()
+    item = await task_queue.get_dead_letter(task_id)
     if not item:
         raise HTTPException(status_code=404, detail="Dead-letter task not found")
     return item
 
 
 @app.delete("/tasks/dead-letter/{task_id}")
 async def delete_dead_letter(task_id: str, _admin_key: AdminKeyDep = None) -> dict[str, Any]:
-    if not _task_queue:
-        raise HTTPException(status_code=503, detail="Task queue not initialized")
-    await _task_queue.delete_dead_letter(task_id)
+    task_queue = await get_task_queue()
+    await task_queue.delete_dead_letter(task_id)
     return {"task_id": task_id, "status": "deleted"}
 
 
 @app.delete("/tasks/dead-letter")
 async def purge_dead_letters(
     older_than_minutes: int | None = None, _admin_key: AdminKeyDep = None
 ) -> dict[str, int]:
-    if not _task_queue:
-        raise HTTPException(status_code=503, detail="Task queue not initialized")
+    task_queue = await get_task_queue()
     if older_than_minutes is None:
-        deleted = await _task_queue.purge_dead_letters()
+        deleted = await task_queue.purge_dead_letters()
     else:
-        deleted = await _task_queue.purge_dead_letters_older_than(timedelta(minutes=older_than_minutes))
+        deleted = await task_queue.purge_dead_letters_older_than(timedelta(minutes=older_than_minutes))
     return {"deleted": deleted}
 
 
 @app.get("/tasks/workers")
 async def worker_status(_admin_key: AdminKeyDep = None) -> dict[str, Any]:
-    if not _task_queue:
-        raise HTTPException(status_code=503, detail="Task queue not initialized")
-    return {"workers": await _task_queue.worker_heartbeats()}
+    task_queue = await get_task_queue()
+    return {"workers": await task_queue.worker_heartbeats()}
 
 
 @app.get("/tasks/{task_id}")
 async def get_task(task_id: str, _admin_key: AdminKeyDep = None) -> dict[str, Any]:
     """Get task status by ID."""
-    if not _task_queue:
-        raise HTTPException(status_code=503, detail="Task queue not initialized")
-    task = await _task_queue.get_task(task_id)
+    task_queue = await get_task_queue()
+    task = await task_queue.get_task(task_id)
     if not task:
         raise HTTPException(status_code=404, detail="Task not found")
     return {
@@ -545,9 +543,8 @@ async def get_task(task_id: str, _admin_key: AdminKeyDep = None) -> dict[str, An
 
 @app.post("/tasks/dead-letter/{task_id}/requeue")
 async def requeue_dead_letter(task_id: str, _admin_key: AdminKeyDep = None) -> dict[str, Any]:
-    if not _task_queue:
-        raise HTTPException(status_code=503, detail="Task queue not initialized")
-    payload = await _task_queue.requeue_dead_letter(task_id)
+    task_queue = await get_task_queue()
+    payload = await task_queue.requeue_dead_letter(task_id)
     if payload is None:
         raise HTTPException(status_code=404, detail="Dead-letter task not found")
     return {"task_id": payload.get("task_id", task_id), "status": "requeued"}

tests/test_tasks_api.py

@@ -13,11 +13,11 @@
 async def _create_client():
     transport = ASGITransport(app=app)
     client = AsyncClient(transport=transport, base_url="http://test")
-    await app_module.startup_event()
+    await app.router.startup()
 
     async def cleanup():
+        await app.router.shutdown()
         await client.aclose()
-        await app_module.shutdown_event()
 
     client.cleanup = cleanup  # type: ignore[attr-defined]
     return client
@@ -30,16 +30,14 @@ async def test_tasks_admin_endpoints_with_memory_backend(monkeypatch):
 
     try:
         client = await _create_client()
-        try:
-            response = await client.get("/tasks/dead-letter")
-            assert response.status_code == 200
-            assert response.json()["dead_letter"] == []
-
-            worker_resp = await client.get("/tasks/workers")
-            assert worker_resp.status_code == 200
-            assert worker_resp.json() == {"workers": {}}
-        finally:
-            await client.cleanup()
+        response = await client.get("/tasks/dead-letter")
+        assert response.status_code == 200
+        assert response.json()["dead_letter"] == []
+
+        worker_resp = await client.get("/tasks/workers")
+        assert worker_resp.status_code == 200
+        assert worker_resp.json() == {"workers": {}}
+        await client.cleanup()
     finally:
         app.dependency_overrides.clear()
 
@@ -116,48 +114,52 @@ async def test_tasks_admin_endpoints_surface_queue_data(monkeypatch):
 
     try:
         client = await _create_client()
+        original_backend = settings.task_queue_backend
+        stub = StubQueue()
+        original_get_queue = app_module.get_task_queue
+        original_queue = app_module._task_queue
+
+        async def fake_get_queue():
+            return stub
+
+        monkeypatch.setattr(app_module, "get_task_queue", fake_get_queue)
+        settings.task_queue_backend = "memory"
         try:
-            original_backend = settings.task_queue_backend
-            original_queue = app_module._task_queue
-            stub = StubQueue()
-            app_module._task_queue = stub
-            settings.task_queue_backend = "memory"
-            try:
-                dead_resp = await client.get(
-                    "/tasks/dead-letter", params={"limit": 5, "offset": 0, "workflow_id": "plan-123"}
-                )
-                assert dead_resp.status_code == 200
-                assert dead_resp.json()["dead_letter"][0]["task_id"] == "dead-1"
-                assert stub.limit == 5
-                assert stub.workflow_id == "plan-123"
-
-                del_resp = await client.delete("/tasks/dead-letter/dead-1")
-                assert del_resp.status_code == 200
-                assert stub.deleted == "dead-1"
-
-                detail_resp = await client.get("/tasks/dead-letter/dead-1")
-                assert detail_resp.status_code == 200
-                assert detail_resp.json()["task_id"] == "dead-1"
-
-                worker_resp = await client.get("/tasks/workers")
-                assert worker_resp.status_code == 200
-                assert worker_resp.json() == {"workers": {"worker:1": {"status": "ok"}}}
-
-                requeue_resp = await client.post("/tasks/dead-letter/dead-1/requeue")
-                assert requeue_resp.status_code == 200
-                assert stub.requeued == "dead-1"
-                assert requeue_resp.json()["status"] == "requeued"
-
-                purge_resp = await client.delete("/tasks/dead-letter")
-                assert purge_resp.status_code == 200
-                assert purge_resp.json()["deleted"] == 1
-
-                purge_resp_age = await client.delete("/tasks/dead-letter", params={"older_than_minutes": 10})
-                assert purge_resp_age.status_code == 200
-            finally:
-                settings.task_queue_backend = original_backend
-                app_module._task_queue = original_queue
+            dead_resp = await client.get(
+                "/tasks/dead-letter", params={"limit": 5, "offset": 0, "workflow_id": "plan-123"}
+            )
+            assert dead_resp.status_code == 200
+            assert dead_resp.json()["dead_letter"][0]["task_id"] == "dead-1"
+            assert stub.limit == 5
+            assert stub.workflow_id == "plan-123"
+
+            del_resp = await client.delete("/tasks/dead-letter/dead-1")
+            assert del_resp.status_code == 200
+            assert stub.deleted == "dead-1"
+
+            detail_resp = await client.get("/tasks/dead-letter/dead-1")
+            assert detail_resp.status_code == 200
+            assert detail_resp.json()["task_id"] == "dead-1"
+
+            worker_resp = await client.get("/tasks/workers")
+            assert worker_resp.status_code == 200
+            assert worker_resp.json() == {"workers": {"worker:1": {"status": "ok"}}}
+
+            requeue_resp = await client.post("/tasks/dead-letter/dead-1/requeue")
+            assert requeue_resp.status_code == 200
+            assert stub.requeued == "dead-1"
+            assert requeue_resp.json()["status"] == "requeued"
+
+            purge_resp = await client.delete("/tasks/dead-letter")
+            assert purge_resp.status_code == 200
+            assert purge_resp.json()["deleted"] == 1
+
+            purge_resp_age = await client.delete("/tasks/dead-letter", params={"older_than_minutes": 10})
+            assert purge_resp_age.status_code == 200
         finally:
+            settings.task_queue_backend = original_backend
+            app_module._task_queue = original_queue
+            monkeypatch.setattr(app_module, "get_task_queue", original_get_queue)
             await client.cleanup()
     finally:
         app.dependency_overrides.clear()