diff --git a/src/main.py b/src/main.py index dcb67ed..99346d8 100644 --- a/src/main.py +++ b/src/main.py @@ -77,6 +77,13 @@ async def stats() -> dict: return s.model_dump(mode="json") +# Initialize the MCP sub-app early so session_manager is available for lifespan. +# streamable_http_app() creates a sub-app with a route at /mcp internally. +# Mounted at "/" (last in the route list), FastAPI routes take priority and +# requests to /mcp fall through to the sub-app — keeping the public URL at /mcp. +_mcp_app = mcp.streamable_http_app() + + # --------------------------------------------------------------------------- # FastAPI app # --------------------------------------------------------------------------- @@ -84,7 +91,8 @@ async def stats() -> dict: async def lifespan(app: FastAPI): logger.info("OpenBrain starting up") await db.get_pool() # Warm the connection pool - yield + async with mcp.session_manager.run(): + yield await db.close_pool() logger.info("OpenBrain shut down") @@ -96,9 +104,6 @@ app = FastAPI( lifespan=lifespan, ) -# Mount MCP server at /mcp (HTTP streamable transport) -app.mount("/mcp", mcp.streamable_http_app()) - # --------------------------------------------------------------------------- # REST endpoints (for direct API access and health checks) @@ -126,3 +131,8 @@ async def api_recent(limit: int = 20, _: str = Depends(require_api_key)) -> list @app.get("/v1/stats", response_model=Stats) async def api_stats(_: str = Depends(require_api_key)) -> Stats: return await brain.stats() + + +# Mount MCP sub-app at "/" last — FastAPI routes above take priority, +# /mcp falls through to the sub-app's internal /mcp route. +app.mount("/", _mcp_app)