Implement host registration, JWT refresh tokens, and multi-host scheduling

Replaces the hardcoded CP_HOST_AGENT_ADDR single-agent setup with a
DB-driven registration system supporting multiple host agents (BYOC).

Key changes:
- Host agents register via one-time token, receive a 7-day JWT + 60-day
  refresh token; heartbeat loop auto-refreshes on 401/403 and pauses all
  sandboxes if refresh fails
- HostClientPool: lazy Connect RPC client cache keyed by host ID, replacing
  the single static agent client throughout the API and service layers
- RoundRobinScheduler: picks an online host for each new sandbox via
  ListActiveHosts; extensible for future scheduling strategies
- HostMonitor (replaces Reconciler): passive heartbeat staleness check marks
  hosts unreachable and sandboxes missing after 90s; active reconciliation
  per online host restores missing-but-alive sandboxes and stops orphans
- Graceful host delete: returns 409 with affected sandbox list without
  ?force=true; force-delete destroys sandboxes then evicts pool client
- Snapshot delete broadcasts to all online hosts (templates have no host_id)
- sandbox.Manager.PauseAll: pauses all running VMs on CP connectivity loss
- New migration: host_refresh_tokens table with token rotation (issue-then-
  revoke ordering to prevent lockout on mid-rotation crash)
- New sandbox status 'missing' (reversible, unlike 'stopped') and host
  status 'unreachable'; both reflected in OpenAPI spec
- Fix: refresh token auth failure now returns 401 (was 400 via generic
  'invalid' substring match in serviceErrToHTTP)

cmd/control-plane/main.go

