From fcd3e386a93a07e791d2c638b62b383b0f3605c0 Mon Sep 17 00:00:00 2001 From: Bruno Bornsztein Date: Mon, 31 Aug 2026 14:28:16 -0500 Subject: [PATCH] feat(ty-on): placement resolver deciding which host a task runs on MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ty is gaining the ability to run a task on another machine. The policy for that — which hosts exist, what they are provisioned for, which one to pick — is fleet-specific and must not live in core, so it ships as a standalone extension with its own module. Nothing outside extensions/ty-on is touched. The binary reads one JSON placement request on stdin and writes one response on stdout. An empty target means "run locally" and is the answer to every question it cannot confidently answer: unknown project, missing or malformed inventory, no reachable host, malformed request, `on` not installed. It never fails a task and never guesses a host — it exits 0 in all cases, and the required reason field explains the choice to the user. Placement reads the same inventory as the `on` CLI (ON_HOSTS, else XDG_CONFIG_HOME, else ~/.config/on/hosts.yaml): no host serving the project means local, one means that host, several are ranked by free memory. Ranking shells out to `on ls` rather than reimplementing the probe, treating `on` as an optional dependency. Only the multi-host rule probes at all, so the common paths answer from a single file read; the probe is bounded by TY_ON_TIMEOUT (3s) because a late answer is worth less than a local one in the spawn path. Tests are table-driven over the resolution rules with the inventory in a temp file via ON_HOSTS, and inject a fake prober so the suite needs no fleet and no `on` on PATH. The `on ls` table parser is pinned against real output. Co-Authored-By: Claude Opus 5 (1M context) --- extensions/ty-on/.gitignore | 6 + extensions/ty-on/README.md | 131 +++++ extensions/ty-on/cmd/main.go | 106 ++++ extensions/ty-on/cmd/main_test.go | 115 +++++ extensions/ty-on/go.mod | 5 + extensions/ty-on/go.sum | 4 + .../ty-on/internal/placement/inventory.go | 89 ++++ .../ty-on/internal/placement/placement.go | 214 ++++++++ .../internal/placement/placement_test.go | 466 ++++++++++++++++++ extensions/ty-on/internal/placement/probe.go | 141 ++++++ 10 files changed, 1277 insertions(+) create mode 100644 extensions/ty-on/.gitignore create mode 100644 extensions/ty-on/README.md create mode 100644 extensions/ty-on/cmd/main.go create mode 100644 extensions/ty-on/cmd/main_test.go create mode 100644 extensions/ty-on/go.mod create mode 100644 extensions/ty-on/go.sum create mode 100644 extensions/ty-on/internal/placement/inventory.go create mode 100644 extensions/ty-on/internal/placement/placement.go create mode 100644 extensions/ty-on/internal/placement/placement_test.go create mode 100644 extensions/ty-on/internal/placement/probe.go 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 +}