kami/orchestra
1
0
Fork 0
Code Issues Pull Requests 1 Actions Packages Projects Releases Wiki Activity

Harden worker federation and operator UI

Browse Source
This commit is contained in:
kami
2026-07-29 13:30:55 +04:00
parent 95a96d87a5
commit 1ca9d64e89
35 changed files with 1195 additions and 581 deletions
Download Patch File Download Diff File Expand all files Collapse all files
internal/federation/client.go
+4 -2
View File
@@ -124,8 +124,8 @@ func (c Client) Ack(ctx context.Context, cursor uint64) error {
}
return err
}
func (c Client) Heartbeat(ctx context.Context) error {
resp, err := c.request(ctx, http.MethodPost, "/v1/federation/workers/"+url.PathEscape(c.WorkerID)+"/heartbeat", nil)
func (c Client) Heartbeat(ctx context.Context, health WorkerHealth) error {
resp, err := c.request(ctx, http.MethodPost, "/v1/federation/workers/"+url.PathEscape(c.WorkerID)+"/heartbeat", health)
if resp != nil {
resp.Body.Close()
}
@@ -146,6 +146,8 @@ func (c Client) PutArtifact(ctx context.Context, b []byte) (string, error) {
return "", err
}
req.Header.Set("Content-Type", "application/octet-stream")
req.Header.Set("X-Orchestra-Worker", c.WorkerID)
req.Header.Set("Authorization", "Bearer "+c.Token)
h := c.HTTP
if h == nil {
h = http.DefaultClient
internal/federation/federation.go
+130 -10
View File
@@ -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
}
internal/federation/federation_test.go
+87
View File
@@ -1,6 +1,8 @@
package federation
import (
"os"
"path/filepath"
"testing"
"time"
)
@@ -27,6 +29,68 @@ func TestCursorIsMonotonicAndAuthenticationIsRequired(t *testing.T) {
}
}
func TestPendingApprovalSurvivesRegistryRestart(t *testing.T) {
path := filepath.Join(t.TempDir(), "federation-state.json")
r := &Registry{StatePath: path}
if err := r.Load(); err != nil {
t.Fatal(err)
}
if err := r.Register(Worker{ID: "w", Token: "t"}, ""); err != nil {
t.Fatal(err)
}
capture, err := r.PutCapture("w", Capture{TaskID: "task", PaneID: "pane", Text: "Allow command?"})
if err != nil {
t.Fatal(err)
}
queued, err := r.Queue("w", Command{TaskID: "task", Kind: "grant_approval", PaneID: "pane", CaptureRevision: capture.Revision})
if err != nil {
t.Fatal(err)
}
info, err := os.Stat(path)
if err != nil {
t.Fatal(err)
}
if info.Mode().Perm() != 0600 {
t.Fatalf("federation state permissions = %o, want 0600", info.Mode().Perm())
}
restarted := &Registry{StatePath: path}
if err := restarted.Load(); err != nil {
t.Fatal(err)
}
if err := restarted.Register(Worker{ID: "w", Token: "intruder"}, ""); err != ErrUnauthorized {
t.Fatalf("recovered worker identity was hijackable: %v", err)
}
// A restart does not mark the worker online; it must prove its retained
// identity by registering again before recovered controls become available.
if err := restarted.Register(Worker{ID: "w", Token: "t"}, ""); err != nil {
t.Fatal(err)
}
gotCapture, ok := restarted.Capture("w", "task")
if !ok || gotCapture.Revision != capture.Revision || gotCapture.Text != capture.Text {
t.Fatalf("capture after restart = %#v, present=%v", gotCapture, ok)
}
commands, err := restarted.Commands("w")
if err != nil || len(commands) != 1 || commands[0].ID != queued.ID {
t.Fatalf("commands after restart = %#v, err=%v", commands, err)
}
if err := restarted.CompleteCommand("w", queued.ID, "acknowledged", ""); err != nil {
t.Fatal(err)
}
again := &Registry{StatePath: path}
if err := again.Load(); err != nil {
t.Fatal(err)
}
if err := again.Register(Worker{ID: "w", Token: "t"}, ""); err != nil {
t.Fatal(err)
}
commands, err = again.Commands("w")
if err != nil || len(commands) != 0 {
t.Fatalf("resolved command recovered as pending: %#v, err=%v", commands, err)
}
}
func TestRegisterRequiresAdmitTokenAndOwnToken(t *testing.T) {
r := &Registry{AdmitToken: "admit-secret"}
if err := r.Register(Worker{ID: "workpc", Token: "secret"}, "wrong"); err != ErrUnauthorized {
@@ -76,6 +140,29 @@ func TestOfflineHookRunsOnceOnTransition(t *testing.T) {
}
}
func TestHeartbeatProjectsWorkerOwnedHealth(t *testing.T) {
r := &Registry{}
if err := r.Register(Worker{ID: "workpc-opencode", Token: "secret"}, ""); err != nil {
t.Fatal(err)
}
checked := time.Now().UTC().Round(0)
errAt := checked.Add(-time.Minute)
if err := r.Heartbeat("workpc-opencode", WorkerHealth{
HerdrStatus: "unreachable", CheckedAt: checked, ActiveTask: "task-1", ActivePane: "pane-1",
LastError: "local herdr: connection refused", ErrorAt: errAt,
}); err != nil {
t.Fatal(err)
}
workers := r.Snapshot()
if len(workers) != 1 {
t.Fatalf("workers=%#v", workers)
}
h := workers[0].Health
if h.HerdrStatus != "unreachable" || h.ActiveTask != "task-1" || h.ActivePane != "pane-1" || h.LastError == "" || !h.CheckedAt.Equal(checked) || !h.ErrorAt.Equal(errAt) {
t.Fatalf("health=%#v", h)
}
}
func TestCaptureRevisionAndCommandQueue(t *testing.T) {
r := &Registry{}
if err := r.Register(Worker{ID: "w", Token: "t"}, ""); err != nil {