orchestra/cmd/orchestra-worker/main.go

// orchestra-worker consumes router-issued leases for one local herdr. Homesrv
// remains the scheduler and CAS authority; this process owns only local Git
// and pane operations.
package main

import (
	"context"
	"encoding/json"
	"fmt"
	"log"
	"orchestra/internal/continuity"
	"orchestra/internal/domain"
	"orchestra/internal/federation"
	"orchestra/internal/herdr"
	"orchestra/internal/orchestrator"
	"os"
	"os/exec"
	"os/signal"
	"path/filepath"
	"strconv"
	"strings"
	"syscall"
	"time"
)

type worker struct {
	api                                    federation.Client
	herdr                                  *herdr.Client
	harnessID, harness, repo, root, remote string
	cursor                                 uint64
	tasks                                  map[string]domain.Task
	sessions                               map[string]herdr.Session
	leases                                 map[string]lease
	statePath                              string
	hard                                   float64
	registration                           federation.Worker
}
type lease struct {
	HandoffRef string `json:"handoff_ref,omitempty"`
}
type workerState struct {
	Cursor   uint64                   `json:"cursor"`
	Sessions map[string]herdr.Session `json:"sessions"`
	Tasks    map[string]domain.Task   `json:"tasks"`
	Leases   map[string]lease         `json:"leases"`
}

func (w *worker) load() {
	b, e := os.ReadFile(w.statePath)
	if e == nil {
		var s workerState
		if json.Unmarshal(b, &s) == nil {
			w.cursor = s.Cursor
			w.sessions = s.Sessions
			w.tasks = s.Tasks
			w.leases = s.Leases
		}
	}
	if w.sessions == nil {
		w.sessions = map[string]herdr.Session{}
	}
	if w.tasks == nil {
		w.tasks = map[string]domain.Task{}
	}
	if w.leases == nil {
		w.leases = map[string]lease{}
	}
}
func (w *worker) save() error {
	b, e := json.Marshal(workerState{Cursor: w.cursor, Sessions: w.sessions, Tasks: w.tasks, Leases: w.leases})
	if e != nil {
		return e
	}
	if e := os.MkdirAll(filepath.Dir(w.statePath), 0700); e != nil {
		return e
	}
	return os.WriteFile(w.statePath, b, 0600)
}

type artifactCAS struct{ api federation.Client }

func (c artifactCAS) PutArtifact(b []byte) (string, error) {
	return c.api.PutArtifact(context.Background(), b)
}
func (c artifactCAS) Artifact(ref string) ([]byte, error) {
	return c.api.Artifact(context.Background(), ref)
}

func created(e domain.Event) (domain.Task, bool) {
	if e.Type != "TaskCreated" {
		return domain.Task{}, false
	}
	var p struct {
		Source      string   `json:"source"`
		ExternalID  string   `json:"external_id"`
		Project     string   `json:"project"`
		Capability  []string `json:"capability"`
		Title       string   `json:"title"`
		Description string   `json:"description"`
	}
	if json.Unmarshal(e.Payload, &p) != nil || p.Source == "" || p.ExternalID == "" || p.Project == "" {
		return domain.Task{}, false
	}
	return domain.Task{ID: e.TaskID, Source: p.Source, ExternalID: p.ExternalID, Project: p.Project, Capability: p.Capability, Title: p.Title, Description: p.Description}, true
}

func (w *worker) syncBase(ctx context.Context) error {
	if out, err := exec.CommandContext(ctx, "git", "-C", w.repo, "fetch", w.remote, "--prune").CombinedOutput(); err != nil {
		return fmt.Errorf("fetch base checkout: %s: %w", out, err)
	}
	branch, err := exec.CommandContext(ctx, "git", "-C", w.repo, "symbolic-ref", "--quiet", "--short", "HEAD").Output()
	if err != nil {
		return fmt.Errorf("identify base branch: %w", err)
	}
	branchName := strings.TrimSpace(string(branch))
	if out, err := exec.CommandContext(ctx, "git", "-C", w.repo, "merge", "--ff-only", w.remote+"/"+branchName).CombinedOutput(); err != nil {
		return fmt.Errorf("fast-forward base checkout: %s: %w", out, err)
	}
	return nil
}

