orchestra/internal/federation/federation.go

package federation

import (
	"crypto/sha256"
	"encoding/json"
	"errors"
	"fmt"
	"os"
	"path/filepath"
	"sync"
	"time"
)

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"`
	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
// reads a remote herdr socket; this is the worker-pulled counterpart.
type Capture struct {
	TaskID   string    `json:"task_id"`
	PaneID   string    `json:"pane_id"`
	Text     string    `json:"text"`
	Revision uint64    `json:"revision"`
	At       time.Time `json:"at"`
}
type Command struct {
	ID              string    `json:"id"`
	TaskID          string    `json:"task_id"`
	Kind            string    `json:"kind"`
	PaneID          string    `json:"pane_id"`
	CaptureRevision uint64    `json:"capture_revision"`
	CreatedAt       time.Time `json:"created_at"`
	Status          string    `json:"status"`
	Error           string    `json:"error,omitempty"`
}

type Registry struct {
	mu sync.Mutex
	// AdmitToken, if set, is a pre-shared secret every registration must
	// present (S10: registration previously accepted a self-declared id and
	// self-chosen token from any caller — admission-control-free). Leave
	// empty only for a deliberately open deployment.
	AdmitToken string
	workers    map[string]Worker
	TTL        time.Duration
	OnOffline  func(Worker)
	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() {
	if r.TTL <= 0 {
		r.TTL = 90 * time.Second
	}
	if r.workers == nil {
		r.workers = map[string]Worker{}
	}
	if r.cursors == nil {
		r.cursors = map[string]uint64{}
	}
	if r.captures == nil {
		r.captures = map[string]Capture{}
	}
	if r.commands == nil {
		r.commands = map[string][]Command{}
	}
}

func captureKey(worker, task string) string { return worker + "\x00" + task }
func (r *Registry) PutCapture(worker string, c Capture) (Capture, error) {
	r.mu.Lock()
	defer r.mu.Unlock()
	r.init()
	if _, ok := r.workers[worker]; !ok {
		return Capture{}, ErrUnknownWorker
	}
	if c.TaskID == "" || c.PaneID == "" {
		return Capture{}, errors.New("task_id and pane_id required")
	}
	k := captureKey(worker, c.TaskID)
	old := r.captures[k]
	if old.Text != c.Text || old.PaneID != c.PaneID {
		c.Revision = old.Revision + 1
	}
	if c.Revision == 0 {
		c.Revision = 1
	}
	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) {
	r.mu.Lock()
	defer r.mu.Unlock()
	r.init()
	c, ok := r.captures[captureKey(worker, task)]
	return c, ok
}
func (r *Registry) Queue(worker string, c Command) (Command, error) {
	r.mu.Lock()
	defer r.mu.Unlock()
	r.init()
	if _, ok := r.workers[worker]; !ok {
		return Command{}, ErrUnknownWorker
	}
	if c.TaskID == "" || c.PaneID == "" || c.CaptureRevision == 0 || (c.Kind != "grant_approval" && c.Kind != "deny_approval") {
		return Command{}, errors.New("invalid control command")
	}
	c.ID = fmt.Sprintf("cmd-%x", sha256.Sum256([]byte(fmt.Sprintf("%s/%s/%s/%d/%d", worker, c.TaskID, c.Kind, c.CaptureRevision, time.Now().UnixNano()))))[:20]
	c.CreatedAt = time.Now().UTC()
	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
}

// CommandRetention is how long a resolved command stays queryable so a worker
// that retries a completion, or an operator reading the UI, still sees its
// outcome. Pending commands are never pruned.
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) bool {
	cutoff := time.Now().UTC().Add(-CommandRetention)
	in := r.commands[worker]
	out := in[:0]
	for _, c := range in {
		if c.Status == "pending" || c.CreatedAt.After(cutoff) {
			out = append(out, c)
		}
	}
	if len(out) == 0 {
		changed := len(in) != 0
		delete(r.commands, worker)
		return changed
	}
	r.commands[worker] = out
	return len(out) != len(in)
}
func (r *Registry) Commands(worker string) ([]Command, error) {
	r.mu.Lock()
	defer r.mu.Unlock()
	r.init()
	if _, ok := r.workers[worker]; !ok {
		return nil, ErrUnknownWorker
	}
	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" {
			out = append(out, c)
		}
	}
	return out, nil
}
func (r *Registry) Command(worker, id string) (Command, bool) {
	r.mu.Lock()
	defer r.mu.Unlock()
	r.init()
	for _, c := range r.commands[worker] {
		if c.ID == id {
			return c, true
		}
	}
	return Command{}, false
}
func (r *Registry) CompleteCommand(worker, id, status, message string) error {
	r.mu.Lock()
	defer r.mu.Unlock()
	r.init()
	for i := range r.commands[worker] {
		if r.commands[worker][i].ID == id {
			if r.commands[worker][i].Status != "pending" {
				return errors.New("command already resolved")
			}
			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
		}
	}
	return errors.New("command not found")
}

