From 09c9671fcb58722b83949018ae0e275f879cdd88 Mon Sep 17 00:00:00 2001
From: Harivansh Rathi <rathiharivansh@gmail.com>
Date: Wed, 8 Apr 2026 14:02:14 -0400
Subject: [PATCH] fix: supervise guest init and desktop services

---
 microagent-desktop-session.sh | 111 +++++++++++++++++++++++++++++++---
 microagent-init.sh            |  73 ++++++++++++++++++----
 2 files changed, 163 insertions(+), 21 deletions(-)

diff --git a/microagent-desktop-session.sh b/microagent-desktop-session.sh
index 8bc0537..992450c 100644
--- a/microagent-desktop-session.sh
+++ b/microagent-desktop-session.sh
@@ -1,23 +1,114 @@
 #!/usr/bin/env bash
-set -euo pipefail
+set -uo pipefail
 
 export DISPLAY=:0
 
-Xvfb "$DISPLAY" -screen 0 1280x800x24 >/tmp/xvfb.log 2>&1 &
-XVFB_PID=$!
+log() {
+  printf '[microagent-desktop] %s\n' "$*" >&2
+}
 
-for _ in $(seq 1 50); do
+pid_running() {
+  local pid="${1:-}"
+  [ -n "$pid" ] && kill -0 "$pid" >/dev/null 2>&1
+}
+
+reap_if_needed() {
+  local pid="${1:-}"
+  if [ -n "$pid" ]; then
+    wait "$pid" >/dev/null 2>&1 || true
+  fi
+}
+
+cleanup() {
+  trap - INT TERM
+  [ -n "${websockify_pid:-}" ] && kill "$websockify_pid" >/dev/null 2>&1 || true
+  [ -n "${x11vnc_pid:-}" ] && kill "$x11vnc_pid" >/dev/null 2>&1 || true
+  [ -n "${xterm_pid:-}" ] && kill "$xterm_pid" >/dev/null 2>&1 || true
+  [ -n "${openbox_pid:-}" ] && kill "$openbox_pid" >/dev/null 2>&1 || true
+  [ -n "${xvfb_pid:-}" ] && kill "$xvfb_pid" >/dev/null 2>&1 || true
+  wait >/dev/null 2>&1 || true
+  exit 0
+}
+
+start_openbox() {
+  reap_if_needed "${openbox_pid:-}"
+  log "starting openbox"
+  runuser -u node -- env DISPLAY="$DISPLAY" openbox >>/tmp/openbox.log 2>&1 &
+  openbox_pid=$!
+}
+
+start_xterm() {
+  reap_if_needed "${xterm_pid:-}"
+  log "starting xterm"
+  runuser -u node -- env DISPLAY="$DISPLAY" xterm -fa Monospace -fs 12 >>/tmp/xterm.log 2>&1 &
+  xterm_pid=$!
+}
+
+start_x11vnc() {
+  reap_if_needed "${x11vnc_pid:-}"
+  log "starting x11vnc"
+  x11vnc -display "$DISPLAY" -rfbport 5900 -forever -shared -nopw >>/tmp/x11vnc.log 2>&1 &
+  x11vnc_pid=$!
+}
+
+start_websockify() {
+  reap_if_needed "${websockify_pid:-}"
+  log "starting websockify on 6080"
+  websockify --web=/usr/share/novnc 6080 localhost:5900 >>/tmp/websockify.log 2>&1 &
+  websockify_pid=$!
+}
+
+trap cleanup INT TERM
+
+log "starting Xvfb"
+Xvfb "$DISPLAY" -screen 0 1280x800x24 >/tmp/xvfb.log 2>&1 &
+xvfb_pid=$!
+
+ready=0
+for _ in $(seq 1 100); do
   if xdpyinfo -display "$DISPLAY" >/dev/null 2>&1; then
+    ready=1
     break
   fi
+  if ! pid_running "$xvfb_pid"; then
+    log "Xvfb exited before display became ready"
+    wait "$xvfb_pid" >/dev/null 2>&1 || true
+    exit 1
+  fi
   sleep 0.1
 done
 
-runuser -u node -- env DISPLAY="$DISPLAY" openbox >/tmp/openbox.log 2>&1 &
-runuser -u node -- env DISPLAY="$DISPLAY" xterm -fa Monospace -fs 12 >/tmp/xterm.log 2>&1 &
+if [ "$ready" -ne 1 ]; then
+  log "Xvfb did not become ready in time"
+  exit 1
+fi
 
-x11vnc -display "$DISPLAY" -rfbport 5900 -forever -shared -nopw >/tmp/x11vnc.log 2>&1 &
-websockify --web=/usr/share/novnc 6080 localhost:5900 >/tmp/websockify.log 2>&1 &
+start_openbox
+start_xterm
+start_x11vnc
+start_websockify
 
