orchestra/internal/router/router_test.go

package router

import (
	"encoding/json"
	"errors"
	"strings"

	"orchestra/internal/authz"
	"orchestra/internal/domain"
	"orchestra/internal/registry"
	"orchestra/internal/store"
	"sync/atomic"
	"testing"
	"time"
)

type reachable struct{}

func (reachable) Reachable(string, time.Duration) bool { return true }

type countedReachability struct {
	calls atomic.Int32
	delay time.Duration
}

func (r *countedReachability) Reachable(string, time.Duration) bool {
	r.calls.Add(1)
	if r.delay > 0 {
		time.Sleep(r.delay)
	}
	return true
}

type countedAvailability struct{ calls atomic.Int32 }

func (a *countedAvailability) Available(registry.Herdr) bool {
	a.calls.Add(1)
	return true
}

type projectAvailability struct{ projects map[string]bool }

func (p projectAvailability) Available(registry.Herdr) bool { return true }
func (p projectAvailability) Supports(h registry.Herdr, project string) bool {
	return p.projects[h.ID+"/"+project]
}

func TestAssignsByAffinityCapabilityAndConcurrency(t *testing.T) {
	s, err := store.Open(t.TempDir())
	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: "h", MachineID: "m", Capabilities: []string{"go"}, Concurrency: 1}},
	})
	if err != nil {
		t.Fatal(err)
	}
	makeTask := func(id string) {
		b, _ := json.Marshal(map[string]any{"source": "test", "external_id": id, "project": "p", "capability": []string{"go"}})
		if err := s.Append(domain.Event{ID: id, TaskID: id, Type: "TaskCreated", Version: 1, Payload: b, Surface: string(authz.System)}); err != nil {
			t.Fatal(err)
		}
	}
	makeTask("a")
	makeTask("b")
	rt := Router{Store: s, Registry: r, Reachability: reachable{}}
	got, err := rt.AssignPending()
	if err != nil || len(got) != 1 {
		t.Fatalf("assigned %d events, err=%v", len(got), err)
	}
	// Store.Tasks() ranges a map, so its order is randomized per call. This
	// assertion used to index two *separate* Tasks() calls, and failed
	// whenever the two orderings disagreed — the flake was in the test, not
	// in assignment. Snapshot once, and assert the actual invariant:
	// concurrency 1 means exactly one of the two tasks is leased.
	leased := 0
	for _, tk := range s.Tasks() {
		if tk.State == domain.StateLeased {
			leased++
		}
	}
	if leased != 1 {
		t.Fatalf("leased %d tasks, want exactly 1 (concurrency 1); got=%+v", leased, got)
	}
}

func TestAssignPendingSnapshotsAndCachesHealth(t *testing.T) {
	s, err := store.Open(t.TempDir())
	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", Address: "h1:1"},
			{ID: "h2", MachineID: "m", Address: "h2:1"},
			{ID: "h3", MachineID: "m", Address: "h3:1"},
		},
	})
	if err != nil {
		t.Fatal(err)
	}
	for i := 0; i < 10; i++ {
		id := domain.NewID()
		payload, _ := json.Marshal(map[string]any{"source": "test", "external_id": id, "project": "p"})
		if err := s.Append(domain.Event{ID: id, TaskID: id, Type: "TaskCreated", Version: 1, Payload: payload, Surface: string(authz.System)}); err != nil {
			t.Fatal(err)
		}
	}
	reach := &countedReachability{}
	availability := &countedAvailability{}
	rt := Router{Store: s, Registry: r, Reachability: reach, Availability: availability, HealthTTL: time.Minute}
	if got, err := rt.AssignPending(); err != nil || len(got) != 10 {
		t.Fatalf("assigned=%d, err=%v", len(got), err)
	}
	if got := reach.calls.Load(); got != 3 {
		t.Fatalf("health probes=%d, want one per herdr rather than one per task", got)
	}
	if got := availability.calls.Load(); got != 1 {
		t.Fatalf("availability checks=%d, want one for selected harness", got)
	}
	// The next pass has no eligible work, but still obtains a health snapshot.
	// It must use the TTL cache rather than re-dial all three herdrs.
	if _, err := rt.AssignPending(); err != nil {
		t.Fatal(err)
	}
	if got := reach.calls.Load(); got != 3 {
		t.Fatalf("health cache missed: probes=%d, want 3", got)
	}
}

