orchestra/cmd/orchestra-worker/main_test.go

package main

import (
	"bufio"
	"context"
	"encoding/json"
	"errors"
	"net"
	"net/http"
	"net/http/httptest"
	"orchestra/internal/domain"
	"orchestra/internal/federation"
	"orchestra/internal/herdr"
	"os"
	"os/exec"
	"path/filepath"
	"testing"
)

func TestWorkerReregistersAfterCoordinatorForgetsIt(t *testing.T) {
	registered := false
	s := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.Method != http.MethodPost || r.URL.Path != "/v1/federation/workers" {
			t.Fatalf("unexpected %s %s", r.Method, r.URL.Path)
		}
		registered = true
		w.WriteHeader(http.StatusCreated)
	}))
	defer s.Close()
	w := &worker{api: federation.Client{BaseURL: s.URL, WorkerID: "h", Token: "t"}, registration: federation.Worker{ID: "h", Capacity: 1}}
	if !w.reRegisterAfterCoordinatorRestart(context.Background(), errors.New("federation: 401 Unauthorized: unknown worker")) {
		t.Fatal("worker did not identify a coordinator restart")
	}
	if !registered {
		t.Fatal("worker did not register")
	}
}

func TestInitialReplayDoesNotResurrectReleasedLease(t *testing.T) {
	s := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		switch r.URL.Path {
		case "/v1/federation/events":
			_, _ = w.Write([]byte(`{"cursor":0,"events":[{"seq":1,"id":"c","type":"TaskCreated","task_id":"t","version":1,"payload":{"source":"s","external_id":"x","project":"p"},"surface":"system"},{"seq":2,"id":"l","type":"TaskLeased","task_id":"t","version":2,"payload":{"harness_id":"h"},"surface":"system"},{"seq":3,"id":"r","type":"TaskReleased","task_id":"t","version":3,"payload":{"reason":"expired"},"surface":"system"}]}`))
		case "/v1/federation/events/ack":
			w.WriteHeader(http.StatusNoContent)
		default:
			t.Fatalf("unexpected %s", r.URL.Path)
		}
	}))
	defer s.Close()
	w := &worker{api: federation.Client{BaseURL: s.URL, WorkerID: "h", Token: "t"}, harnessID: "h", tasks: map[string]domain.Task{}, sessions: map[string]herdr.Session{}, leases: map[string]lease{}, statePath: t.TempDir() + "/state.json", hard: .75}
	if err := w.once(context.Background()); err != nil {
		t.Fatal(err)
	}
	if len(w.leases) != 0 || len(w.sessions) != 0 {
		t.Fatalf("replayed lease survived: leases=%v sessions=%v", w.leases, w.sessions)
	}
	if w.cursor != 3 {
		t.Fatalf("cursor=%d want 3", w.cursor)
	}
}

func TestWorkerHydratesLeaseMissingTaskCache(t *testing.T) {
	s := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		switch r.URL.Path {
		case "/v1/federation/events":
			_, _ = w.Write([]byte(`{"cursor":2,"events":[{"seq":2,"id":"l","type":"TaskLeased","task_id":"t","version":2,"payload":{"harness_id":"h"},"surface":"system"}]}`))
		case "/v1/tasks":
			_, _ = w.Write([]byte(`[{"id":"t","source":"s","external_id":"x","project":"p","state":"leased"}]`))
		case "/v1/federation/events/ack":
			w.WriteHeader(http.StatusNoContent)
		default:
			t.Fatalf("unexpected %s", r.URL.Path)
		}
	}))
	defer s.Close()
	w := &worker{api: federation.Client{BaseURL: s.URL, WorkerID: "h", Token: "t"}, harnessID: "h", cursor: 1, tasks: map[string]domain.Task{}, sessions: map[string]herdr.Session{"t": {}}, leases: map[string]lease{}, statePath: t.TempDir() + "/state.json", hard: .75}
	if err := w.once(context.Background()); err != nil {
		t.Fatal(err)
	}
	if _, ok := w.tasks["t"]; !ok || w.cursor != 2 {
		t.Fatalf("task was not hydrated: tasks=%v cursor=%d", w.tasks, w.cursor)
	}
}

func TestWorkerStartsRouterIssuedLeaseInLocalGitWorktree(t *testing.T) {
	remote := filepath.Join(t.TempDir(), "remote.git")
	if out, err := exec.Command("git", "init", "--bare", remote).CombinedOutput(); err != nil {
		t.Fatalf("remote: %v %s", err, out)
	}
	seed := t.TempDir()
	for _, a := range [][]string{{"init", seed}, {"-C", seed, "config", "user.email", "t@t"}, {"-C", seed, "config", "user.name", "t"}, {"-C", seed, "commit", "--allow-empty", "-m", "init"}, {"-C", seed, "remote", "add", "origin", remote}, {"-C", seed, "push", "-u", "origin", "HEAD:master"}} {
		if out, err := exec.Command("git", a...).CombinedOutput(); err != nil {
			t.Fatalf("git %v: %v %s", a, err, out)
		}
	}
	repo := filepath.Join(t.TempDir(), "repo")
	if out, err := exec.Command("git", "clone", remote, repo).CombinedOutput(); err != nil {
		t.Fatal(string(out))
	}
	ln, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatal(err)
	}
	defer ln.Close()
	go func() {
		for {
			c, e := ln.Accept()
			if e != nil {
				return
			}
			go func() {
				defer c.Close()
				var r herdr.Request
				if json.NewDecoder(bufio.NewReader(c)).Decode(&r) != nil {
					return
				}
				var result string
				switch r.Method {
				case "worktree.create":
					result = `{"path":"` + filepath.Join(filepath.Dir(repo), "worktrees", "task") + `","root_pane":{"pane_id":"p1"}}`
				case "agent.start":
					result = `{}`
				case "pane.get":
					result = `{"pane":{"agent":"opencode","agent_status":"idle"}}`
				case "pane.read":
					result = `{"read":{"text":""}}`
				case "agent.prompt":
					result = `{}`
				default:
					result = `{}`
				}
				_ = json.NewEncoder(c).Encode(herdr.Response{ID: r.ID, Result: json.RawMessage(result)})
			}()
		}
	}()
	root := filepath.Join(filepath.Dir(repo), "worktrees")
	w := &worker{herdr: &herdr.Client{Path: ln.Addr().String()}, repo: repo, root: root, remote: "origin", harness: "opencode", tasks: map[string]domain.Task{}, sessions: map[string]herdr.Session{}, leases: map[string]lease{}, statePath: filepath.Join(t.TempDir(), "state.json"), hard: .75}
	// New() is needed for its pane map; override the address for the fake.
	w.herdr = herdr.New(ln.Addr().String())
	task := domain.Task{ID: "task", Source: "s", ExternalID: "x", Project: "p", Title: "test", Description: "do work"}
	if err := w.start(context.Background(), task, ""); err != nil {
		t.Fatal(err)
	}
	if _, ok := w.sessions[task.ID]; !ok {
		t.Fatal("session not persisted")
	}
	if _, err := os.Stat(filepath.Join(root, "task", "TASK.md")); err != nil {
		t.Fatalf("TASK.md: %v", err)
	}
}