-trap 'kill $XVFB_PID || true; kill 0 || true; exit 0' INT TERM
-wait -n
+while true; do
+  if ! pid_running "$xvfb_pid"; then
+    log "Xvfb exited; stopping desktop session"
+    wait "$xvfb_pid" >/dev/null 2>&1 || true
+    exit 1
+  fi
+  if ! pid_running "${openbox_pid:-}"; then
+    log "openbox exited; restarting"
+    start_openbox
+  fi
+  if ! pid_running "${xterm_pid:-}"; then
+    log "xterm exited; restarting"
+    start_xterm
+  fi
+  if ! pid_running "${x11vnc_pid:-}"; then
+    log "x11vnc exited; restarting"
+    start_x11vnc
+  fi
+  if ! pid_running "${websockify_pid:-}"; then
+    log "websockify exited; restarting"
+    start_websockify
+  fi
+  sleep 1
+done
diff --git a/microagent-init.sh b/microagent-init.sh
index 527322e..02ed89a 100644
--- a/microagent-init.sh
+++ b/microagent-init.sh
@@ -1,5 +1,11 @@
 #!/usr/bin/env bash
-set -euo pipefail
+set -uo pipefail
+
+PATH=/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin
+
+log() {
+  printf '[microagent-init] %s\n' "$*" >&2
+}
 
 mountpoint -q /proc || mount -t proc proc /proc
 mountpoint -q /sys || mount -t sysfs sysfs /sys
@@ -9,32 +15,77 @@ mountpoint -q /run || mount -t tmpfs tmpfs /run
 mkdir -p /tmp /var/tmp /run/sshd /var/log
 chmod 1777 /tmp /var/tmp
 
-# PID 1 is the guest supervisor: it brings up the minimum runtime surface and
-# stays alive to keep the VM services attached to a single lifecycle root.
+cleanup() {
+  trap - INT TERM
+  [ -n "${rng_pid:-}" ] && kill "$rng_pid" >/dev/null 2>&1 || true
+  [ -n "${sshd_pid:-}" ] && kill "$sshd_pid" >/dev/null 2>&1 || true
+  [ -n "${desktop_pid:-}" ] && kill "$desktop_pid" >/dev/null 2>&1 || true
+  wait >/dev/null 2>&1 || true
+  exit 0
+}
+
+pid_running() {
+  local pid="${1:-}"
+  [ -n "$pid" ] && kill -0 "$pid" >/dev/null 2>&1
+}
+
+reap_if_needed() {
+  local pid="${1:-}"
+  if [ -n "$pid" ]; then
+    wait "$pid" >/dev/null 2>&1 || true
+  fi
+}
+
+start_sshd() {
+  reap_if_needed "${sshd_pid:-}"
+  log "starting sshd on 2222"
+  /usr/sbin/sshd -D -e >>/var/log/sshd.log 2>&1 &
+  sshd_pid=$!
+}
+
+start_desktop() {
+  reap_if_needed "${desktop_pid:-}"
+  log "starting noVNC desktop on 6080"
+  /usr/local/bin/microagent-desktop-session >>/var/log/desktop.log 2>&1 &
+  desktop_pid=$!
+}
+
+trap cleanup INT TERM
+
+log "bringing up guest network"
 if ! /usr/local/bin/microagent-network-up >/var/log/network.log 2>&1; then
   cat /var/log/network.log >&2 || true
   exit 1
 fi
 
 if [ ! -f /etc/ssh/ssh_host_ed25519_key ]; then
+  log "generating ssh host keys"
   ssh-keygen -A
 fi
 
 if [ -f /etc/microagent/authorized_keys ]; then
+  log "installing baked authorized_keys for node"
   install -d -m 0700 -o node -g node /home/node/.ssh
   install -m 0600 -o node -g node /etc/microagent/authorized_keys /home/node/.ssh/authorized_keys
 fi
 
 if command -v jitterentropy-rngd >/dev/null 2>&1; then
+  log "starting jitterentropy-rngd"
   jitterentropy-rngd -v >/var/log/jitterentropy.log 2>&1 &
+  rng_pid=$!
 fi
 
-/usr/sbin/sshd -D -e >/var/log/sshd.log 2>&1 &
-/usr/local/bin/microagent-desktop-session >/var/log/desktop.log 2>&1 &
+start_sshd
+start_desktop
 
-trap 'kill 0 || true; exit 0' INT TERM
-wait -n
-status=$?
-kill 0 || true
-wait || true
-exit "$status"
+while true; do
+  if ! pid_running "${sshd_pid:-}"; then
+    log "sshd exited; restarting"
+    start_sshd
+  fi
+  if ! pid_running "${desktop_pid:-}"; then
+    log "desktop session exited; restarting"
+    start_desktop
+  fi
+  sleep 1
+done