added healthcheck for default web sockets

Author: mishushakov

template/e2b.toml

@@ -13,6 +13,7 @@
 team_id = "460355b3-4f64-48f9-9a16-4442817f79f5"
 memory_mb = 1_024
 start_cmd = "/root/.jupyter/start-up.sh"
+ready_cmd = "curl -s -o /dev/null -w '200' http://localhost:49999/health | grep -q '200'"
 dockerfile = "e2b.Dockerfile"
 template_name = "code-interpreter-v1"
 template_id = "nlhz8vlwyupq845jsdg9"

template/server/main.py

@@ -65,7 +65,17 @@ async def lifespan(app: FastAPI):
 
 @app.get("/health")
 async def get_health():
-    return "OK"
+    websockets_ready = True
+    for ws in default_websockets.values():
+        if not ws in websockets or not websockets[ws].is_connected():
+            websockets_ready = False
+            print(f"Context {ws} is not connected")
+            break
+
+    if not websockets_ready:
+        return PlainTextResponse("Not ready", status_code=503)
+
+    return PlainTextResponse("OK", status_code=200)
 
 
 @app.post("/execute")

template/server/messaging.py

@@ -542,6 +542,10 @@ async def _process_message(self, data: dict):
         else:
             logger.warning(f"[UNHANDLED MESSAGE TYPE]: {data['msg_type']}")
 
+    def is_connected(self) -> bool:
+        """Check if the websocket is connected and open"""
+        return self._ws is not None and self._ws.open
+
     async def close(self):
         logger.debug(f"Closing WebSocket {self.context_id}")