Add federation worker and canonical handoffs

2026-07-28 16:17:18 +04:00
parent 58793a5aa3
commit 2cecbc4015
22 changed files with 1429 additions and 108 deletions
@@ -0,0 +1,357 @@
+// 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:
+		}
+	}
+}
@@ -0,0 +1,151 @@
+package main
+
+import (
+	"bufio"
+	"context"
+	"encoding/json"
+	"errors"
+	"net"
+	"net/http"
+	"net/http/httptest"
+	"orchestra/internal/domain"
+	"orchestra/internal/federation"
+	"orchestra/internal/herdr"
+	"os"
+	"os/exec"
+	"path/filepath"
+	"testing"
+)
+
+func TestWorkerReregistersAfterCoordinatorForgetsIt(t *testing.T) {
+	registered := false
+	s := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if r.Method != http.MethodPost || r.URL.Path != "/v1/federation/workers" {
+			t.Fatalf("unexpected %s %s", r.Method, r.URL.Path)
+		}
+		registered = true
+		w.WriteHeader(http.StatusCreated)
+	}))
+	defer s.Close()
+	w := &worker{api: federation.Client{BaseURL: s.URL, WorkerID: "h", Token: "t"}, registration: federation.Worker{ID: "h", Capacity: 1}}
+	if !w.reRegisterAfterCoordinatorRestart(context.Background(), errors.New("federation: 401 Unauthorized: unknown worker")) {
+		t.Fatal("worker did not identify a coordinator restart")
+	}
+	if !registered {
+		t.Fatal("worker did not register")
+	}
+}
+
+func TestInitialReplayDoesNotResurrectReleasedLease(t *testing.T) {
+	s := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		switch r.URL.Path {
+		case "/v1/federation/events":
+			_, _ = w.Write([]byte(`{"cursor":0,"events":[{"seq":1,"id":"c","type":"TaskCreated","task_id":"t","version":1,"payload":{"source":"s","external_id":"x","project":"p"},"surface":"system"},{"seq":2,"id":"l","type":"TaskLeased","task_id":"t","version":2,"payload":{"harness_id":"h"},"surface":"system"},{"seq":3,"id":"r","type":"TaskReleased","task_id":"t","version":3,"payload":{"reason":"expired"},"surface":"system"}]}`))
+		case "/v1/federation/events/ack":
+			w.WriteHeader(http.StatusNoContent)
+		default:
+			t.Fatalf("unexpected %s", r.URL.Path)
+		}
+	}))
+	defer s.Close()
+	w := &worker{api: federation.Client{BaseURL: s.URL, WorkerID: "h", Token: "t"}, harnessID: "h", tasks: map[string]domain.Task{}, sessions: map[string]herdr.Session{}, leases: map[string]lease{}, statePath: t.TempDir() + "/state.json", hard: .75}
+	if err := w.once(context.Background()); err != nil {
+		t.Fatal(err)
+	}
+	if len(w.leases) != 0 || len(w.sessions) != 0 {
+		t.Fatalf("replayed lease survived: leases=%v sessions=%v", w.leases, w.sessions)
+	}
+	if w.cursor != 3 {
+		t.Fatalf("cursor=%d want 3", w.cursor)
+	}
+}
+
+func TestWorkerHydratesLeaseMissingTaskCache(t *testing.T) {
+	s := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		switch r.URL.Path {
+		case "/v1/federation/events":
+			_, _ = w.Write([]byte(`{"cursor":2,"events":[{"seq":2,"id":"l","type":"TaskLeased","task_id":"t","version":2,"payload":{"harness_id":"h"},"surface":"system"}]}`))
+		case "/v1/tasks":
+			_, _ = w.Write([]byte(`[{"id":"t","source":"s","external_id":"x","project":"p","state":"leased"}]`))
+		case "/v1/federation/events/ack":
+			w.WriteHeader(http.StatusNoContent)
+		default:
+			t.Fatalf("unexpected %s", r.URL.Path)
+		}
+	}))
+	defer s.Close()
+	w := &worker{api: federation.Client{BaseURL: s.URL, WorkerID: "h", Token: "t"}, harnessID: "h", cursor: 1, tasks: map[string]domain.Task{}, sessions: map[string]herdr.Session{"t": {}}, leases: map[string]lease{}, statePath: t.TempDir() + "/state.json", hard: .75}
+	if err := w.once(context.Background()); err != nil {
+		t.Fatal(err)
+	}
+	if _, ok := w.tasks["t"]; !ok || w.cursor != 2 {
+		t.Fatalf("task was not hydrated: tasks=%v cursor=%d", w.tasks, w.cursor)
+	}
+}
+
+func TestWorkerStartsRouterIssuedLeaseInLocalGitWorktree(t *testing.T) {
+	remote := filepath.Join(t.TempDir(), "remote.git")
+	if out, err := exec.Command("git", "init", "--bare", remote).CombinedOutput(); err != nil {
+		t.Fatalf("remote: %v %s", err, out)
+	}
+	seed := t.TempDir()
+	for _, a := range [][]string{{"init", seed}, {"-C", seed, "config", "user.email", "t@t"}, {"-C", seed, "config", "user.name", "t"}, {"-C", seed, "commit", "--allow-empty", "-m", "init"}, {"-C", seed, "remote", "add", "origin", remote}, {"-C", seed, "push", "-u", "origin", "HEAD:master"}} {
+		if out, err := exec.Command("git", a...).CombinedOutput(); err != nil {
+			t.Fatalf("git %v: %v %s", a, err, out)
+		}
+	}
+	repo := filepath.Join(t.TempDir(), "repo")
+	if out, err := exec.Command("git", "clone", remote, repo).CombinedOutput(); err != nil {
+		t.Fatal(string(out))
+	}
+	ln, err := net.Listen("tcp", "127.0.0.1:0")
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer ln.Close()
+	go func() {
+		for {
+			c, e := ln.Accept()
+			if e != nil {
+				return
+			}
+			go func() {
+				defer c.Close()
+				var r herdr.Request
+				if json.NewDecoder(bufio.NewReader(c)).Decode(&r) != nil {
+					return
+				}
+				var result string
+				switch r.Method {
+				case "worktree.create":
+					result = `{"path":"` + filepath.Join(filepath.Dir(repo), "worktrees", "task") + `","root_pane":{"pane_id":"p1"}}`
+				case "agent.start":
+					result = `{}`
+				case "pane.get":
+					result = `{"pane":{"agent":"opencode","agent_status":"idle"}}`
+				case "pane.read":
+					result = `{"read":{"text":""}}`
+				case "agent.prompt":
+					result = `{}`
+				default:
+					result = `{}`
+				}
+				_ = json.NewEncoder(c).Encode(herdr.Response{ID: r.ID, Result: json.RawMessage(result)})
+			}()
+		}
+	}()
+	root := filepath.Join(filepath.Dir(repo), "worktrees")
+	w := &worker{herdr: &herdr.Client{Path: ln.Addr().String()}, repo: repo, root: root, remote: "origin", harness: "opencode", tasks: map[string]domain.Task{}, sessions: map[string]herdr.Session{}, leases: map[string]lease{}, statePath: filepath.Join(t.TempDir(), "state.json"), hard: .75}
+	// New() is needed for its pane map; override the address for the fake.
+	w.herdr = herdr.New(ln.Addr().String())
+	task := domain.Task{ID: "task", Source: "s", ExternalID: "x", Project: "p", Title: "test", Description: "do work"}
+	if err := w.start(context.Background(), task, ""); err != nil {
+		t.Fatal(err)
+	}
+	if _, ok := w.sessions[task.ID]; !ok {
+		t.Fatal("session not persisted")
+	}
+	if _, err := os.Stat(filepath.Join(root, "task", "TASK.md")); err != nil {
+		t.Fatalf("TASK.md: %v", err)
+	}
+}