Harden worker federation and operator UI

internal/federation/federation.go

@@ -2,8 +2,11 @@ package federation
 
 import (
 	"crypto/sha256"
+	"encoding/json"
 	"errors"
 	"fmt"
+	"os"
+	"path/filepath"
 	"sync"
 	"time"
 )
@@ -12,12 +15,26 @@ var ErrUnknownWorker = errors.New("unknown worker")
 var ErrUnauthorized = errors.New("worker authentication failed")
 
 type Worker struct {
-	ID       string    `json:"id"`
-	Address  string    `json:"address"`
-	Capacity int       `json:"capacity"`
-	LastSeen time.Time `json:"last_seen"`
-	Online   bool      `json:"online"`
-	Token    string    `json:"-"`
+	ID       string       `json:"id"`
+	Address  string       `json:"address"`
+	Capacity int          `json:"capacity"`
+	LastSeen time.Time    `json:"last_seen"`
+	Online   bool         `json:"online"`
+	Health   WorkerHealth `json:"health"`
+	Token    string       `json:"-"`
+}
+
+// WorkerHealth is reported by the worker that owns the local herdr socket.
+// It intentionally does not reuse coordinator TCP-probe state: a remote
+// socket is meaningful only from the machine where the worker and checkout
+// live.
+type WorkerHealth struct {
+	HerdrStatus string    `json:"herdr_status"` // reachable, unreachable, or unknown
+	CheckedAt   time.Time `json:"checked_at,omitempty"`
+	ActiveTask  string    `json:"active_task_id,omitempty"`
+	ActivePane  string    `json:"active_pane_id,omitempty"`
+	LastError   string    `json:"last_error,omitempty"`
+	ErrorAt     time.Time `json:"error_at,omitempty"`
 }
 
 // Capture is published by a worker that owns the pane. The coordinator never