func TestHealthProbesRunInParallel(t *testing.T) {
	s, err := store.Open(t.TempDir())
	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", Address: "h1:1"}, {ID: "h2", MachineID: "m", Address: "h2:1"}, {ID: "h3", MachineID: "m", Address: "h3:1"}},
	})
	if err != nil {
		t.Fatal(err)
	}
	payload, _ := json.Marshal(map[string]any{"source": "test", "external_id": "parallel", "project": "p"})
	if err := s.Append(domain.Event{ID: "parallel", TaskID: "parallel", Type: "TaskCreated", Version: 1, Payload: payload, Surface: string(authz.System)}); err != nil {
		t.Fatal(err)
	}
	reach := &countedReachability{delay: 100 * time.Millisecond}
	started := time.Now()
	if _, err := (&Router{Store: s, Registry: r, Reachability: reach}).AssignPending(); err != nil {
		t.Fatal(err)
	}
	if elapsed := time.Since(started); elapsed > 250*time.Millisecond {
		t.Fatalf("health probes took %s; expected parallel probes, not ~300ms serial", elapsed)
	}
}

func TestAssignPendingRequiresWorkerProjectSupport(t *testing.T) {
	s, err := store.Open(t.TempDir())
	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: "h", MachineID: "m", Concurrency: 1}}})
	if err != nil {
		t.Fatal(err)
	}
	b, _ := json.Marshal(map[string]any{"source": "test", "external_id": "project-check", "project": "p"})
	if err := s.Append(domain.Event{ID: "create", TaskID: "t", Type: "TaskCreated", Version: 1, Payload: b, Surface: string(authz.System)}); err != nil {
		t.Fatal(err)
	}
	rt := Router{Store: s, Registry: r, Reachability: reachable{}, Availability: projectAvailability{projects: map[string]bool{}}}
	if got, err := rt.AssignPending(); err != nil || len(got) != 0 {
		t.Fatalf("unsupported project lease = %#v, %v", got, err)
	}
	if got, _ := s.Task("t"); got.State != domain.StateQueued {
		t.Fatalf("unsupported project state=%s", got.State)
	}
	rt.Availability = projectAvailability{projects: map[string]bool{"h/p": true}}
	if got, err := rt.AssignPending(); err != nil || len(got) != 1 {
		t.Fatalf("supported project lease = %#v, %v", got, err)
	}
}