@@ -17,7 +17,8 @@ import (
 	"git.omukk.dev/wrenn/sandbox/internal/auth/oauth"
 	"git.omukk.dev/wrenn/sandbox/internal/config"
 	"git.omukk.dev/wrenn/sandbox/internal/db"
-	"git.omukk.dev/wrenn/sandbox/proto/hostagent/gen/hostagentv1connect"
+	"git.omukk.dev/wrenn/sandbox/internal/lifecycle"
+	"git.omukk.dev/wrenn/sandbox/internal/scheduler"
 )
 
 func main() {
@@ -66,12 +67,11 @@ func main() {
 	}
 	slog.Info("connected to redis")
 
-	// Connect RPC client for the host agent.
-	agentHTTP := &http.Client{Timeout: 10 * time.Minute}
-	agentClient := hostagentv1connect.NewHostAgentServiceClient(
-		agentHTTP,
-		cfg.HostAgentAddr,
-	)
+	// Host client pool — manages Connect RPC clients to host agents.
+	hostPool := lifecycle.NewHostClientPool()
+
+	// Scheduler — picks a host for each new sandbox (round-robin for now).
+	hostScheduler := scheduler.NewRoundRobinScheduler(queries)
 
 	// OAuth provider registry.
 	oauthRegistry := oauth.NewRegistry()
@@ -87,11 +87,11 @@ func main() {
 	}
 
 	// API server.
-	srv := api.New(queries, agentClient, pool, rdb, []byte(cfg.JWTSecret), oauthRegistry, cfg.OAuthRedirectURL)
+	srv := api.New(queries, hostPool, hostScheduler, pool, rdb, []byte(cfg.JWTSecret), oauthRegistry, cfg.OAuthRedirectURL)
 
-	// Start reconciler.
-	reconciler := api.NewReconciler(queries, agentClient, "default", 5*time.Second)
-	reconciler.Start(ctx)
+	// Start host monitor (passive + active reconciliation every 30s).
+	monitor := api.NewHostMonitor(queries, hostPool, 30*time.Second)
+	monitor.Start(ctx)
 
 	httpServer := &http.Server{
 		Addr:    cfg.ListenAddr,
@@ -114,7 +114,7 @@ func main() {
 		}
 	}()
 
-	slog.Info("control plane starting", "addr", cfg.ListenAddr, "agent", cfg.HostAgentAddr)
+	slog.Info("control plane starting", "addr", cfg.ListenAddr)
 	if err := httpServer.ListenAndServe(); err != nil && err != http.ErrServerClosed {
 		slog.Error("http server error", "error", err)
 		os.Exit(1)

cmd/host-agent/main.go

@@ -18,7 +18,7 @@ import (
 )
 
 func main() {
-	registrationToken := flag.String("register", "", "One-time registration token from the control plane")
+	registrationToken := flag.String("register", "", "One-time registration token from the control plane (required on first run)")
 	advertiseAddr := flag.String("address", "", "Externally-reachable address (ip:port) for this host agent")
 	flag.Parse()
 
@@ -42,7 +42,16 @@ func main() {
 	listenAddr := envOrDefault("AGENT_LISTEN_ADDR", ":50051")
 	rootDir := envOrDefault("AGENT_FILES_ROOTDIR", "/var/lib/wrenn")
 	cpURL := os.Getenv("AGENT_CP_URL")
-	tokenFile := filepath.Join(rootDir, "host-token")
+	tokenFile := filepath.Join(rootDir, "host.jwt")
+
+	if cpURL == "" {
+		slog.Error("AGENT_CP_URL environment variable is required")
+		os.Exit(1)
+	}
+	if *advertiseAddr == "" {
+		slog.Error("--address flag is required (externally-reachable ip:port)")
+		os.Exit(1)
+	}
 
 	cfg := sandbox.Config{
 		KernelPath:   filepath.Join(rootDir, "kernels", "vmlinux"),
@@ -58,34 +67,34 @@ func main() {
 
 	mgr.StartTTLReaper(ctx)
 
-	if *advertiseAddr == "" {
-		slog.Error("--address flag is required (externally-reachable ip:port)")
+	// Register with the control plane and start heartbeating.
+	hostToken, err := hostagent.Register(ctx, hostagent.RegistrationConfig{
+		CPURL:             cpURL,
+		RegistrationToken: *registrationToken,
+		TokenFile:         tokenFile,
+		Address:           *advertiseAddr,
+	})
+	if err != nil {
+		slog.Error("host registration failed", "error", err)
 		os.Exit(1)
 	}
 
-	// Register with the control plane (if configured).
-	if cpURL != "" {
-		hostToken, err := hostagent.Register(ctx, hostagent.RegistrationConfig{
-			CPURL:             cpURL,
-			RegistrationToken: *registrationToken,
-			TokenFile:         tokenFile,
-			Address:           *advertiseAddr,
-		})
-		if err != nil {
-			slog.Error("host registration failed", "error", err)
-			os.Exit(1)
-		}
-
-		hostID, err := hostagent.HostIDFromToken(hostToken)
-		if err != nil {
-			slog.Error("failed to extract host ID from token", "error", err)
-			os.Exit(1)
-		}
-
-		slog.Info("host registered", "host_id", hostID)
-		hostagent.StartHeartbeat(ctx, cpURL, hostID, hostToken, 30*time.Second)
+	hostID, err := hostagent.HostIDFromToken(hostToken)
+	if err != nil {
+		slog.Error("failed to extract host ID from token", "error", err)
+		os.Exit(1)
 	}
 
+	slog.Info("host registered", "host_id", hostID)
+
+	// Start heartbeat loop. On CP rejection: try JWT refresh. If that fails,
+	// pause all running sandboxes to ensure they're not left orphaned.
+	hostagent.StartHeartbeat(ctx, cpURL, tokenFile, hostID, 30*time.Second, func() {
+		pauseCtx, cancel := context.WithTimeout(context.Background(), 2*time.Minute)
+		defer cancel()
+		mgr.PauseAll(pauseCtx)
+	})
+
 	srv := hostagent.NewServer(mgr)
 	path, handler := hostagentv1connect.NewHostAgentServiceHandler(srv)
 
@@ -115,7 +124,7 @@ func main() {
 		}
 	}()
 
-	slog.Info("host agent starting", "addr", listenAddr)
+	slog.Info("host agent starting", "addr", listenAddr, "host_id", hostID)
 	if err := httpServer.ListenAndServe(); err != nil && err != http.ErrServerClosed {
 		slog.Error("http server error", "error", err)
 		os.Exit(1)