@@ -53,6 +70,88 @@ type Registry struct {
 	cursors    map[string]uint64
 	captures   map[string]Capture // worker/task
 	commands   map[string][]Command
+	// StatePath preserves worker-owned pane captures and pending approval
+	// commands across coordinator restarts. A worker must still re-register to
+	// be online before it can read or act on recovered state.
+	StatePath string
+}
+
+type persistedState struct {
+	Captures map[string]Capture         `json:"captures"`
+	Commands map[string][]Command       `json:"commands"`
+	Workers  map[string]persistedWorker `json:"workers"`
+}
+
+// persistedWorker deliberately includes the per-worker token. The state file
+// is mode 0600, and retaining this binding prevents an arbitrary process from
+// registering a recovered worker ID and executing its pending approval.
+type persistedWorker struct {
+	ID       string `json:"id"`
+	Address  string `json:"address"`
+	Capacity int    `json:"capacity"`
+	Token    string `json:"token"`
+}
+
+// Load restores durable capture/command state. Call this before accepting
+// federation requests; an unreadable state file is unsafe because it could
+// otherwise make a pending approval silently disappear.
+func (r *Registry) Load() error {
+	r.mu.Lock()
+	defer r.mu.Unlock()
+	r.init()
+	if r.StatePath == "" {
+		return nil
+	}
+	b, err := os.ReadFile(r.StatePath)
+	if errors.Is(err, os.ErrNotExist) {
+		return nil
+	}
+	if err != nil {
+		return err
+	}
+	var state persistedState
+	if err := json.Unmarshal(b, &state); err != nil {
+		return fmt.Errorf("invalid federation state: %w", err)
+	}
+	if state.Captures != nil {
+		r.captures = state.Captures
+	}
+	if state.Commands != nil {
+		r.commands = state.Commands
+	}
+	for id, w := range state.Workers {
+		if id == "" || w.ID != id || w.Token == "" {
+			return fmt.Errorf("invalid federation worker %q", id)
+		}
+		r.workers[id] = Worker{ID: w.ID, Address: w.Address, Capacity: w.Capacity, Token: w.Token}
+	}
+	return nil
+}
+
+// persistLocked atomically replaces the state file. Callers hold r.mu.
+func (r *Registry) persistLocked() error {
+	if r.StatePath == "" {
+		return nil
+	}
+	workers := make(map[string]persistedWorker, len(r.workers))
+	for id, w := range r.workers {
+		workers[id] = persistedWorker{ID: w.ID, Address: w.Address, Capacity: w.Capacity, Token: w.Token}
+	}
+	b, err := json.Marshal(persistedState{Captures: r.captures, Commands: r.commands, Workers: workers})
+	if err != nil {
+		return err
+	}
+	if err := os.MkdirAll(filepath.Dir(r.StatePath), 0755); err != nil {
+		return err
+	}
+	tmp := r.StatePath + ".tmp"
+	if err := os.WriteFile(tmp, b, 0600); err != nil {
+		return err
+	}
+	if err := os.Rename(tmp, r.StatePath); err != nil {
+		return err
+	}
+	return os.Chmod(r.StatePath, 0600)
 }
 
 func (r *Registry) init() {
@@ -94,6 +193,9 @@ func (r *Registry) PutCapture(worker string, c Capture) (Capture, error) {
 	}
 	c.At = time.Now().UTC()
 	r.captures[k] = c
+	if err := r.persistLocked(); err != nil {
+		return Capture{}, fmt.Errorf("persist capture: %w", err)
+	}
 	return c, nil
 }
 func (r *Registry) Capture(worker, task string) (Capture, bool) {
@@ -118,6 +220,9 @@ func (r *Registry) Queue(worker string, c Command) (Command, error) {
 	c.Status = "pending"
 	r.commands[worker] = append(r.commands[worker], c)
 	r.pruneCommands(worker)
+	if err := r.persistLocked(); err != nil {
+		return Command{}, fmt.Errorf("persist command: %w", err)
+	}
 	return c, nil
 }
 
@@ -129,7 +234,7 @@ const CommandRetention = 30 * time.Minute
 // pruneCommands drops resolved commands past CommandRetention. B21: this list
 // was append-only, so resolved commands accumulated for the process lifetime
 // and every worker poll rescanned the entire history. Callers hold r.mu.
-func (r *Registry) pruneCommands(worker string) {
+func (r *Registry) pruneCommands(worker string) bool {
 	cutoff := time.Now().UTC().Add(-CommandRetention)
 	in := r.commands[worker]
 	out := in[:0]
@@ -139,10 +244,12 @@ func (r *Registry) pruneCommands(worker string) {
 		}
 	}
 	if len(out) == 0 {
+		changed := len(in) != 0
 		delete(r.commands, worker)
-		return
+		return changed
 	}
 	r.commands[worker] = out
+	return len(out) != len(in)
 }
 func (r *Registry) Commands(worker string) ([]Command, error) {
 	r.mu.Lock()
@@ -151,7 +258,11 @@ func (r *Registry) Commands(worker string) ([]Command, error) {
 	if _, ok := r.workers[worker]; !ok {
 		return nil, ErrUnknownWorker
 	}
-	r.pruneCommands(worker)
+	if r.pruneCommands(worker) {
+		if err := r.persistLocked(); err != nil {
+			return nil, fmt.Errorf("persist pruned commands: %w", err)
+		}
+	}
 	var out []Command
 	for _, c := range r.commands[worker] {
 		if c.Status == "pending" {
@@ -182,6 +293,9 @@ func (r *Registry) CompleteCommand(worker, id, status, message string) error {
 			}
 			r.commands[worker][i].Status = status
 			r.commands[worker][i].Error = message
+			if err := r.persistLocked(); err != nil {
+				return fmt.Errorf("persist command resolution: %w", err)
+			}
 			return nil
 		}
 	}
@@ -214,6 +328,9 @@ func (r *Registry) Register(w Worker, admitToken string) error {
 	if _, ok := r.cursors[w.ID]; !ok {
 		r.cursors[w.ID] = 0
 	}
+	if err := r.persistLocked(); err != nil {
+		return fmt.Errorf("persist worker registration: %w", err)
+	}
 	return nil
 }
 func (r *Registry) Authenticate(id, token string) error {
@@ -251,7 +368,7 @@ func (r *Registry) Ack(id string, cursor uint64) error {
 	r.cursors[id] = cursor
 	return nil
 }
-func (r *Registry) Heartbeat(id string) error {
+func (r *Registry) Heartbeat(id string, health ...WorkerHealth) error {
 	r.mu.Lock()
 	defer r.mu.Unlock()
 	r.init()
@@ -261,6 +378,9 @@ func (r *Registry) Heartbeat(id string) error {
 	}
 	w.LastSeen = time.Now().UTC()
 	w.Online = true
+	if len(health) > 0 {
+		w.Health = health[0]
+	}
 	r.workers[id] = w
 	return nil
 }