orchestra/internal/orchestrator/orchestrator.go

// Package orchestrator connects router lease events to an opaque herdr
// session. It is deliberately small: scheduling remains in router and the
// adapter remains the only component that knows how to drive a harness.
package orchestrator

import (
	"context"
	"encoding/json"
	"fmt"
	"orchestra/internal/domain"
	"orchestra/internal/herdr"
	"orchestra/internal/store"
	"os"
	"os/exec"
	"path/filepath"
	"sync"
	"time"
)

type Worktrees interface {
	Create(context.Context, domain.Task) (string, error)
}
type Adapters interface {
	Adapter(string) (herdr.Adapter, error)
}

// GitWorktrees creates one isolated checkout per task. The root is expected
// to be a clone containing the project's remote; callers may set a separate
// root per deployment.
type GitWorktrees struct {
	Root string
	Repo string
}

func (w GitWorktrees) Create(ctx context.Context, t domain.Task) (string, error) {
	if w.Root == "" || w.Repo == "" {
		return "", fmt.Errorf("worktree: root and repo required")
	}
	if err := os.MkdirAll(w.Root, 0755); err != nil {
		return "", err
	}
	p := filepath.Join(w.Root, t.ID)
	if _, err := os.Stat(p); err == nil {
		return p, nil
	}
	branch := "orchestra/" + t.ID
	cmd := exec.CommandContext(ctx, "git", "-C", w.Repo, "worktree", "add", "-b", branch, p, "HEAD")
	if out, err := cmd.CombinedOutput(); err != nil {
		return "", fmt.Errorf("%s: %w", string(out), err)
	}
	return p, nil
}

type AdapterFactory struct{ Herdrs map[string]herdr.Adapter }

func (f AdapterFactory) Adapter(id string) (herdr.Adapter, error) {
	a, ok := f.Herdrs[id]
	if !ok {
		return nil, fmt.Errorf("adapter %q not registered", id)
	}
	return a, nil
}

type Coordinator struct {
	Store     *store.Store
	Worktrees Worktrees
	Adapters  Adapters
	StatePath string
	mu        sync.Mutex
	sessions  map[string]herdr.Session
	loaded    bool
}

func (c *Coordinator) loadSessions() {
	c.mu.Lock()
	defer c.mu.Unlock()
	if c.loaded {
		return
	}
	c.loaded = true
	c.sessions = map[string]herdr.Session{}
	if c.StatePath == "" {
		return
	}
	b, err := os.ReadFile(c.StatePath)
	if err != nil {
		return
	}
	_ = json.Unmarshal(b, &c.sessions)
}

func (c *Coordinator) saveSessionsLocked() error {
	if c.StatePath == "" {
		return nil
	}
	b, err := json.Marshal(c.sessions)
	if err != nil {
		return err
	}
	tmp := c.StatePath + ".tmp"
	if err = os.WriteFile(tmp, b, 0600); err != nil {
		return err
	}
	return os.Rename(tmp, c.StatePath)
}

// Reconcile drops mappings whose task lease did not survive restart and kills
// their recoverable herdr sessions so an orphan cannot keep consuming a slot.
func (c *Coordinator) Reconcile(ctx context.Context) error {
	c.loadSessions()
	c.mu.Lock()
	defer c.mu.Unlock()
	for taskID, session := range c.sessions {
		t, ok := c.Store.Task(taskID)
		if ok && t.State == domain.StateLeased {
			continue
		}
		if a, err := c.Adapters.Adapter(session.Harness); err == nil {
			_ = a.Kill(ctx, session)
		}
		delete(c.sessions, taskID)
	}
	return c.saveSessionsLocked()
}

