diff --git a/services/mcp-docugen/src/mcp_docugen/main.py b/services/mcp-docugen/src/mcp_docugen/main.py index 04de7e3..deff841 100644 --- a/services/mcp-docugen/src/mcp_docugen/main.py +++ b/services/mcp-docugen/src/mcp_docugen/main.py @@ -49,24 +49,27 @@ async def build_app(settings: Settings | None = None) -> FastAPI: ) mcp = build_mcp_server(template_store, renderer) + mcp_asgi = mcp.streamable_http_app() @asynccontextmanager async def lifespan(app: FastAPI): - cleanup_task = asyncio.create_task( - _periodic_cleanup(generation_store, interval_seconds=24 * 3600) - ) - try: - yield - finally: - cleanup_task.cancel() + async with mcp.session_manager.run(): + cleanup_task = asyncio.create_task( + _periodic_cleanup(generation_store, interval_seconds=24 * 3600) + ) try: - await cleanup_task - except asyncio.CancelledError: - pass + yield + finally: + cleanup_task.cancel() + try: + await cleanup_task + except asyncio.CancelledError: + pass app = build_http_app(template_store, generation_store) app.router.lifespan_context = lifespan - app.mount("/mcp", mcp.streamable_http_app()) + # streamable_http_app() espone /mcp internamente; monto a root. + app.mount("/", mcp_asgi) app.add_middleware( ApiKeyAuthMiddleware,