func (w *worker) start(ctx context.Context, t domain.Task, ref string) error {
	var wt string
	var h continuity.Handoff
	var err error
	// Synchronize the local base before any worktree operation. A worker never
	// treats a coordinator-side path as truth; the Git remote is the only
	// cross-machine transport.
	if err := w.syncBase(ctx); err != nil {
		return err
	}
	if ref != "" {
		b, err := w.api.Artifact(ctx, ref)
		if err != nil {
			return err
		}
		h, err = continuity.Decode(b)
		if err != nil {
			return err
		}
		if out, err := exec.CommandContext(ctx, "git", "-C", w.repo, "fetch", w.remote, "--prune").CombinedOutput(); err != nil {
			return fmt.Errorf("fetch pickup anchor: %s: %w", out, err)
		}
		wt = filepath.Join(w.root, t.ID)
		if _, err := os.Stat(wt); os.IsNotExist(err) {
			if out, err := exec.CommandContext(ctx, "git", "-C", w.repo, "worktree", "add", "-b", "orchestra/"+t.ID, wt, h.Anchor.GitSHA).CombinedOutput(); err != nil {
				return fmt.Errorf("create pickup worktree: %s: %w", out, err)
			}
		}
		if err = continuity.ValidatePickup(wt, h, taskHash(t)); err != nil {
			return err
		}
	} else {
		wt, err = (orchestrator.GitWorktrees{Repo: w.repo, Root: w.root}).Create(ctx, t)
		if err != nil {
			return err
		}
	}
	if _, err = w.herdr.Worktree(ctx, w.repo, wt, "orchestra/"+t.ID); err != nil {
		return err
	}
	s, err := w.herdr.StartAgent(ctx, wt, wt, "orchestra/"+t.ID, w.harness, t.ID)
	if err != nil {
		return err
	}
	p := "Begin Orchestra task " + t.ID + ".\nTitle: " + t.Title + "\nInstructions:\n" + t.Description + "\nWork only in this worktree. Do not edit TASK.md."
	if ref != "" {
		p += "\nA validated handoff exists. Read TASK.md and inspect local git history before continuing."
	}
	w.sessions[t.ID] = s
	if err := w.save(); err != nil {
		return err
	}
	// A prompt response can be lost after herdr accepted it. Persist the
	// session first so the worker can reconcile/release it after restart.
	return w.herdr.Prompt(ctx, s.PaneID, p, 0)
}

func taskHash(t domain.Task) string { b := continuity.RenderTaskFile(t); return domain.Hash(b) }

func (w *worker) releaseReady(ctx context.Context) {
	for id, s := range w.sessions {
		if report, err := os.ReadFile(filepath.Join(s.Worktree, ".orchestra-report.md")); err == nil && len(report) > 0 {
			if ref, err := w.api.PutArtifact(ctx, report); err == nil {
				if err = w.api.Complete(ctx, id, ref); err == nil {
					delete(w.sessions, id)
					delete(w.leases, id)
					_ = w.save()
					continue
				}
				log.Printf("complete %s: %v", id, err)
			} else {
				log.Printf("upload completion %s: %v", id, err)
			}
		}
		if _, err := os.Stat(filepath.Join(s.Worktree, herdr.HandoffReportFile)); err != nil {
			if !s.HandoffRequested {
				a := herdr.CLIAdapter{Client: w.herdr, Harness: w.harness}
				if occ, occErr := a.Occupancy(s); occErr == nil && occ >= w.hard {
					if boundary, boundaryErr := a.AtTurnBoundary(ctx, s); boundaryErr == nil && boundary {
						if err := a.RequestHandoff(ctx, s); err == nil {
							s.HandoffRequested, s.HandoffReason = true, "threshold"
							w.sessions[id] = s
							_ = w.save()
						}
					}
				}
			}
			continue
		}
		a := herdr.CLIAdapter{Client: w.herdr, Harness: w.harness, CAS: artifactCAS{w.api}, Remote: w.remote}
		ref, err := a.Release(ctx, s)
		if err != nil {
			log.Printf("release %s: %v", id, err)
			continue
		}
		sha, err := herdr.HeadSHA(s.Worktree)
		if err != nil {
			log.Printf("release %s anchor: %v", id, err)
			continue
		}
		if err = w.api.Release(ctx, id, ref, sha); err != nil {
			log.Printf("publish release %s: %v", id, err)
			continue
		}
		// release_agent only removes herdr's binding. Closing the released pane
		// after the handoff is durable prevents the next StartAgent from
		// inheriting the predecessor's still-running terminal process.
		if err := a.Kill(ctx, s); err != nil {
			log.Printf("close released pane %s: %v", id, err)
		}
		delete(w.sessions, id)
		_ = w.save()
	}
}

