diff --git a/extensions/ty-on/.gitignore b/extensions/ty-on/.gitignore new file mode 100644 index 00000000..b507ea27 --- /dev/null +++ b/extensions/ty-on/.gitignore @@ -0,0 +1,6 @@ +# Binary +ty-on + +# Test artifacts +*.test +coverage.out diff --git a/extensions/ty-on/README.md b/extensions/ty-on/README.md new file mode 100644 index 00000000..91962efa --- /dev/null +++ b/extensions/ty-on/README.md @@ -0,0 +1,131 @@ +# ty-on + +Placement resolver for TaskYou. Decides which machine a task should run on, and +answers "this one, here" or "run it locally". + +## Why this is an extension + +ty can run a task on another machine. The *policy* for that — which hosts exist, +what they are provisioned for, which one to pick — is specific to whoever owns +the fleet, so it does not belong in ty. A normal ty user never sees any of it. + +This extension is the policy half. It touches nothing in ty's core: ty invokes +it, and it answers. + +## The contract + +ty-on is a binary that reads one JSON request on stdin and writes one JSON +response on stdout. It is invoked once per task, before the executor is spawned. + +Request: + +```json +{ + "event": "task.placement", + "task": { + "id": 5225, + "title": "Some task", + "project": "taskyou", + "repo_path": "/Users/bruno/Projects/workflow", + "executor": "claude" + } +} +``` + +Response: + +```json +{ + "target": "ol-agents", + "workdir": "~/projects/engineering", + "reason": "most free memory of 2 hosts serving offerlab (ol-agents 26.5G, mona 11.3G)" +} +``` + +`target` names a host in the `on` inventory. `workdir` is that project's +checkout path on that host — a remote path, so a leading `~` is left alone for +the remote shell to expand. + +**An empty `target` means "run locally"**, and it is the answer to every +question ty-on cannot confidently answer: unknown project, missing inventory, no +reachable host, malformed request, `on` not installed. ty-on never fails a task +and never guesses a host — it exits 0 in all cases. + +`reason` is always populated and is shown to the user, so it is written to +explain a surprising placement without further digging. + +## Placement rules + +The inventory is the same one the [`on`](https://github.com/bborn/on) CLI reads: +`$ON_HOSTS`, else `$XDG_CONFIG_HOME/on/hosts.yaml`, else +`~/.config/on/hosts.yaml`. + +```yaml +hosts: + ol-agents: + ssh: ol-agents + workdir: ~/projects + capabilities: [agent, ruby, node] + repos: + offerlab: ~/projects/engineering +``` + +Given a task's project: + +1. Find the hosts whose `repos` map contains that project. +2. **None** → local. The fleet has no checkout to run in. +3. **One** → that host. No probing: this path answers from the file alone, so it + works on a machine that does not have `on` installed at all. +4. **Several** → the one with the most free memory. `on ls` already probes the + fleet in parallel, so ty-on shells out to it rather than reimplementing the + probe. Hosts `on` could not reach are dropped; ties break on host name so the + answer is stable. + +`on` is an optional dependency. If it is missing, or fails, or is slow, the task +stays local with a reason saying so. + +### Speed + +This runs in the task spawn path, so it is built to be fast or to get out of the +way. Rules 1–3 are a single file read. Rule 4 costs one `on ls` (an SSH round +trip per host, in parallel), bounded by `TY_ON_TIMEOUT` — past that budget ty-on +prefers a local placement to a late answer. + +## Usage + +```console +$ go build -o ty-on ./cmd +$ echo '{"event":"task.placement","task":{"project":"taskyou"}}' | ./ty-on +{"target":"mona","workdir":"~/Projects/taskyou","reason":"only host serving taskyou"} +``` + +`ty-on --help` prints the same summary; `ty-on --version` prints the version. + +## Environment + +| Variable | Default | Meaning | +| --- | --- | --- | +| `ON_HOSTS` | — | Host inventory path. Overrides the default lookup, and is passed through to `on ls` so both read the same file. | +| `XDG_CONFIG_HOME` | — | When set and `ON_HOSTS` is not, the inventory is `$XDG_CONFIG_HOME/on/hosts.yaml`. | +| `TY_ON_TIMEOUT` | `3s` | Budget for the `on ls` probe. Unset or unparseable falls back to the default. | + +## Development + +```console +$ go test ./... +$ golangci-lint run --config ../../.golangci.yml ./... +``` + +Tests inject a fake prober rather than shelling out, so the suite passes on a +machine with no fleet and no `on` installed. The `on ls` table parser is pinned +against real output. + +## Not in scope + +ty-on decides *where*; it does not move anything. Syncing the working tree, +creating worktrees, and opening SSH sessions are all `on`'s job, and invoking +this resolver is ty's. + +Host `capabilities` are parsed but not yet used for filtering — the rules above +are deliberately the whole policy. Matching an executor against a host's +capabilities is the obvious next lever. diff --git a/extensions/ty-on/cmd/main.go b/extensions/ty-on/cmd/main.go new file mode 100644 index 00000000..d81e3103 --- /dev/null +++ b/extensions/ty-on/cmd/main.go @@ -0,0 +1,106 @@ +// Command ty-on decides which host a TaskYou task should run on. +// +// It reads one JSON placement request on stdin and writes one JSON response on +// stdout: +// +// $ echo '{"event":"task.placement","task":{"project":"taskyou"}}' | ty-on +// {"target":"mona","workdir":"~/Projects/taskyou","reason":"only host serving taskyou"} +// +// An empty target means "run locally", and is the answer to every question this +// resolver cannot confidently answer. It exits 0 in all cases: it is called in +// the task spawn path and must never fail a task. +package main + +import ( + "context" + "encoding/json" + "fmt" + "io" + "os" + "time" + + "github.com/bborn/workflow/extensions/ty-on/internal/placement" +) + +// version is injected at build time via -ldflags "-X main.version=...". +var version = "dev" + +// maxRequest caps how much stdin we will read. A placement request is a few +// hundred bytes; anything near this is a malformed caller. +const maxRequest = 1 << 20 + +const usage = `ty-on — placement resolver for TaskYou + +Reads one JSON placement request on stdin, writes one JSON response on stdout. + + echo '{"event":"task.placement","task":{"project":"taskyou"}}' | ty-on + +Reads the same host inventory as the "on" CLI: $ON_HOSTS, else +$XDG_CONFIG_HOME/on/hosts.yaml, else ~/.config/on/hosts.yaml. + +Flags: + -h, --help show this help + -v, --version print the version + +Environment: + ON_HOSTS host inventory path + TY_ON_TIMEOUT budget for the "on ls" probe (default 3s) +` + +func main() { + for _, arg := range os.Args[1:] { + switch arg { + case "-h", "--help", "help": + fmt.Print(usage) + return + case "-v", "--version", "version": + fmt.Println(version) + return + } + } + + emit(resolve(context.Background(), os.Stdin)) +} + +// resolve turns whatever is on stdin into a placement response. Every failure +// mode becomes a local placement carrying an explanation. +func resolve(ctx context.Context, stdin io.Reader) placement.Response { + body, err := io.ReadAll(io.LimitReader(stdin, maxRequest)) + if err != nil { + return placement.Local("placement request could not be read: %v", err) + } + if len(body) == 0 { + return placement.Local("empty placement request") + } + + var req placement.Request + if err := json.Unmarshal(body, &req); err != nil { + return placement.Local("placement request is not valid JSON: %v", err) + } + + return placement.Resolver{Timeout: timeout()}.Resolve(ctx, req) +} + +// timeout reads the probe budget from TY_ON_TIMEOUT, falling back to the +// default when it is unset or nonsense. +func timeout() time.Duration { + raw := os.Getenv("TY_ON_TIMEOUT") + if raw == "" { + return placement.DefaultTimeout + } + d, err := time.ParseDuration(raw) + if err != nil || d <= 0 { + return placement.DefaultTimeout + } + return d +} + +func emit(resp placement.Response) { + out, err := json.Marshal(resp) + if err != nil { + // Response is three strings; this cannot fail in practice, but a + // hand-written fallback still beats writing nothing at all. + out = []byte(`{"target":"","workdir":"","reason":"placement response could not be encoded"}`) + } + fmt.Fprintf(os.Stdout, "%s\n", out) +} diff --git a/extensions/ty-on/cmd/main_test.go b/extensions/ty-on/cmd/main_test.go new file mode 100644 index 00000000..1990a2d7 --- /dev/null +++ b/extensions/ty-on/cmd/main_test.go @@ -0,0 +1,115 @@ +package main + +import ( + "context" + "encoding/json" + "path/filepath" + "strings" + "testing" + "time" + + "github.com/bborn/workflow/extensions/ty-on/internal/placement" +) + +// noInventory points the resolver at a path that does not exist, so these tests +// exercise the stdin/stdout contract without depending on a real fleet. +func noInventory(t *testing.T) { + t.Helper() + t.Setenv("ON_HOSTS", filepath.Join(t.TempDir(), "absent.yaml")) + t.Setenv("PATH", t.TempDir()) +} + +func TestResolveReadsTheRequestContract(t *testing.T) { + tests := []struct { + name string + stdin string + // wantReason is a substring the reason must contain. + wantReason string + }{ + { + name: "a well-formed request is understood", + stdin: `{"event":"task.placement","task":{"id":5225,"title":"Some task","project":"taskyou","repo_path":"/Users/bruno/Projects/workflow","executor":"claude"}}`, + wantReason: "no host inventory at", + }, + { + name: "malformed JSON falls back to local", + stdin: `{"event":"task.placement",`, + wantReason: "placement request is not valid JSON", + }, + { + name: "a JSON scalar falls back to local", + stdin: `"nope"`, + wantReason: "placement request is not valid JSON", + }, + { + name: "empty stdin falls back to local", + stdin: "", + wantReason: "empty placement request", + }, + { + name: "a request with no task falls back to local", + stdin: `{"event":"task.placement"}`, + wantReason: "task has no project", + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + noInventory(t) + + got := resolve(context.Background(), strings.NewReader(tc.stdin)) + + if got.Target != "" { + t.Errorf("target = %q, want a local placement", got.Target) + } + if !strings.Contains(got.Reason, tc.wantReason) { + t.Errorf("reason = %q, want it to contain %q", got.Reason, tc.wantReason) + } + }) + } +} + +// Whatever happens, the response must be one JSON object carrying all three +// fields — core parses it unconditionally. +func TestResolveAlwaysEncodesTheFullResponse(t *testing.T) { + noInventory(t) + + out, err := json.Marshal(resolve(context.Background(), strings.NewReader("garbage"))) + if err != nil { + t.Fatalf("marshal response: %v", err) + } + + var fields map[string]any + if err := json.Unmarshal(out, &fields); err != nil { + t.Fatalf("response is not a JSON object: %v (%s)", err, out) + } + for _, key := range []string{"target", "workdir", "reason"} { + if _, ok := fields[key]; !ok { + t.Errorf("response is missing %q: %s", key, out) + } + } + if fields["reason"] == "" { + t.Errorf("reason is empty: %s", out) + } +} + +func TestTimeout(t *testing.T) { + tests := []struct { + raw string + want time.Duration + }{ + {"", placement.DefaultTimeout}, + {"750ms", 750 * time.Millisecond}, + {"10s", 10 * time.Second}, + {"nonsense", placement.DefaultTimeout}, + {"0s", placement.DefaultTimeout}, + {"-5s", placement.DefaultTimeout}, + } + + for _, tc := range tests { + t.Setenv("TY_ON_TIMEOUT", tc.raw) + if got := timeout(); got != tc.want { + t.Errorf("TY_ON_TIMEOUT=%q: timeout() = %s, want %s", tc.raw, got, tc.want) + } + } +} diff --git a/extensions/ty-on/go.mod b/extensions/ty-on/go.mod new file mode 100644 index 00000000..316eed1b --- /dev/null +++ b/extensions/ty-on/go.mod @@ -0,0 +1,5 @@ +module github.com/bborn/workflow/extensions/ty-on + +go 1.24.4 + +require gopkg.in/yaml.v3 v3.0.1 diff --git a/extensions/ty-on/go.sum b/extensions/ty-on/go.sum new file mode 100644 index 00000000..a62c313c --- /dev/null +++ b/extensions/ty-on/go.sum @@ -0,0 +1,4 @@ +gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM= +gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= +gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= +gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= diff --git a/extensions/ty-on/internal/placement/inventory.go b/extensions/ty-on/internal/placement/inventory.go new file mode 100644 index 00000000..c7edb35b --- /dev/null +++ b/extensions/ty-on/internal/placement/inventory.go @@ -0,0 +1,89 @@ +package placement + +import ( + "fmt" + "os" + "path/filepath" + "sort" + + "gopkg.in/yaml.v3" +) + +// Inventory mirrors the parts of `on`'s hosts.yaml this resolver cares about. +// Unknown keys are ignored, so `on` can grow its schema without breaking us. +// +// See github.com/bborn/on for the full documented shape. +type Inventory struct { + // Repos maps a project name to its clone URL. Used by `on` when a host + // does not have the project yet; the resolver only reads Hosts. + Repos map[string]string `yaml:"repos"` + Hosts map[string]Host `yaml:"hosts"` +} + +// Host is one machine in the fleet. +type Host struct { + SSH string `yaml:"ssh"` + Workdir string `yaml:"workdir"` + Capabilities []string `yaml:"capabilities"` + // Repos maps a project name to that project's checkout path on this host. + Repos map[string]string `yaml:"repos"` +} + +// Candidate is a host that has a checkout of the project being placed. +type Candidate struct { + // Name is the inventory key, which is also the name `on` accepts. + Name string + Host Host + // Checkout is the project's path on the host, from the host's repos map. + Checkout string +} + +// InventoryPath resolves the inventory location the same way `on` does: +// ON_HOSTS wins, then $XDG_CONFIG_HOME/on/hosts.yaml, then ~/.config/on/hosts.yaml. +func InventoryPath() string { + if p := os.Getenv("ON_HOSTS"); p != "" { + return p + } + if dir := os.Getenv("XDG_CONFIG_HOME"); dir != "" { + return filepath.Join(dir, "on", "hosts.yaml") + } + home, err := os.UserHomeDir() + if err != nil { + // Nothing sensible to fall back to; the caller will report the miss. + return filepath.Join(".config", "on", "hosts.yaml") + } + return filepath.Join(home, ".config", "on", "hosts.yaml") +} + +// LoadInventory reads and parses the inventory at path. Its errors are written +// to be shown to a user verbatim as a placement reason. +func LoadInventory(path string) (*Inventory, error) { + data, err := os.ReadFile(path) //nolint:gosec // path is operator-controlled config, not user input + if err != nil { + if os.IsNotExist(err) { + return nil, fmt.Errorf("no host inventory at %s, nothing to place onto", path) + } + return nil, fmt.Errorf("host inventory at %s could not be read: %v", path, err) + } + + var inv Inventory + if err := yaml.Unmarshal(data, &inv); err != nil { + return nil, fmt.Errorf("host inventory at %s is not valid YAML: %v", path, err) + } + return &inv, nil +} + +// Serving returns the hosts that have a checkout of project, sorted by name so +// the answer does not depend on Go's map iteration order. +func (inv *Inventory) Serving(project string) []Candidate { + var out []Candidate + for name, host := range inv.Hosts { + checkout, ok := host.Repos[project] + if !ok || checkout == "" { + continue + } + out = append(out, Candidate{Name: name, Host: host, Checkout: checkout}) + } + sort.Slice(out, func(i, j int) bool { return out[i].Name < out[j].Name }) + return out +} diff --git a/extensions/ty-on/internal/placement/placement.go b/extensions/ty-on/internal/placement/placement.go new file mode 100644 index 00000000..0dde19d0 --- /dev/null +++ b/extensions/ty-on/internal/placement/placement.go @@ -0,0 +1,214 @@ +// Package placement decides which host a TaskYou task should run on. +// +// It is deliberately conservative: every path that cannot confidently name a +// host resolves to a local placement (empty target) with a reason explaining +// why. Nothing in here returns an error to the caller — the caller is in the +// task spawn path and must never be failed by a policy decision. +package placement + +import ( + "context" + "errors" + "fmt" + "sort" + "strings" + "time" +) + +// Event is the only event this resolver answers. +const Event = "task.placement" + +// Request is the JSON document ty writes to the resolver's stdin. +type Request struct { + Event string `json:"event"` + Task Task `json:"task"` +} + +// Task describes the task being placed. +type Task struct { + ID int64 `json:"id"` + Title string `json:"title"` + Project string `json:"project"` + RepoPath string `json:"repo_path"` + Executor string `json:"executor"` +} + +// Response is the JSON document the resolver writes to stdout. +// +// An empty Target means "run locally". Reason is always populated and is shown +// to the user, so it should be specific enough to explain a surprising choice. +type Response struct { + Target string `json:"target"` + Workdir string `json:"workdir"` + Reason string `json:"reason"` +} + +// Local builds a "run here" response with the given reason. +func Local(format string, args ...any) Response { + return Response{Reason: fmt.Sprintf(format, args...)} +} + +// DefaultTimeout bounds the `on ls` probe. Ranking several hosts costs an SSH +// round trip per host; past this we prefer a local placement to a slow answer. +const DefaultTimeout = 3 * time.Second + +// Resolver holds the knobs the resolution rules depend on. The zero value is +// usable: it reads the inventory `on` would read and probes with the default +// timeout. +type Resolver struct { + // InventoryPath overrides the inventory location. Empty means "ask the + // same environment `on` asks" (ON_HOSTS, XDG_CONFIG_HOME, ~/.config). + InventoryPath string + + // Timeout bounds the `on ls` probe. Zero means DefaultTimeout. + Timeout time.Duration + + // Prober ranks hosts by free memory. Nil means shell out to `on ls`. + Prober Prober +} + +// Resolve applies the placement rules to req and always returns a usable +// response. +func (r Resolver) Resolve(ctx context.Context, req Request) Response { + if req.Event != "" && req.Event != Event { + return Local("unsupported event %q, expected %q", req.Event, Event) + } + project := req.Task.Project + if project == "" { + return Local("task has no project, nothing to match against the host inventory") + } + + path := r.InventoryPath + if path == "" { + path = InventoryPath() + } + + inv, err := LoadInventory(path) + if err != nil { + return Local("%s", err) + } + + candidates := inv.Serving(project) + switch len(candidates) { + case 0: + return Local("no host in %s serves %s (%s)", path, project, hostSummary(inv)) + case 1: + c := candidates[0] + return Response{ + Target: c.Name, + Workdir: c.Checkout, + Reason: fmt.Sprintf("only host serving %s", project), + } + } + + return r.rank(ctx, project, path, candidates) +} + +// rank picks between two or more hosts that all serve the project. +func (r Resolver) rank(ctx context.Context, project, path string, candidates []Candidate) Response { + prober := r.Prober + if prober == nil { + prober = OnProber{InventoryPath: path} + } + + timeout := r.Timeout + if timeout <= 0 { + timeout = DefaultTimeout + } + ctx, cancel := context.WithTimeout(ctx, timeout) + defer cancel() + + stats, err := prober.Probe(ctx) + if err != nil { + // A late answer is worth less than a local one: the caller is waiting + // to spawn the executor. + if errors.Is(ctx.Err(), context.DeadlineExceeded) { + return Local("%d hosts serve %s but comparing them took longer than %s, so this task stays local", + len(candidates), project, timeout) + } + return Local("%d hosts serve %s but they could not be compared: %s", + len(candidates), project, err) + } + + // Keep only candidates the probe could actually reach and measure. + type ranked struct { + Candidate + freeKB int64 + } + var reachable []ranked + for _, c := range candidates { + s, ok := stats[c.Name] + if !ok || !s.Reachable { + continue + } + reachable = append(reachable, ranked{Candidate: c, freeKB: s.FreeKB}) + } + + switch len(reachable) { + case 0: + return Local("%d hosts serve %s but none of them are reachable (%s)", + len(candidates), project, names(candidates)) + case 1: + only := reachable[0] + return Response{ + Target: only.Name, + Workdir: only.Checkout, + Reason: fmt.Sprintf("only reachable host of %d serving %s (%s)", + len(candidates), project, names(candidates)), + } + } + + // Most free memory wins; ties break on host name so the answer is stable. + sort.SliceStable(reachable, func(i, j int) bool { + if reachable[i].freeKB != reachable[j].freeKB { + return reachable[i].freeKB > reachable[j].freeKB + } + return reachable[i].Name < reachable[j].Name + }) + + // Spell out what each contender had, so a surprising winner is explicable. + detail := make([]string, 0, len(reachable)) + for _, h := range reachable { + detail = append(detail, h.Name+" "+humanKB(h.freeKB)) + } + reason := fmt.Sprintf("most free memory of %d hosts serving %s (%s)", + len(reachable), project, strings.Join(detail, ", ")) + if skipped := len(candidates) - len(reachable); skipped > 0 { + reason += fmt.Sprintf("; %d unreachable", skipped) + } + + best := reachable[0] + return Response{Target: best.Name, Workdir: best.Checkout, Reason: reason} +} + +func names(candidates []Candidate) string { + out := make([]string, len(candidates)) + for i, c := range candidates { + out[i] = c.Name + } + return strings.Join(out, ", ") +} + +func hostSummary(inv *Inventory) string { + switch n := len(inv.Hosts); n { + case 0: + return "inventory lists no hosts" + case 1: + return "1 host in inventory" + default: + return fmt.Sprintf("%d hosts in inventory", n) + } +} + +// humanKB renders kilobytes the way `on ls` does, so a reason string can be +// checked against `on ls` output by eye. +func humanKB(kb int64) string { + switch { + case kb >= 1<<20: + return fmt.Sprintf("%.1fG", float64(kb)/(1<<20)) + case kb >= 1<<10: + return fmt.Sprintf("%dM", kb/(1<<10)) + default: + return fmt.Sprintf("%dK", kb) + } +} diff --git a/extensions/ty-on/internal/placement/placement_test.go b/extensions/ty-on/internal/placement/placement_test.go new file mode 100644 index 00000000..2d3fe5f4 --- /dev/null +++ b/extensions/ty-on/internal/placement/placement_test.go @@ -0,0 +1,466 @@ +package placement + +import ( + "context" + "errors" + "os" + "path/filepath" + "strings" + "testing" + "time" +) + +// fleet is an inventory where offerlab has exactly one host and taskyou has two, +// so the one-candidate and several-candidate rules can both be exercised. +const fleet = ` +repos: + taskyou: git@github.com:bborn/taskyou.git + +hosts: + ol-agents: + ssh: ol-agents + workdir: ~/projects + capabilities: [agent, ruby] + repos: + offerlab: ~/projects/engineering + + mona: + ssh: mona + workdir: ~/Projects + capabilities: [agent, docker] + repos: + taskyou: ~/Projects/taskyou + + rex: + ssh: rex + workdir: /root + capabilities: [agent] + repos: + taskyou: /root/taskyou +` + +// Free memory as `on ls` would report it, in the kilobytes HostStat carries. +const ( + mem11585M = 11585 * 1024 + mem2048M = 2048 * 1024 +) + +// writeInventory writes body to a temp file and points ON_HOSTS at it. +func writeInventory(t *testing.T, body string) string { + t.Helper() + path := filepath.Join(t.TempDir(), "hosts.yaml") + if err := os.WriteFile(path, []byte(body), 0o600); err != nil { + t.Fatalf("write inventory: %v", err) + } + t.Setenv("ON_HOSTS", path) + return path +} + +// stubProber stands in for `on ls` so no test needs the real CLI on PATH. +type stubProber struct { + stats map[string]HostStat + err error + // blocks makes Probe wait for the resolver's deadline instead of answering. + blocks bool + calls int +} + +func (p *stubProber) Probe(ctx context.Context) (map[string]HostStat, error) { + p.calls++ + if p.blocks { + <-ctx.Done() + return nil, ctx.Err() + } + return p.stats, p.err +} + +// unusedProber fails the test if the resolver reaches for a probe it should not +// need — the zero- and one-candidate rules must answer without measuring anything. +type unusedProber struct{ t *testing.T } + +func (p unusedProber) Probe(context.Context) (map[string]HostStat, error) { + p.t.Helper() + p.t.Error("probed the fleet for a placement that needs no comparison") + return nil, errors.New("should not be called") +} + +func reachable(kb int64) HostStat { return HostStat{Reachable: true, FreeKB: kb} } + +func unreachable() HostStat { return HostStat{} } + +func TestResolve(t *testing.T) { + tests := []struct { + name string + // inventory is written to a temp file and exposed via ON_HOSTS. Empty + // means "point ON_HOSTS at a path that does not exist". + inventory string + project string + event string + // prober stands in for `on ls`. Nil means the placement must not probe. + prober Prober + wantTarget string + wantWorkdir string + // wantReason is a substring the reason must contain. + wantReason string + }{ + { + name: "unknown project falls back to local", + inventory: fleet, + project: "influencekit", + wantTarget: "", + wantReason: "serves influencekit (3 hosts in inventory)", + }, + { + name: "single host serving the project wins outright", + inventory: fleet, + project: "offerlab", + wantTarget: "ol-agents", + wantWorkdir: "~/projects/engineering", + wantReason: "only host serving offerlab", + }, + { + name: "several hosts are ranked by free memory", + inventory: fleet, + project: "taskyou", + prober: &stubProber{stats: map[string]HostStat{ + "mona": reachable(mem11585M), + "rex": reachable(mem2048M), + // Not a candidate: roomiest host, but no taskyou checkout. + "ol-agents": reachable(26565 * 1024), + }}, + wantTarget: "mona", + wantWorkdir: "~/Projects/taskyou", + wantReason: "most free memory of 2 hosts serving taskyou (mona 11.3G, rex 2.0G)", + }, + { + name: "unreachable candidates are skipped", + inventory: fleet, + project: "taskyou", + prober: &stubProber{stats: map[string]HostStat{ + "mona": unreachable(), + "rex": reachable(mem2048M), + }}, + wantTarget: "rex", + wantWorkdir: "/root/taskyou", + wantReason: "only reachable host of 2 serving taskyou (mona, rex)", + }, + { + name: "no reachable candidate falls back to local", + inventory: fleet, + project: "taskyou", + prober: &stubProber{stats: map[string]HostStat{ + "mona": unreachable(), + "rex": unreachable(), + }}, + wantTarget: "", + wantReason: "none of them are reachable (mona, rex)", + }, + { + name: "a candidate the probe never mentions is skipped", + inventory: fleet, + project: "taskyou", + prober: &stubProber{stats: map[string]HostStat{"mona": reachable(mem11585M)}}, + wantTarget: "mona", + // rex was in the inventory but not in the probe: treated as absent. + wantWorkdir: "~/Projects/taskyou", + wantReason: "only reachable host of 2 serving taskyou", + }, + { + name: "missing inventory falls back to local", + inventory: "", + project: "taskyou", + wantTarget: "", + wantReason: "no host inventory at", + }, + { + name: "malformed inventory falls back to local", + inventory: "hosts: [this is not: a mapping\n", + project: "taskyou", + wantTarget: "", + wantReason: "is not valid YAML", + }, + { + name: "empty inventory falls back to local", + inventory: "hosts: {}\n", + project: "taskyou", + wantTarget: "", + wantReason: "inventory lists no hosts", + }, + { + name: "an unusable probe falls back to local", + inventory: fleet, + project: "taskyou", + prober: &stubProber{err: errors.New("the on CLI is not installed")}, + wantTarget: "", + wantReason: "2 hosts serve taskyou but they could not be compared: the on CLI is not installed", + }, + { + name: "a task without a project falls back to local", + inventory: fleet, + project: "", + wantTarget: "", + wantReason: "task has no project", + }, + { + name: "an unknown event falls back to local", + inventory: fleet, + project: "taskyou", + event: "task.started", + wantTarget: "", + wantReason: `unsupported event "task.started"`, + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + if tc.inventory == "" { + t.Setenv("ON_HOSTS", filepath.Join(t.TempDir(), "absent.yaml")) + } else { + writeInventory(t, tc.inventory) + } + + req := request(tc.project) + if tc.event != "" { + req.Event = tc.event + } + + prober := tc.prober + if prober == nil { + prober = unusedProber{t} + } + + got := Resolver{Prober: prober}.Resolve(context.Background(), req) + + if got.Target != tc.wantTarget { + t.Errorf("target = %q, want %q (reason: %s)", got.Target, tc.wantTarget, got.Reason) + } + if got.Workdir != tc.wantWorkdir { + t.Errorf("workdir = %q, want %q", got.Workdir, tc.wantWorkdir) + } + if got.Reason == "" { + t.Fatal("reason is empty; every placement must explain itself") + } + if !strings.Contains(got.Reason, tc.wantReason) { + t.Errorf("reason = %q, want it to contain %q", got.Reason, tc.wantReason) + } + }) + } +} + +// A slow probe must not hold up the spawn path: we prefer local to a late answer. +func TestResolveGivesUpOnASlowProbe(t *testing.T) { + writeInventory(t, fleet) + prober := &stubProber{blocks: true} + + start := time.Now() + got := Resolver{Prober: prober, Timeout: 50 * time.Millisecond}. + Resolve(context.Background(), request("taskyou")) + elapsed := time.Since(start) + + if got.Target != "" { + t.Errorf("target = %q, want a local placement", got.Target) + } + if want := "comparing them took longer than 50ms"; !strings.Contains(got.Reason, want) { + t.Errorf("reason = %q, want it to contain %q", got.Reason, want) + } + if elapsed > time.Second { + t.Errorf("took %s, want the probe to be abandoned promptly", elapsed) + } +} + +// Ties must not depend on Go's map iteration order. +func TestResolveBreaksMemoryTiesByName(t *testing.T) { + writeInventory(t, fleet) + prober := &stubProber{stats: map[string]HostStat{ + "mona": reachable(4096 * 1024), + "rex": reachable(4096 * 1024), + }} + + for i := 0; i < 10; i++ { + got := Resolver{Prober: prober}.Resolve(context.Background(), request("taskyou")) + if got.Target != "mona" { + t.Fatalf("run %d: target = %q, want the alphabetically first of the tied hosts", i, got.Target) + } + } +} + +// The single-candidate rule must answer from the inventory alone, so a fleet +// without `on` installed still places tasks. +func TestResolveDoesNotProbeForASingleCandidate(t *testing.T) { + writeInventory(t, fleet) + prober := &stubProber{stats: map[string]HostStat{"ol-agents": reachable(mem2048M)}} + + got := Resolver{Prober: prober}.Resolve(context.Background(), request("offerlab")) + + if got.Target != "ol-agents" { + t.Errorf("target = %q, want ol-agents (reason: %s)", got.Target, got.Reason) + } + if prober.calls != 0 { + t.Errorf("probed %d times, want 0", prober.calls) + } +} + +// `on` is an optional dependency. With it absent, the default prober must +// report that plainly and the task must stay local rather than fail. +func TestResolveWhenOnIsNotInstalled(t *testing.T) { + writeInventory(t, fleet) + // An empty PATH: exec.LookPath cannot find `on` anywhere. + t.Setenv("PATH", t.TempDir()) + + got := Resolver{}.Resolve(context.Background(), request("taskyou")) + + if got.Target != "" { + t.Errorf("target = %q, want a local placement", got.Target) + } + if want := "the on CLI is not installed"; !strings.Contains(got.Reason, want) { + t.Errorf("reason = %q, want it to contain %q", got.Reason, want) + } +} + +func TestOnProberReportsAMissingBinary(t *testing.T) { + t.Setenv("PATH", t.TempDir()) + + _, err := OnProber{}.Probe(context.Background()) + + if err == nil { + t.Fatal("Probe() succeeded, want an error naming the missing CLI") + } + if want := "the on CLI is not installed"; !strings.Contains(err.Error(), want) { + t.Errorf("err = %q, want it to contain %q", err, want) + } +} + +func TestServingIsSortedAndIgnoresBlankCheckouts(t *testing.T) { + inv, err := LoadInventory(writeInventory(t, ` +hosts: + zeta: + repos: {taskyou: /srv/taskyou} + alpha: + repos: {taskyou: /home/alpha/taskyou} + blank: + repos: {taskyou: ""} + other: + repos: {offerlab: /srv/offerlab} +`)) + if err != nil { + t.Fatalf("load inventory: %v", err) + } + + got := inv.Serving("taskyou") + want := []string{"alpha", "zeta"} + if len(got) != len(want) { + t.Fatalf("serving = [%s], want %v", names(got), want) + } + for i, name := range want { + if got[i].Name != name { + t.Errorf("serving[%d] = %q, want %q", i, got[i].Name, name) + } + } +} + +func TestInventoryPathPrefersONHOSTS(t *testing.T) { + t.Setenv("ON_HOSTS", "/tmp/custom-hosts.yaml") + if got := InventoryPath(); got != "/tmp/custom-hosts.yaml" { + t.Errorf("InventoryPath() = %q, want the ON_HOSTS value", got) + } + + t.Setenv("ON_HOSTS", "") + t.Setenv("XDG_CONFIG_HOME", "/xdg") + if got, want := InventoryPath(), filepath.Join("/xdg", "on", "hosts.yaml"); got != want { + t.Errorf("InventoryPath() = %q, want %q", got, want) + } +} + +// The real `on ls` table, verbatim, so the parser is pinned to the actual +// output shape rather than to a tidied-up version of it. +const lsTable = `HOST SSH CORES AVAIL TOTAL LOAD +mona mona 4 11585M 15887M 0.02 72% free +ol-agents ol-agents 16 26565M 31337M 0.12 84% free +down down - - - - ssh: Could not resolve hostname down: nodename nor servname provided +` + +func TestParseOnLS(t *testing.T) { + stats := ParseOnLS(lsTable) + + if len(stats) != 3 { + t.Fatalf("parsed %d rows, want 3: %v", len(stats), stats) + } + if s := stats["ol-agents"]; !s.Reachable || s.FreeKB != 26565*1024 { + t.Errorf("ol-agents = %+v, want reachable with 26565M free", s) + } + if s := stats["mona"]; !s.Reachable || s.FreeKB != mem11585M { + t.Errorf("mona = %+v, want reachable with 11585M free", s) + } + if s := stats["down"]; s.Reachable { + t.Errorf("down = %+v, want unreachable", s) + } + if _, ok := stats["HOST"]; ok { + t.Error("the header row was parsed as a host") + } +} + +func TestParseOnLSIgnoresJunk(t *testing.T) { + if stats := ParseOnLS(""); len(stats) != 0 { + t.Errorf("empty output parsed as %v, want no hosts", stats) + } + if stats := ParseOnLS("on: no inventory found\n"); len(stats) != 0 { + t.Errorf("short line parsed as %v, want no hosts", stats) + } +} + +func TestParseSize(t *testing.T) { + tests := []struct { + in string + want int64 + wantOK bool + }{ + {"8777M", 8777 * 1024, true}, + {"512K", 512, true}, + {"2G", 2 * 1024 * 1024, true}, + {"1.5G", 1536 * 1024, true}, + {"1T", 1024 * 1024 * 1024, true}, + {"4096", 4096, true}, + {"-", 0, false}, + {"", 0, false}, + {"ssh:", 0, false}, + {"-1M", 0, false}, + } + + for _, tc := range tests { + got, ok := parseSize(tc.in) + if ok != tc.wantOK || got != tc.want { + t.Errorf("parseSize(%q) = (%d, %t), want (%d, %t)", tc.in, got, ok, tc.want, tc.wantOK) + } + } +} + +func TestHumanKB(t *testing.T) { + tests := []struct { + in int64 + want string + }{ + {11585 * 1024, "11.3G"}, + {2048 * 1024, "2.0G"}, + {512 * 1024, "512M"}, + {900, "900K"}, + } + + for _, tc := range tests { + if got := humanKB(tc.in); got != tc.want { + t.Errorf("humanKB(%d) = %q, want %q", tc.in, got, tc.want) + } + } +} + +func request(project string) Request { + return Request{ + Event: Event, + Task: Task{ + ID: 5225, + Title: "Some task", + Project: project, + RepoPath: "/Users/bruno/Projects/workflow", + Executor: "claude", + }, + } +} diff --git a/extensions/ty-on/internal/placement/probe.go b/extensions/ty-on/internal/placement/probe.go new file mode 100644 index 00000000..7c49d712 --- /dev/null +++ b/extensions/ty-on/internal/placement/probe.go @@ -0,0 +1,141 @@ +package placement + +import ( + "bytes" + "context" + "fmt" + "os" + "os/exec" + "strconv" + "strings" +) + +// HostStat is what a probe learned about one host. +type HostStat struct { + // Reachable is false when the probe could not measure the host at all. + Reachable bool + // FreeKB is available memory in kilobytes. Only meaningful when Reachable. + FreeKB int64 +} + +// Prober measures the fleet so two candidate hosts can be compared. +type Prober interface { + // Probe returns a stat per host name, keyed the way the inventory keys it. + Probe(ctx context.Context) (map[string]HostStat, error) +} + +// OnProber shells out to the `on` CLI, which already knows how to probe the +// fleet in parallel. `on` is an optional dependency: if it is not installed we +// report that rather than reimplementing the probe. +type OnProber struct { + // InventoryPath is passed through as ON_HOSTS so `on` reads exactly the + // inventory this resolver read. + InventoryPath string + + // Binary overrides the executable name. Empty means "on". + Binary string +} + +// Probe runs `on ls` and parses its table. +func (p OnProber) Probe(ctx context.Context) (map[string]HostStat, error) { + bin := p.Binary + if bin == "" { + bin = "on" + } + + path, err := exec.LookPath(bin) + if err != nil { + return nil, fmt.Errorf("the %s CLI is not installed", bin) + } + + cmd := exec.CommandContext(ctx, path, "ls") //nolint:gosec // fixed argv, path from LookPath + if p.InventoryPath != "" { + cmd.Env = append(os.Environ(), "ON_HOSTS="+p.InventoryPath) + } + var stdout, stderr bytes.Buffer + cmd.Stdout = &stdout + cmd.Stderr = &stderr + + if err := cmd.Run(); err != nil { + // A timeout surfaces as a kill here; the caller inspects ctx and turns + // that into the "stays local" reason, so just report what we saw. + if msg := firstLine(stderr.String()); msg != "" { + return nil, fmt.Errorf("%s ls failed: %s", bin, msg) + } + return nil, fmt.Errorf("%s ls failed: %v", bin, err) + } + + stats := ParseOnLS(stdout.String()) + if len(stats) == 0 { + return nil, fmt.Errorf("%s ls reported no hosts", bin) + } + return stats, nil +} + +// ParseOnLS reads the table `on ls` prints: +// +// HOST SSH CORES AVAIL TOTAL LOAD +// ol-agents ol-agents 16 26565M 31337M 0.12 84% free +// rex rex - - - - ssh: ... +// +// A row whose AVAIL column is missing or unparseable (a dash, an SSH error) is +// recorded as unreachable rather than dropped, so callers can tell "the host is +// down" apart from "the host is not in the fleet". +func ParseOnLS(out string) map[string]HostStat { + stats := make(map[string]HostStat) + for _, line := range strings.Split(out, "\n") { + fields := strings.Fields(line) + if len(fields) < 4 || fields[0] == "HOST" || !isCount(fields[2]) { + // Not a host row: the header, a blank line, or a stray message. + continue + } + name := fields[0] + freeKB, ok := parseSize(fields[3]) + stats[name] = HostStat{Reachable: ok, FreeKB: freeKB} + } + return stats +} + +// isCount reports whether s is the CORES cell of a host row: a core count, or +// a dash when the host could not be reached. Anything else means the line is +// not part of the table. +func isCount(s string) bool { + if s == "-" { + return true + } + _, err := strconv.Atoi(s) + return err == nil +} + +// parseSize reads a size cell such as "8777M", "26.5G" or "-" into kilobytes. +func parseSize(s string) (int64, bool) { + if s == "" || s == "-" { + return 0, false + } + + mult := int64(1) + switch s[len(s)-1] { + case 'K', 'k': + s = s[:len(s)-1] + case 'M', 'm': + mult, s = 1<<10, s[:len(s)-1] + case 'G', 'g': + mult, s = 1<<20, s[:len(s)-1] + case 'T', 't': + mult, s = 1<<30, s[:len(s)-1] + } + + n, err := strconv.ParseFloat(s, 64) + if err != nil || n < 0 { + return 0, false + } + return int64(n * float64(mult)), true +} + +func firstLine(s string) string { + s = strings.TrimSpace(s) + if i := strings.IndexByte(s, '\n'); i >= 0 { + s = s[:i] + } + return s +}