Merge pull request #358 from pollen-robotics/357-add-keep-alive-mechanism-to-allow-for-automatic-shutdown

Implement health check mecanism.

Author: FabienDanieau

src/reachy_mini/daemon/app/dashboard/static/js/health_check.js

@@ -0,0 +1,9 @@
+window.addEventListener("load", () => {
+    const healthCheckInterval = setInterval(async () => {
+        const response = await fetch("/health-check", { method: "POST" });
+        if (!response.ok) {
+            console.error("Health check failed");
+            clearInterval(healthCheckInterval);
+        }
+    }, 2500);
+});

src/reachy_mini/daemon/app/dashboard/templates/base.html

@@ -8,6 +8,7 @@
     <link rel="stylesheet" href="/static/css/app.css">
 
     {% block extra_js %}{% endblock %}
+    <script src="/static/js/health_check.js"></script>
 </head>
 
 

src/reachy_mini/daemon/app/main.py

@@ -8,6 +8,7 @@
 """
 
 import argparse
+import asyncio
 import logging
 from contextlib import asynccontextmanager
 from dataclasses import dataclass
@@ -44,6 +45,7 @@ class Args:
     check_collision: bool = False
 
     autostart: bool = True
+    timeout_health_check: float | None = None
 
     wake_up_on_start: bool = True
     goto_sleep_on_stop: bool = True
@@ -54,7 +56,7 @@ class Args:
     localhost_only: bool | None = None
 
 
-def create_app(args: Args) -> FastAPI:
+def create_app(args: Args, health_check_event: asyncio.Event | None = None) -> FastAPI:
     """Create and configure the FastAPI application."""
     localhost_only = (
         args.localhost_only
@@ -108,6 +110,14 @@ async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]:
 
     app.include_router(router)
 
+    if health_check_event is not None:
+
+        @app.post("/health-check")
+        async def health_check() -> dict[str, str]:
+            """Health check endpoint to reset the health check timer."""
+            health_check_event.set()
+            return {"status": "ok"}
+
     app.add_middleware(
         CORSMiddleware,
         allow_origins=["*"],  # or restrict to your HF domain
@@ -142,8 +152,36 @@ def run_app(args: Args) -> None:
     """Run the FastAPI app with Uvicorn."""
     logging.basicConfig(level=logging.INFO)
 
-    app = create_app(args)
-    uvicorn.run(app, host=args.fastapi_host, port=args.fastapi_port)
+    health_check_event = asyncio.Event()
+    app = create_app(args, health_check_event)
+
+    config = uvicorn.Config(app, host=args.fastapi_host, port=args.fastapi_port)
+    server = uvicorn.Server(config)
+
+    async def health_check_timeout(timeout_seconds: float) -> None:
+        while True:
+            try:
+                await asyncio.wait_for(
+                    health_check_event.wait(),
+                    timeout=timeout_seconds,
+                )
+                health_check_event.clear()
+            except asyncio.TimeoutError:
+                logging.warning("Health check timeout reached, stopping app.")
+                server.should_exit = True
+                break
+
+    loop = asyncio.get_event_loop()
+    if args.timeout_health_check is not None:
+        loop.create_task(health_check_timeout(args.timeout_health_check))
+
+    try:
+        loop.run_until_complete(server.serve())
+    except KeyboardInterrupt:
+        logging.info("Received Ctrl-C, shutting down gracefully.")
+    finally:
+        # Optional: additional cleanup here
+        pass
 
 
 def main() -> None:
@@ -198,6 +236,12 @@ def main() -> None:
         dest="autostart",
         help="Do not automatically start the daemon on launch (default: False).",
     )
+    parser.add_argument(
+        "--timeout-health-check",
+        type=float,
+        default=None,
+        help="Set the health check timeout in seconds (default: None).",
+    )
     parser.add_argument(
         "--wake-up-on-start",
         action="store_true",