func (w *worker) once(ctx context.Context) error {
	es, _, err := w.api.Events(ctx, w.cursor)
	if err != nil {
		return err
	}
	for _, e := range es {
		if t, ok := created(e); ok {
			w.tasks[t.ID] = t
		}
		if e.Type == "TaskLeased" {
			var p struct {
				HarnessID  string `json:"harness_id"`
				HandoffRef string `json:"handoff_ref"`
			}
			if json.Unmarshal(e.Payload, &p) == nil && p.HarnessID == w.harnessID {
				w.leases[e.TaskID] = lease{HandoffRef: p.HandoffRef}
			}
		}
		if e.Type == "TaskReleased" || e.Type == "TaskCompleted" || e.Type == "TaskFailed" || e.Type == "TaskBlocked" {
			delete(w.leases, e.TaskID)
			delete(w.sessions, e.TaskID)
		}
		if e.Seq > w.cursor {
			w.cursor = e.Seq
		}
	}
	// State is only a cache. If a lease survived but its TaskCreated event is
	// older than the worker's cursor (or the event has been compacted), hydrate
	// the authoritative task projection before deciding whether to start.
	for taskID := range w.leases {
		if _, ok := w.tasks[taskID]; !ok {
			tasks, err := w.api.Tasks(ctx)
			if err != nil {
				return fmt.Errorf("hydrate leased task %s: %w", taskID, err)
			}
			for _, task := range tasks {
				w.tasks[task.ID] = task
			}
			break
		}
	}
	// Project the whole batch before launching. This prevents a new worker
	// from resurrecting every historical lease during its initial replay.
	for taskID, l := range w.leases {
		if _, started := w.sessions[taskID]; started {
			continue
		}
		if t, ok := w.tasks[taskID]; ok {
			if err := w.start(ctx, t, l.HandoffRef); err != nil {
				log.Printf("lease %s: %v", t.ID, err)
			}
		}
	}
	w.releaseReady(ctx)
	if err := w.save(); err != nil {
		return err
	}
	return w.api.Ack(ctx, w.cursor)
}

// reRegisterAfterCoordinatorRestart restores the coordinator's in-memory
// worker registry. A worker must survive a server restart without operator
// intervention; its persisted cursor and sessions remain valid.
func (w *worker) reRegisterAfterCoordinatorRestart(ctx context.Context, cause error) bool {
	if cause == nil || !strings.Contains(cause.Error(), "401 Unauthorized: unknown worker") {
		return false
	}
	if err := w.api.Register(ctx, w.registration); err != nil {
		log.Printf("re-register: %v", err)
		return false
	}
	log.Printf("re-registered after coordinator restart")
	return true
}
func required(k string) string {
	v := os.Getenv(k)
	if v == "" {
		log.Fatalf("%s required", k)
	}
	return v
}
func main() {
	id, token := required("ORCHESTRA_WORKER_ID"), required("ORCHESTRA_WORKER_TOKEN")
	hard := .75
	if v, err := strconv.ParseFloat(os.Getenv("ORCHESTRA_OCCUPANCY_HARD"), 64); err == nil && v > 0 && v < 1 {
		hard = v
	}
	w := &worker{api: federation.Client{BaseURL: required("ORCHESTRA_URL"), WorkerID: id, Token: token, AdmitToken: os.Getenv("ORCHESTRA_FEDERATION_ADMIT_TOKEN")}, harnessID: required("ORCHESTRA_WORKER_HERDR_ID"), harness: required("ORCHESTRA_WORKER_HARNESS"), repo: required("ORCHESTRA_REPO"), root: required("ORCHESTRA_WORKTREE_ROOT"), remote: required("ORCHESTRA_GIT_REMOTE"), tasks: map[string]domain.Task{}, sessions: map[string]herdr.Session{}, leases: map[string]lease{}, statePath: os.Getenv("ORCHESTRA_WORKER_STATE"), hard: hard, registration: federation.Worker{ID: id, Address: os.Getenv("ORCHESTRA_WORKER_ADDRESS"), Capacity: 1}}
	if w.statePath == "" {
		w.statePath = filepath.Join(w.root, ".orchestra-worker-state.json")
	}
	w.load()
	w.herdr = herdr.New(required("ORCHESTRA_WORKER_HERDR"))
	if id != w.harnessID {
		log.Fatal("ORCHESTRA_WORKER_ID must equal ORCHESTRA_WORKER_HERDR_ID so leases and offline recovery have one owner")
	}
	if err := w.api.Register(context.Background(), w.registration); err != nil {
		log.Fatal(err)
	}
	ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
	defer stop()
	ticker := time.NewTicker(5 * time.Second)
	defer ticker.Stop()
	for {
		if err := w.api.Heartbeat(ctx); err != nil {
			log.Printf("heartbeat: %v", err)
			if w.reRegisterAfterCoordinatorRestart(ctx, err) {
				continue
			}
		}
		if err := w.once(ctx); err != nil {
			log.Printf("poll: %v", err)
			w.reRegisterAfterCoordinatorRestart(ctx, err)
		}
		select {
		case <-ctx.Done():
			return
		case <-ticker.C:
		}
	}
}