// Monitor performs conservative hard-threshold rotation. The adapter owns
// the handoff creation; the coordinator only publishes its content address
// and frees the lease for pickup by the router.
func (c *Coordinator) Monitor(ctx context.Context, hard float64, interval time.Duration) error {
	if err := c.Reconcile(ctx); err != nil {
		return err
	}
	if interval <= 0 {
		interval = 30 * time.Second
	}
	t := time.NewTicker(interval)
	defer t.Stop()
	for {
		select {
		case <-ctx.Done():
			return ctx.Err()
		case <-t.C:
			c.rotate(ctx, hard)
		}
	}
}
func (c *Coordinator) rotate(ctx context.Context, hard float64) {
	c.loadSessions()
	c.mu.Lock()
	sessions := make(map[string]herdr.Session, len(c.sessions))
	for id, s := range c.sessions {
		sessions[id] = s
	}
	c.mu.Unlock()
	for taskID, session := range sessions {
		task, ok := c.Store.Task(taskID)
		if !ok || task.State != domain.StateLeased {
			continue
		}
		a, err := c.Adapters.Adapter(session.Harness)
		if err != nil {
			continue
		}
		reason := "threshold"
		if signal, ok := a.(herdr.RotationSignal); ok {
			if r, signalErr := signal.RotationSignal(ctx, session); signalErr == nil && r != "" {
				reason = r
			}
		}
		occupancy, err := a.Occupancy(session)
		if err != nil || (occupancy < hard && reason == "threshold") {
			continue
		}
		if boundary, ok := a.(herdr.TurnBoundary); ok {
			atBoundary, boundaryErr := boundary.AtTurnBoundary(ctx, session)
			if boundaryErr == nil && !atBoundary {
				continue
			}
		}
		ref, err := a.Release(ctx, session)
		if err != nil {
			continue
		}
		if ref == "" {
			continue
		}
		b, _ := json.Marshal(map[string]string{"handoff_ref": ref, "reason": reason})
		e := domain.Event{ID: domain.NewID(), Type: "TaskReleased", TaskID: taskID, Version: task.Version + 1, Payload: b}
		if c.Store.Append(e) == nil {
			c.mu.Lock()
			delete(c.sessions, taskID)
			_ = c.saveSessionsLocked()
			c.mu.Unlock()
		}
	}
}

func (c *Coordinator) Start(ctx context.Context, e domain.Event) error {
	if e.Type != "TaskLeased" {
		return nil
	}
	if c.Store == nil || c.Worktrees == nil || c.Adapters == nil {
		return fmt.Errorf("orchestrator: dependencies required")
	}
	c.loadSessions()
	t, ok := c.Store.Task(e.TaskID)
	if !ok {
		return domain.ErrNotFound
	}
	var p struct {
		HarnessID  string `json:"harness_id"`
		HandoffRef string `json:"handoff_ref"`
	}
	if err := json.Unmarshal(e.Payload, &p); err != nil || p.HarnessID == "" {
		return fmt.Errorf("orchestrator: invalid lease")
	}
	w, err := c.Worktrees.Create(ctx, t)
	if err != nil {
		return c.block(t, "worktree: "+err.Error())
	}
	a, err := c.Adapters.Adapter(p.HarnessID)
	if err != nil {
		return c.block(t, "adapter: "+err.Error())
	}
	s, err := a.Lease(ctx, t.ID, w)
	if err != nil {
		return c.block(t, "lease: "+err.Error())
	}
	if p.HandoffRef != "" {
		if err = a.Bootstrap(ctx, s, p.HandoffRef); err != nil {
			_ = a.Kill(ctx, s)
			return c.block(t, "bootstrap: "+err.Error())
		}
	}
	c.mu.Lock()
	if c.sessions == nil {
		c.sessions = map[string]herdr.Session{}
	}
	c.sessions[t.ID] = s
	err = c.saveSessionsLocked()
	c.mu.Unlock()
	return err
}

func (c *Coordinator) block(t domain.Task, reason string) error {
	b, _ := json.Marshal(map[string]string{"blocker": reason})
	return c.Store.Append(domain.Event{ID: domain.NewID(), Type: "TaskBlocked", TaskID: t.ID, Version: t.Version + 1, Payload: b})
}

func (c *Coordinator) Session(taskID string) (herdr.Session, bool) {
	c.loadSessions()
	c.mu.Lock()
	defer c.mu.Unlock()
	s, ok := c.sessions[taskID]
	return s, ok
}