// TestRotationDoesNotCountAgainstRetryLimit guards B4: rotation is
// TaskReleased carrying a valid handoff_ref (spec §5.3: "rotation =
// intra-task lease transfer"), never a failure. A task healthy enough to
// rotate repeatedly must survive past MaxAttempts, which is a retry policy
// for genuine failures (expiry/crash releases without a handoff_ref), not
// for lease transfers.
func TestRotationDoesNotCountAgainstRetryLimit(t *testing.T) {
	s, err := store.Open(t.TempDir())
	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: "h", MachineID: "m", Capabilities: []string{"go"}, Concurrency: 1}},
	})
	if err != nil {
		t.Fatal(err)
	}
	b, _ := json.Marshal(map[string]any{"source": "test", "external_id": "a", "project": "p", "capability": []string{"go"}})
	if err := s.Append(domain.Event{ID: "a", TaskID: "a", Type: "TaskCreated", Version: 1, Payload: b, Surface: string(authz.System)}); err != nil {
		t.Fatal(err)
	}
	rt := Router{Store: s, Registry: r, Reachability: reachable{}, Retry: RetryPolicy{MaxAttempts: 3}}
	handoffRef, err := s.PutArtifact([]byte("handoff"))
	if err != nil {
		t.Fatal(err)
	}

	for i := 0; i < 5; i++ {
		got, err := rt.AssignPending()
		if err != nil {
			t.Fatal(err)
		}
		task, ok := s.Task("a")
		if !ok {
			t.Fatal("task missing")
		}
		if task.State == domain.StateFailed {
			t.Fatalf("task failed after %d rotations, retry limit wrongly counted rotation as a failure", i)
		}
		if task.State != domain.StateLeased {
			if len(got) == 0 {
				t.Fatalf("round %d: task not leased and nothing assigned (state=%v)", i, task.State)
			}
			continue
		}
		rb, _ := json.Marshal(map[string]any{
			"handoff_ref":      handoffRef,
			"reason":           "threshold",
			"anchor_sha":       "0123456789abcdef0123456789abcdef01234567",
			"harness_id":       task.Lease.HarnessID,
			"lease_epoch":      task.Lease.Epoch,
			"expected_version": task.Version,
		})
		release := domain.Event{ID: domain.NewID(), Type: "TaskReleased", TaskID: "a", Version: task.Version + 1, Payload: rb, Surface: string(authz.System)}
		if err := s.Append(release); err != nil {
			t.Fatal(err)
		}
		if _, err := rt.HandleEvent(release); err != nil {
			t.Fatal(err)
		}
	}
	task, _ := s.Task("a")
	if task.State == domain.StateFailed {
		t.Fatal("task failed after 5 rotations, want still alive")
	}
}

// TestQuotaWindowsAreIndependent proves the 5-hour rolling window and the
// weekly window (spec §7.2, §9 item 1) are each conservative-80%-full gates
// on their own — a harness can be fine on one window and excluded by the
// other, and receipts outside a window must not count toward it.
func TestQuotaWindowsAreIndependent(t *testing.T) {
	s, err := store.Open(t.TempDir())
	if err != nil {
		t.Fatal(err)
	}
	now := time.Now().UTC()
	report := func(at time.Time, consumed float64) {
		p, _ := json.Marshal(map[string]any{"harness_id": "h1", "consumed": consumed})
		if err := s.Append(domain.Event{ID: domain.NewID(), Type: "QuotaReported", TaskID: "quota", Version: 1, Surface: string(authz.System), Payload: p, At: at}); err != nil {
			t.Fatal(err)
		}
	}

	// Case 1: only weekly limit configured. A receipt older than 5h but
	// within the week still counts toward the weekly gate.
	report(now.Add(-6*time.Hour), 85)
	weeklyOnly := QuotaAvailability{Store: s, Limits: map[string]QuotaWindowLimits{"h1": {Weekly: 100}}, Now: func() time.Time { return now }}
	if weeklyOnly.Available(registry.Herdr{ID: "h1"}) {
		t.Fatal("weekly window should be exhausted at 85/100 (>=80%)")
	}

	// Case 2: only a 5h limit configured. The same 6h-old receipt is outside
	// the 5h window and must not count; a fresh zero receipt proves the native
	// usage source is known for this window.
	report(now, 0)
	fiveHourOnly := QuotaAvailability{Store: s, Limits: map[string]QuotaWindowLimits{"h1": {FiveHour: 100}}, Now: func() time.Time { return now }}
	if !fiveHourOnly.Available(registry.Herdr{ID: "h1"}) {
		t.Fatal("receipt outside the 5h window incorrectly counted against it")
	}

	// Case 3: a fresh receipt inside the 5h window trips the 5h gate even
	// though the weekly gate (fed by both receipts) also trips — both are
	// independently enforced, and either failing excludes the harness.
	report(now.Add(-time.Minute), 90)
	both := QuotaAvailability{Store: s, Limits: map[string]QuotaWindowLimits{"h1": {FiveHour: 100, Weekly: 500}}, Now: func() time.Time { return now }}
	if both.Available(registry.Herdr{ID: "h1"}) {
		t.Fatal("5h window should be exhausted at 90/100 (>=80%) regardless of weekly headroom")
	}
	unknown := QuotaAvailability{Store: s, Limits: map[string]QuotaWindowLimits{"unknown": {FiveHour: 100}}, Now: func() time.Time { return now }}
	if unknown.Available(registry.Herdr{ID: "unknown"}) {
		t.Fatal("bounded harness without a native usage receipt must fail closed")
	}
}

