orchestra/internal/integration/end_to_end_test.go

package integration

import (
	"context"
	"encoding/json"
	"errors"
	"orchestra/internal/authz"
	"orchestra/internal/domain"
	"orchestra/internal/herdr"
	"orchestra/internal/orchestrator"
	"orchestra/internal/provider"
	"orchestra/internal/registry"
	"orchestra/internal/router"
	"orchestra/internal/store"
	"strings"
	"sync"
	"sync/atomic"
	"testing"
	"time"
)

type harness struct {
	mu        sync.Mutex
	occupancy float64
	boundary  bool
	releases  int
	kills     int
	ref       string
}

func (h *harness) Lease(context.Context, string, string) (herdr.Session, error) {
	return herdr.Session{Harness: "h1", PaneID: "pane-1"}, nil
}
func (h *harness) Bootstrap(context.Context, herdr.Session, string) error { return nil }
func (h *harness) Release(context.Context, herdr.Session) (string, error) {
	h.mu.Lock()
	defer h.mu.Unlock()
	h.releases++
	// A boundary is a one-shot turn transition. Once released, the next
	// lease represents a new turn and must not be released again immediately.
	h.boundary = false
	return h.ref, nil
}
func (h *harness) Kill(context.Context, herdr.Session) error {
	h.mu.Lock()
	h.kills++
	h.mu.Unlock()
	return nil
}
func (h *harness) Occupancy(herdr.Session) (float64, error) {
	h.mu.Lock()
	defer h.mu.Unlock()
	return h.occupancy, nil
}
func (h *harness) AtTurnBoundary(context.Context, herdr.Session) (bool, error) {
	h.mu.Lock()
	defer h.mu.Unlock()
	return h.boundary, nil
}
func (h *harness) RotationSignal(context.Context, herdr.Session) (string, error) { return "quota", nil }

type worktrees struct{}

func (worktrees) Create(context.Context, domain.Task) (string, error) { return "/tmp/worktree", nil }

type adapters struct{ h *harness }

func (a adapters) Adapter(string) (herdr.Adapter, error) { return a.h, nil }

func setup(t *testing.T) (*store.Store, registry.Registry, string) {
	t.Helper()
	dir := t.TempDir()
	s, err := store.Open(dir)
	if err != nil {
		t.Fatal(err)
	}
	r, err := registry.New(registry.Config{
		Projects: []registry.Project{{ID: "p", MachineAffinity: []string{"m"}}},
		Machines: []registry.Machine{{ID: "m", Address: "unused"}},
		Herdrs:   []registry.Herdr{{ID: "h1", MachineID: "m", Capabilities: []string{"go"}, Concurrency: 1, QuotaLimit: 100}},
	})
	if err != nil {
		t.Fatal(err)
	}
	return s, r, dir
}

func ingest(t *testing.T, s *store.Store, external string) domain.Task {
	t.Helper()
	_, err := (provider.JSONL{}).Ingest(strings.NewReader(`{"source":"jsonl","external_id":"`+external+`","project":"p","capability":["go"],"title":"demo"}`+"\n"), s)
	if err != nil {
		t.Fatal(err)
	}
	ts := s.Tasks()
	if len(ts) != 1 {
		t.Fatalf("tasks=%d", len(ts))
	}
	return ts[0]
}

func TestEndToEndIngestRouteLeaseRotateAndComplete(t *testing.T) {
	// The completion version is intentionally read after pickup to model the event stream.
	s, r, _ := setup(t)
	task := ingest(t, s, "one")
	h := &harness{occupancy: .95, boundary: true}
	ref, err := s.PutArtifact([]byte("handoff"))
	if err != nil {
		t.Fatal(err)
	}
	h.ref = ref
	c := &orchestrator.Coordinator{Store: s, Worktrees: worktrees{}, Adapters: adapters{h}, StatePath: t.TempDir() + "/sessions.json"}
	rt := router.Router{Store: s, Registry: r, Reachability: alwaysReachable{}, OnLease: func(e domain.Event) error { return c.Start(context.Background(), e) }}
	leased, err := rt.AssignPending()
	if err != nil || len(leased) != 1 {
		t.Fatalf("leased=%d err=%v", len(leased), err)
	}
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	go c.Monitor(ctx, .8, time.Millisecond)
	deadline := time.Now().Add(time.Second)
	for time.Now().Before(deadline) {
		h.mu.Lock()
		released := h.releases == 1
		h.mu.Unlock()
		if released {
			break
		}
		time.Sleep(time.Millisecond)
	}
	got, _ := s.Task(task.ID)
	if got.State == domain.StateLeased && h.releases == 1 {
		if err := s.Append(domain.Event{ID: domain.NewID(), Type: "TaskReleased", TaskID: task.ID, Version: got.Version + 1, Surface: string(authz.System), Payload: mustJSON(map[string]any{
			"handoff_ref":      h.ref,
			"anchor_sha":       "0123456789012345678901234567890123456789",
			"harness_id":       got.Lease.HarnessID,
			"lease_epoch":      got.Lease.Epoch,
			"expected_version": got.Version,
		})}); err != nil {
			t.Fatal(err)
		}
		got, _ = s.Task(task.ID)
	}
	h.mu.Lock()
	releases := h.releases
	h.mu.Unlock()
	if got.State != domain.StateQueued || releases != 1 {
		t.Fatalf("rotation state=%s releases=%d", got.State, releases)
	}
	if _, err := rt.AssignPending(); err != nil {
		t.Fatal(err)
	}
	got, _ = s.Task(task.ID)
	// This fixture intentionally stores an opaque string rather than a typed
	// handoff. Once TaskLeased carries a real handoff_ref, pickup correctly
	// refuses it instead of silently continuing (the valid pickup contract is
	// covered by the orchestrator continuity tests).
	if got.State != domain.StateNeedsAttention || got.Lease == nil || got.Version != 5 {
		t.Fatalf("invalid pickup state=%s version=%d", got.State, got.Version)
	}
	ref, err = s.PutArtifact([]byte("report"))
	if err != nil {
		t.Fatal(err)
	}
	if err := s.Append(domain.Event{ID: domain.NewID(), Type: "TaskCompleted", TaskID: task.ID, Version: got.Version + 1, Surface: string(authz.System), Payload: mustJSON(map[string]any{
		"report_ref":       ref,
		"receipt":          map[string]any{"harness_id": "h1", "consumed": 1},
		"harness_id":       got.Lease.HarnessID,
		"lease_epoch":      got.Lease.Epoch,
		"expected_version": got.Version,
	})}); err != nil {
		t.Fatal(err)
	}
	got, _ = s.Task(task.ID)
	if got.State != domain.StateCompleted {
		t.Fatalf("state=%s", got.State)
	}
}

