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
@@ -46,6 +46,49 @@ func herdrAddress(rr registry.Registry, h registry.Herdr) string {
 	return net.JoinHostPort(host, defaultHerdrPort)
 }

+type federatedAvailability struct {
+	base         router.Availability
+	workers      *federation.Registry
+	localMachine string
+}
+
+// federatedReachability keeps the legacy TCP probe for local herdrs, while
+// avoiding a coordinator-side probe of a remote worker's herdr socket. A
+// remote harness is reachable precisely when its worker is registered (as
+// enforced by federatedAvailability); probing its raw herdr endpoint here
+// would reintroduce the cross-machine Design A dependency.
+type federatedReachability struct {
+	base   registry.Reachability
+	remote map[string]bool
+}
+
+func (r federatedReachability) Reachable(address string, timeout time.Duration) bool {
+	if r.remote[address] {
+		return true
+	}
+	return r.base.Reachable(address, timeout)
+}
+
+func remoteHerdrAddresses(rr registry.Registry, localMachine string) map[string]bool {
+	remote := map[string]bool{}
+	for _, h := range rr.Herdrs() {
+		if h.MachineID != localMachine {
+			remote[herdrAddress(rr, h)] = true
+		}
+	}
+	return remote
+}
+
+func (a federatedAvailability) Available(h registry.Herdr) bool {
+	if a.base != nil && !a.base.Available(h) {
+		return false
+	}
+	if a.localMachine == "" || h.MachineID == a.localMachine {
+		return true
+	}
+	return a.workers.Available(h.ID)
+}
+
 func main() {
 	dir := os.Getenv("ORCHESTRA_DATA")
 	if dir == "" {
@@ -58,11 +101,17 @@ func main() {
 	var rr registry.Registry
 	var rt *router.Router
 	var coordinator *orchestrator.Coordinator
+	localMachine := os.Getenv("ORCHESTRA_MACHINE_ID")
+	workers := &federation.Registry{AdmitToken: os.Getenv("ORCHESTRA_FEDERATION_ADMIT_TOKEN")}
 	if config := os.Getenv("ORCHESTRA_CONFIG"); config != "" {
 		if rr, err = registry.Load(config); err != nil {
 			log.Fatalf("load orchestra config: %v", err)
 		}
-		rt = &router.Router{Store: s, Registry: rr, Reachability: registry.TCPReachability{}, Timeout: time.Second, Retry: router.RetryPolicy{MaxAttempts: 3, Backoff: time.Minute}}
+		reachability := registry.Reachability(registry.TCPReachability{})
+		if localMachine != "" {
+			reachability = federatedReachability{base: reachability, remote: remoteHerdrAddresses(rr, localMachine)}
+		}
+		rt = &router.Router{Store: s, Registry: rr, Reachability: reachability, Timeout: time.Second, Retry: router.RetryPolicy{MaxAttempts: 3, Backoff: time.Minute}}
 		limits := map[string]router.QuotaWindowLimits{}
 		for _, h := range rr.Herdrs() {
 			w := router.QuotaWindowLimits{FiveHour: h.QuotaLimit5h, Weekly: h.QuotaLimitWeekly}
@@ -77,6 +126,9 @@ func main() {
 		if len(limits) > 0 {
 			rt.Availability = router.QuotaAvailability{Store: s, Limits: limits}
 		}
+		if localMachine != "" {
+			rt.Availability = federatedAvailability{base: rt.Availability, workers: workers, localMachine: localMachine}
+		}
 		if repo, root := os.Getenv("ORCHESTRA_REPO"), os.Getenv("ORCHESTRA_WORKTREE_ROOT"); repo != "" && root != "" {
 			adapters := map[string]herdr.Adapter{}
 			for _, h := range rr.Herdrs() {
@@ -115,7 +167,22 @@ func main() {
 				Default:  orchestrator.GitWorktrees{Root: root, Repo: repo},
 			}
 			coordinator = &orchestrator.Coordinator{Store: s, StatePath: filepath.Join(dir, "runtime-sessions.json"), Worktrees: worktrees, Adapters: orchestrator.AdapterFactory{Herdrs: adapters}}
-			rt.OnLease = func(e domain.Event) error { return coordinator.Start(context.Background(), e) }
+			rt.OnLease = func(e domain.Event) error {
+				// In federated mode the coordinator must never inspect a remote
+				// checkout. Its worker consumes the router-issued lease event and
+				// performs all Git/herdr operations on that machine (§2.1).
+				if localMachine != "" {
+					var p struct {
+						HarnessID string `json:"harness_id"`
+					}
+					if json.Unmarshal(e.Payload, &p) == nil {
+						if h, ok := rr.Herdr(p.HarnessID); ok && h.MachineID != localMachine {
+							return nil
+						}
+					}
+				}
+				return coordinator.Start(context.Background(), e)
+			}
 			hard := 0.75
 			if v, parseErr := strconv.ParseFloat(os.Getenv("ORCHESTRA_OCCUPANCY_HARD"), 64); parseErr == nil && v > 0 && v < 1 {
 				hard = v
@@ -131,7 +198,6 @@ func main() {
 		}
 	}
 	mux := http.NewServeMux()
-	workers := &federation.Registry{AdmitToken: os.Getenv("ORCHESTRA_FEDERATION_ADMIT_TOKEN")}
 	workers.OnOffline = func(w federation.Worker) {
 		for _, t := range s.Tasks() {
 			if t.State == domain.StateLeased && t.Lease != nil && t.Lease.HarnessID == w.ID {
@@ -143,6 +209,13 @@ func main() {
 			}
 		}
 	}
+	go func() {
+		ticker := time.NewTicker(15 * time.Second)
+		defer ticker.Stop()
+		for range ticker.C {
+			_ = workers.Snapshot()
+		}
+	}()
 	providerHealth := map[string]*provider.Supervisor{}
 	surface := func(r *http.Request) authz.Surface {
 		v := authz.ParseSurface(r.Header.Get("X-Orchestra-Surface"))
@@ -245,6 +318,33 @@ func main() {
 		w.WriteHeader(http.StatusCreated)
 		json.NewEncoder(w).Encode(map[string]string{"ref": ref})
 	})
+	// Artifacts are normally write-only to public surfaces. A federation
+	// worker may read a handoff only after authenticating as the worker that
+	// will validate it against its own checkout (§2.1, §6.2).
+	mux.HandleFunc("/v1/artifacts/", func(w http.ResponseWriter, r *http.Request) {
+		if r.Method != http.MethodGet {
+			http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
+			return
+		}
+		wid := r.Header.Get("X-Orchestra-Worker")
+		token := strings.TrimPrefix(r.Header.Get("Authorization"), "Bearer ")
+		if err := workers.Authenticate(wid, token); err != nil {
+			http.Error(w, err.Error(), http.StatusUnauthorized)
+			return
+		}
+		ref := strings.TrimPrefix(r.URL.Path, "/v1/artifacts/")
+		if len(ref) != 64 {
+			http.Error(w, "artifact ref required", http.StatusBadRequest)
+			return
+		}
+		b, err := s.Artifact(ref)
+		if err != nil {
+			http.Error(w, "artifact not found", http.StatusNotFound)
+			return
+		}
+		w.Header().Set("Content-Type", "application/json")
+		_, _ = w.Write(b)
+	})
 	// /v1/harness/complete is the automatic TaskCompleted producer (AUDIT.md
 	// B3): a harness-side hook posts here when the agent has declared the
 	// task done (see deploy/hooks/orchestra-stop.sh), not on every turn
@@ -715,6 +815,11 @@ func main() {
 			http.Error(w, err.Error(), status)
 			return
 		}
+		if rt != nil {
+			if _, err := rt.AssignPending(); err != nil {
+				log.Printf("route after worker registration: %v", err)
+			}
+		}
 		w.WriteHeader(http.StatusCreated)
 		json.NewEncoder(w).Encode(worker)
 	})
@@ -768,7 +873,7 @@ func main() {
 		w.WriteHeader(http.StatusNoContent)
 	})
 	mux.HandleFunc("/v1/federation/workers/", func(w http.ResponseWriter, r *http.Request) {
-		if r.Method != http.MethodPost || (!strings.HasSuffix(r.URL.Path, "/heartbeat") && !strings.HasSuffix(r.URL.Path, "/claim") && !strings.HasSuffix(r.URL.Path, "/handoff")) {
+		if r.Method != http.MethodPost || (!strings.HasSuffix(r.URL.Path, "/heartbeat") && !strings.HasSuffix(r.URL.Path, "/handoff") && !strings.HasSuffix(r.URL.Path, "/complete")) {
 			http.Error(w, "not found", 404)
 			return
 		}
@@ -804,22 +909,24 @@ func main() {
 			http.Error(w, "task not found", 404)
 			return
 		}
-		if strings.HasSuffix(r.URL.Path, "/claim") {
-			if b.TTLSeconds <= 0 {
-				b.TTLSeconds = 1800
+		if t.State != domain.StateLeased || t.Lease == nil || t.Lease.HarnessID != parts[3] {
+			http.Error(w, "lease not owned", 409)
+			return
+		}
+		if strings.HasSuffix(r.URL.Path, "/complete") {
+			if b.HandoffRef == "" {
+				http.Error(w, "report_ref required", 400)
+				return
 			}
-			e, err := s.Lease(b.TaskID, parts[3], time.Duration(b.TTLSeconds)*time.Second)
-			if err != nil {
+			p, _ := json.Marshal(map[string]any{"report_ref": b.HandoffRef, "receipt": map[string]any{"harness_id": parts[3], "consumed": 0}})
+			e := domain.Event{ID: id(), Type: "TaskCompleted", TaskID: b.TaskID, Version: t.Version + 1, Payload: p, Surface: string(authz.System)}
+			if err := s.Append(e); err != nil {
 				http.Error(w, err.Error(), 409)
 				return
 			}
 			json.NewEncoder(w).Encode(e)
 			return
 		}
-		if t.State != domain.StateLeased || t.Lease == nil || t.Lease.HarnessID != parts[3] {
-			http.Error(w, "lease not owned", 409)
-			return
-		}
 		if b.HandoffRef == "" {
 			http.Error(w, "handoff_ref required", 400)
 			return
@@ -0,0 +1,35 @@
+package main
+
+import (
+	"os"
+	"path/filepath"
+	"testing"
+	"time"
+
+	"orchestra/internal/registry"
+)
+
+type unreachable struct{}
+
+func (unreachable) Reachable(string, time.Duration) bool { return false }
+
+func TestFederatedReachabilityDefersRemoteHerdrToWorkerHeartbeat(t *testing.T) {
+	path := filepath.Join(t.TempDir(), "config.json")
+	if err := os.WriteFile(path, []byte(`{
+  "machines":[{"id":"homesrv","address":"192.168.1.104:9145"},{"id":"workpc","address":"192.168.1.105:9145"}],
+  "herdrs":[{"id":"local","machine_id":"homesrv","harness":"opencode"},{"id":"remote","machine_id":"workpc","harness":"opencode"}]
+	}`), 0o600); err != nil {
+		t.Fatal(err)
+	}
+	r, err := registry.Load(path)
+	if err != nil {
+		t.Fatal(err)
+	}
+	check := federatedReachability{base: unreachable{}, remote: remoteHerdrAddresses(r, "homesrv")}
+	if check.Reachable("192.168.1.105:9245", time.Second) != true {
+		t.Fatal("remote herdr should be admitted for worker heartbeat gating")
+	}
+	if check.Reachable("192.168.1.104:9245", time.Second) {
+		t.Fatal("local herdr should still require its TCP probe")
+	}
+}