// Every eligibility gate must say why. A silent `continue` is
// indistinguishable from an empty queue: during burn-in a task sat queued
// while every gate checked out by hand, and the router reported nothing.
func TestAssignPendingRecordsWhyItPlacedNothing(t *testing.T) {
	s, err := store.Open(t.TempDir())
	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: "h", MachineID: "m", Concurrency: 1}},
	})
	if err != nil {
		t.Fatal(err)
	}
	b, _ := json.Marshal(map[string]any{"source": "test", "external_id": "why", "project": "p"})
	if err := s.Append(domain.Event{ID: "create", TaskID: "t", Type: "TaskCreated", Version: 1, Payload: b, Surface: string(authz.System)}); err != nil {
		t.Fatal(err)
	}

	// The live shape that went undiagnosed: a worker that has not declared the
	// project. Everything else about it looks healthy.
	rt := Router{Store: s, Registry: r, Reachability: reachable{}, Availability: projectAvailability{projects: map[string]bool{}}}
	if got, err := rt.AssignPending(); err != nil || len(got) != 0 {
		t.Fatalf("lease = %#v, %v", got, err)
	}
	reasons := rt.Rejections()
	if len(reasons) == 0 {
		t.Fatal("the router placed nothing and said nothing")
	}
	var named bool
	for _, rej := range reasons {
		if rej.TaskID != "t" {
			t.Fatalf("rejection for the wrong task: %+v", rej)
		}
		if rej.HerdrID == "h" && strings.Contains(rej.Reason, "has not declared project") {
			named = true
		}
	}
	if !named {
		t.Fatalf("no rejection names the failing gate: %+v", reasons)
	}

	// A pre-lease refusal is the gate that fails closed on purpose, and it must
	// also be visible rather than looking like "no candidates".
	rt.Availability = projectAvailability{projects: map[string]bool{"h/p": true}}
	s.PreLease = func(string) error { return errors.New("gitea unreachable") }
	if got, err := rt.AssignPending(); err != nil || len(got) != 0 {
		t.Fatalf("lease despite pre-lease refusal = %#v, %v", got, err)
	}
	found := false
	for _, rej := range rt.Rejections() {
		if strings.Contains(rej.Reason, "lease refused") && strings.Contains(rej.Reason, "gitea unreachable") {
			found = true
		}
	}
	if !found {
		t.Fatalf("pre-lease refusal not reported: %+v", rt.Rejections())
	}

	// A queued task in retry backoff is skipped before the candidate loop, so
	// it needs its own reason: it looks assignable and nothing else records it.
	if err := s.Append(domain.Event{ID: "backoff", TaskID: "t", Type: "TaskCorrected", Version: 2, Payload: backoffPayload(time.Now().Add(time.Hour)), Surface: string(authz.System)}); err == nil {
		if got, _ := s.Task("t"); !got.NextRetryAt.IsZero() {
			if _, err := rt.AssignPending(); err != nil {
				t.Fatal(err)
			}
			backoff := false
			for _, rej := range rt.Rejections() {
				if strings.Contains(rej.Reason, "retry backoff") {
					backoff = true
				}
			}
			if !backoff {
				t.Fatalf("a task in retry backoff was silently skipped: %+v", rt.Rejections())
			}
		}
	}

	// A successful pass leaves nothing behind to misread.
	s.PreLease = nil
	if got, err := rt.AssignPending(); err != nil || len(got) != 1 {
		t.Fatalf("lease = %#v, %v", got, err)
	}
	if reasons := rt.Rejections(); len(reasons) != 0 {
		t.Fatalf("stale rejections after a successful pass: %+v", reasons)
	}
}

func backoffPayload(at time.Time) []byte {
	b, _ := json.Marshal(map[string]any{"next_retry_at": at.UTC().Format(time.RFC3339Nano)})
	return b
}