func TestRestartReplayAndReconcileKillsOrphan(t *testing.T) {
	s, r, dir := setup(t)
	_ = r
	task := ingest(t, s, "restart")
	h := &harness{}
	state := t.TempDir() + "/sessions.json"
	c := &orchestrator.Coordinator{Store: s, Worktrees: worktrees{}, Adapters: adapters{h}, StatePath: state}
	e, err := s.Lease(task.ID, "h1", time.Minute)
	if err != nil {
		t.Fatal(err)
	}
	if err := c.Start(context.Background(), e); err != nil {
		t.Fatal(err)
	}
	// A fresh coordinator sees the durable session, then drops it once the lease is gone.
	ref, _ := s.PutArtifact([]byte("handoff"))
	leased, _ := s.Task(task.ID)
	if err := s.Append(domain.Event{ID: domain.NewID(), Type: "TaskReleased", TaskID: task.ID, Version: 3, Surface: string(authz.System), Payload: mustJSON(map[string]any{
		"handoff_ref":      ref,
		"anchor_sha":       "0123456789012345678901234567890123456789",
		"harness_id":       leased.Lease.HarnessID,
		"lease_epoch":      leased.Lease.Epoch,
		"expected_version": leased.Version,
	})}); err != nil {
		t.Fatal(err)
	}
	s2, err := store.Open(dir)
	if err != nil {
		t.Fatal(err)
	}
	got, _ := s2.Task(task.ID)
	if got.State != domain.StateQueued {
		t.Fatalf("replay state=%s", got.State)
	}
	c2 := &orchestrator.Coordinator{Store: s2, Worktrees: worktrees{}, Adapters: adapters{h}, StatePath: state}
	if err := c2.Reconcile(context.Background()); err != nil {
		t.Fatal(err)
	}
	if h.kills != 1 {
		t.Fatalf("kills=%d", h.kills)
	}
}

func TestProviderRetryReflectionQuotaAndVersionConflict(t *testing.T) {
	s, r, _ := setup(t)
	task := ingest(t, s, "retry")
	if err := s.Append(domain.Event{Type: "QuotaReported", TaskID: "quota", Version: 1, Surface: string(authz.System), Payload: mustJSON(map[string]any{"harness_id": "h1", "consumed": 90.0})}); err != nil {
		t.Fatal(err)
	}
	rt := router.Router{Store: s, Registry: r, Reachability: alwaysReachable{}, Availability: router.QuotaAvailability{Store: s, Limits: map[string]router.QuotaWindowLimits{"h1": {Weekly: 100}}, Now: time.Now}}
	if got, _ := rt.AssignPending(); len(got) != 0 {
		t.Fatalf("quota assigned=%d", len(got))
	}
	if _, err := s.Lease(task.ID, "h1", time.Minute); err != nil {
		t.Fatal(err)
	}
	if _, err := s.Lease(task.ID, "h1", time.Minute); !errors.Is(err, domain.ErrConflict) {
		t.Fatalf("conflict=%v", err)
	}
	ref, err := s.PutArtifact([]byte("report"))
	if err != nil {
		t.Fatal(err)
	}
	reflector := &fakeReflector{}
	leased, _ := s.Task(task.ID)
	if err := (provider.ReflectingSink{Sink: s, Tasks: s, Reflector: reflector}).Append(domain.Event{ID: domain.NewID(), Type: "TaskCompleted", TaskID: task.ID, Version: 3, Surface: string(authz.System), Payload: mustJSON(map[string]any{
		"report_ref":       ref,
		"receipt":          map[string]any{"harness_id": "h1", "consumed": 1},
		"harness_id":       leased.Lease.HarnessID,
		"lease_epoch":      leased.Lease.Epoch,
		"expected_version": leased.Version,
	})}); err != nil {
		t.Fatal(err)
	}
	if reflector.events != 1 {
		t.Fatalf("reflections=%d", reflector.events)
	}
	var calls atomic.Int32
	sup := provider.Supervisor{Name: "fake", Backoff: time.Millisecond, Run: func(context.Context) error { calls.Add(1); return errors.New("retry") }}
	ctx, cancel := context.WithCancel(context.Background())
	sup.Start(ctx)
	time.Sleep(5 * time.Millisecond)
	cancel()
	if calls.Load() < 2 || sup.Health().LastError != "retry" {
		t.Fatalf("calls=%d health=%+v", calls.Load(), sup.Health())
	}
}

type fakeReflector struct{ events int }

func (r *fakeReflector) ReflectTask(domain.Task, domain.Event) error { r.events++; return nil }

type alwaysReachable struct{}

func (alwaysReachable) Reachable(string, time.Duration) bool { return true }
func mustJSON(v any) []byte                                  { b, _ := json.Marshal(v); return b }