// Register admits a worker. admitToken must match r.AdmitToken whenever one
// is configured. Re-registering an ID that's already claimed requires that
// worker's own current token, so a caller can't self-declare someone else's
// id and hijack an existing worker's identity/capacity.
func (r *Registry) Register(w Worker, admitToken string) error {
	if w.ID == "" {
		return errors.New("worker id required")
	}
	if w.Token == "" {
		return errors.New("worker token required")
	}
	r.mu.Lock()
	defer r.mu.Unlock()
	r.init()
	if r.AdmitToken != "" && admitToken != r.AdmitToken {
		return ErrUnauthorized
	}
	if existing, ok := r.workers[w.ID]; ok && existing.Token != w.Token {
		return ErrUnauthorized
	}
	w.LastSeen = time.Now().UTC()
	w.Online = true
	r.workers[w.ID] = w
	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 {
	r.mu.Lock()
	defer r.mu.Unlock()
	r.init()
	w, ok := r.workers[id]
	if !ok {
		return ErrUnknownWorker
	}
	if w.Token == "" || token == "" || w.Token != token {
		return ErrUnauthorized
	}
	return nil
}
func (r *Registry) Cursor(id string) (uint64, error) {
	r.mu.Lock()
	defer r.mu.Unlock()
	r.init()
	if _, ok := r.workers[id]; !ok {
		return 0, ErrUnknownWorker
	}
	return r.cursors[id], nil
}
func (r *Registry) Ack(id string, cursor uint64) error {
	r.mu.Lock()
	defer r.mu.Unlock()
	r.init()
	if _, ok := r.workers[id]; !ok {
		return ErrUnknownWorker
	}
	if cursor < r.cursors[id] {
		return errors.New("cursor moved backwards")
	}
	r.cursors[id] = cursor
	return nil
}
func (r *Registry) Heartbeat(id string, health ...WorkerHealth) error {
	r.mu.Lock()
	defer r.mu.Unlock()
	r.init()
	w, ok := r.workers[id]
	if !ok {
		return ErrUnknownWorker
	}
	w.LastSeen = time.Now().UTC()
	w.Online = true
	if len(health) > 0 {
		w.Health = health[0]
	}
	r.workers[id] = w
	return nil
}

// Available refreshes TTL state and reports whether a registered worker owns
// this harness id. Router admission uses it so a reachable TCP bridge alone
// can never make an offline worker eligible for a lease.
func (r *Registry) Available(id string) bool {
	r.mu.Lock()
	defer r.mu.Unlock()
	r.init()
	w, ok := r.workers[id]
	if !ok {
		return false
	}
	w.Online = time.Since(w.LastSeen) <= r.TTL
	r.workers[id] = w
	return w.Online
}
func (r *Registry) Snapshot() []Worker {
	r.mu.Lock()
	defer r.mu.Unlock()
	r.init()
	now := time.Now().UTC()
	out := make([]Worker, 0, len(r.workers))
	for id, w := range r.workers {
		wasOnline := w.Online
		w.Online = now.Sub(w.LastSeen) <= r.TTL
		r.workers[id] = w
		if wasOnline && !w.Online && r.OnOffline != nil {
			go r.OnOffline(w)
		}
		out = append(out, w)
	}
	return out
}