From 10dae9b60ed7ad1de9bd9988af01552c997d32c8 Mon Sep 17 00:00:00 2001 From: KoLDXr00T Date: Tue, 20 Oct 2026 20:09:52 +0100 Subject: [PATCH] router: one MCP server over local and remote box hosts, with compact output box mcp fronts the hosts in ~/.box/hosts.yaml besides the local box. A host is reached only over ssh: box __ensure-serve starts its own box serve, and one ssh forward per host brings that owner-only socket here, into a 0700 directory. No network listener is opened. A sandbox there is host/name; a bare name is the local one, else the only host with it. Where several remote hosts have it, up and run are placed (where it is up or warm, then the least loaded under max_up); anything acting on one copy's state must name the host. An unreachable host is reported, not fatal, and retried after a pause; after a quiet spell a host is checked before use, since its box serve exits when idle or upgraded. box hosts shows them. The MCP tools spend fewer of the agent's tokens: - command output loses escapes and progress redraws, collapses repeated lines, and is cut to its first and last lines (8 KiB a stream by default, --output-budget); the rest is kept in memory for a new output tool, by line range or regex - read_file, diff and list_sandboxes answer "unchanged" when the fresh result is identical to what the agent already holds; every read still reaches the sandbox, so it is never stale, and full: true shows it again - the sandbox list is one line per sandbox instead of indented JSON Host lists are validated before ssh sees them: no destination ssh could read as an option, no shell metacharacters in the remote binary path. --- README.md | 12 +- SECURITY.md | 16 ++ internal/cli/hosts.go | 161 ++++++++++++ internal/cli/root.go | 4 +- internal/cli/up.go | 22 +- internal/mcp/compact.go | 271 ++++++++++++++++++++ internal/mcp/compact_test.go | 211 +++++++++++++++ internal/mcp/mcp.go | 180 ++++++++++--- internal/mcp/mcp_test.go | 2 +- internal/mcp/seen.go | 48 ++++ internal/router/hosts.go | 396 ++++++++++++++++++++++++++++ internal/router/hosts_test.go | 253 ++++++++++++++++++ internal/router/router.go | 453 +++++++++++++++++++++++++++++++++ internal/router/router_test.go | 194 ++++++++++++++ 14 files changed, 2176 insertions(+), 47 deletions(-) create mode 100644 internal/cli/hosts.go create mode 100644 internal/mcp/compact.go create mode 100644 internal/mcp/compact_test.go create mode 100644 internal/mcp/seen.go create mode 100644 internal/router/hosts.go create mode 100644 internal/router/hosts_test.go create mode 100644 internal/router/router.go create mode 100644 internal/router/router_test.go diff --git a/README.md b/README.md index 36c7898..6f01a3e 100644 --- a/README.md +++ b/README.md @@ -61,7 +61,17 @@ box apply try-a # or: box discard try-a - **SDKs:** [Python](sdk/python/), [TypeScript](sdk/typescript/), [Go](sdk/go/), [Rust](sdk/rust/) - **MCP:** `claude mcp add box -- box mcp`. Agents can fork but not - apply, unless you start it with `--allow-apply`. + apply, unless you start it with `--allow-apply`. Long output is shown by + its first and last lines, with the rest one `output` call away, and a + re-read of something unchanged costs one line, not the content. +- **Other machines:** list them in `~/.box/hosts.yaml` and `box mcp` fronts + them too, over ssh; a sandbox there is `host/name`. `box hosts` shows them. + +```yaml +hosts: + - name: gpu1 + ssh: me@gpu1.example.com +``` ```python from sdbox import Sandbox diff --git a/SECURITY.md b/SECURITY.md index fdebee8..2c9e576 100644 --- a/SECURITY.md +++ b/SECURITY.md @@ -124,6 +124,22 @@ agent or another sandbox's. The SDK server listens on a Unix socket created owner-only (`umask` before `listen`), and is never placed in a shared directory. +## Remote hosts (`hosts.yaml`) + +`box mcp` can front other machines listed in `~/.box/hosts.yaml`. It reaches +each one only over ssh: it runs `box __ensure-serve` there, which starts +that host's own `box serve` on its owner-only socket, and forwards that socket +to one in a `0700` directory here. box opens no network listener for this. +ssh is the boundary: whoever can log in to the host can run commands in its +sandboxes, as whoever can open the socket can locally. ssh runs in batch +mode and checks host keys as your ssh config says; box never relaxes that. + +The host list is checked before ssh sees it: a destination must be +`user@host` or a config alias, never something ssh would read as an option, +and the remote binary path may hold no character a shell gives meaning to. +An agent connected to `box mcp` can use every sandbox on every listed host, +and `--allow-apply` lets it apply forks on all of them. + ## Host-side handling of guest-written files A guest can write anything into `/data`: symlinks to `/etc/passwd`, `..` diff --git a/internal/cli/hosts.go b/internal/cli/hosts.go new file mode 100644 index 0000000..bf61919 --- /dev/null +++ b/internal/cli/hosts.go @@ -0,0 +1,161 @@ +package cli + +import ( + "encoding/json" + "errors" + "fmt" + "net" + "net/http" + "net/http/httptest" + "os" + "os/exec" + "syscall" + "text/tabwriter" + "time" + + "github.com/spf13/cobra" + + "box/internal/router" + "box/internal/sandbox" + "box/internal/server" +) + +// remoteIdle is how long a box serve started for a router waits unused +// before exiting. The router checks after a quiet spell and starts it again. +const remoteIdle = 30 * time.Minute + +// apiFor is the API box mcp and box hosts speak through: the local one, or a +// router over it and the hosts in the host list. A missing default host +// list means local only; a missing file named with --hosts is an error. +func apiFor(hostsFile string, localOnly bool) (http.Handler, func(), error) { + local := server.NewHandler(Version) + if localOnly { + return local, func() {}, nil + } + path, explicit := hostsFile, hostsFile != "" + if !explicit { + p, err := router.ConfigPath() + if err != nil { + return nil, nil, err + } + path = p + } + cfg, err := router.LoadConfig(path) + if errors.Is(err, os.ErrNotExist) && !explicit { + return local, func() {}, nil + } + if err != nil { + return nil, nil, err + } + rt, closeAll, err := router.Open(Version, local, cfg) + if err != nil { + return nil, nil, err + } + return rt, closeAll, nil +} + +func hostsCmd() *cobra.Command { + var hostsFile string + c := &cobra.Command{ + Use: "hosts", Short: "the remote hosts box mcp reaches", GroupID: groupInspect, + Long: "Lists the hosts in ~/.box/hosts.yaml, connects to each and shows\n" + + "what it has. box mcp fronts all of them: a sandbox on another host is\n" + + "named host/name. A host is reached over ssh, so whoever can log in\n" + + "there can run commands in its sandboxes.\n\n" + + " hosts:\n" + + " - name: gpu1\n" + + " ssh: me@gpu1.example.com # user@host or an ssh config alias\n" + + " max_up: 8 # optional: placement keeps at most 8 up\n" + + " box: /usr/local/bin/box # optional: box on that host", + Args: cobra.NoArgs, + ValidArgsFunction: completeNothing, + RunE: func(*cobra.Command, []string) error { + api, closeAll, err := apiFor(hostsFile, false) + if err != nil { + return err + } + defer closeAll() + rec := httptest.NewRecorder() + api.ServeHTTP(rec, httptest.NewRequest("GET", "/v1/hosts", nil)) + if rec.Code == http.StatusNotFound { + fmt.Println("no hosts configured; add them to ~/.box/hosts.yaml (see box hosts --help)") + return nil + } + var hosts []struct { + Name, Version, Error string + Sandboxes, Up, Warm int + MaxUp int `json:"max_up"` + } + if err := json.Unmarshal(rec.Body.Bytes(), &hosts); err != nil { + return err + } + tw := tabwriter.NewWriter(os.Stdout, 0, 4, 2, ' ', 0) + fmt.Fprintln(tw, "HOST\tSANDBOXES\tUP\tWARM\tVERSION\tSTATUS") + down := 0 + for _, h := range hosts { + status := "ok" + if h.Error != "" { + status, down = "unreachable: "+h.Error, down+1 + } + up := fmt.Sprint(h.Up) + if h.MaxUp > 0 { + up += fmt.Sprintf("/%d", h.MaxUp) + } + fmt.Fprintf(tw, "%s\t%d\t%s\t%d\t%s\t%s\n", h.Name, h.Sandboxes, up, h.Warm, h.Version, status) + } + tw.Flush() + if down > 0 { + exitCode = 1 + return fmt.Errorf("%s unreachable", plural(down, "host", "hosts")) + } + return nil + }, + } + c.Flags().StringVar(&hostsFile, "hosts", "", "host list (default ~/.box/hosts.yaml)") + return c +} + +// ensureServeCmd starts box serve if it is not running and prints its socket +// path. A router runs it on a remote host over ssh, then forwards that path. +func ensureServeCmd() *cobra.Command { + return &cobra.Command{ + Use: "__ensure-serve", Hidden: true, + Args: cobra.NoArgs, + ValidArgsFunction: completeNothing, + RunE: func(*cobra.Command, []string) error { + sock, err := sandbox.SocketPath() + if err != nil { + return err + } + if !listening(sock) { + exe, err := os.Executable() + if err != nil { + return err + } + // No stdio: ssh waits for every holder of its stdout to close. + cmd := exec.Command(exe, "serve", "--socket", sock, "--idle", remoteIdle.String()) + cmd.SysProcAttr = &syscall.SysProcAttr{Setsid: true} // outlive this ssh session + if err := cmd.Start(); err != nil { + return err + } + cmd.Process.Release() + for i := 0; i < 250 && !listening(sock); i++ { + time.Sleep(20 * time.Millisecond) + } + if !listening(sock) { + return fmt.Errorf("started box serve but it never listened on %s", sock) + } + } + fmt.Println(sock) + return nil + }, + } +} + +func listening(sock string) bool { + c, err := net.Dial("unix", sock) + if err == nil { + c.Close() + } + return err == nil +} diff --git a/internal/cli/root.go b/internal/cli/root.go index 14c9ca5..47bb4e3 100644 --- a/internal/cli/root.go +++ b/internal/cli/root.go @@ -59,13 +59,13 @@ func newRoot() *cobra.Command { root.AddCommand( newCmd(), editCmd(), buildCmd(), lsCmd(), - runCmd(), shellCmd(), upCmd(), execCmd(), downCmd(), serveCmd(), mcpCmd(), + runCmd(), shellCmd(), upCmd(), execCmd(), downCmd(), serveCmd(), mcpCmd(), hostsCmd(), resetCmd(), snapshotCmd(), restoreCmd(), forkCmd(), diffCmd(), applyCmd(), discardCmd(), dataCmd(), envCmd(), logsCmd(), verifyCmd(), inspectCmd(), doctorCmd(), renameCmd(), destroyCmd(), nukeCmd(), - agentCmd(), prepareCmd(), tendCmd(), watchDataCmd(), overlayCmd(), krunCmd(), exportCmd(), netnsCmd(), fcCmd(), tarOutCmd(), tarInCmd(), + agentCmd(), prepareCmd(), tendCmd(), watchDataCmd(), overlayCmd(), krunCmd(), exportCmd(), netnsCmd(), fcCmd(), tarOutCmd(), tarInCmd(), ensureServeCmd(), ) return root } diff --git a/internal/cli/up.go b/internal/cli/up.go index 7529aa5..bd77dc2 100644 --- a/internal/cli/up.go +++ b/internal/cli/up.go @@ -186,23 +186,39 @@ func doctorCmd() *cobra.Command { } func mcpCmd() *cobra.Command { - var allowApply bool + var allowApply, localOnly bool + var hostsFile string + var budget int c := &cobra.Command{ Use: "mcp", Short: "MCP server for AI agents (stdio)", GroupID: groupRun, Long: "Serves box over the Model Context Protocol on stdin and stdout,\n" + "for Claude Code, Claude Desktop, Cursor and other MCP clients. Tools:\n" + "list_sandboxes, run, up, exec, down, read_file, write_file, fork,\n" + - "diff, discard. Sandboxes themselves are created and built with the CLI.\n\n" + + "diff, discard, output. Sandboxes themselves are created and built\n" + + "with the CLI.\n\n" + + "With hosts in ~/.box/hosts.yaml (see box hosts), it fronts those too:\n" + + "a sandbox on another host is host/name.\n\n" + + "Long command output is shown by its first and last lines, the rest\n" + + "kept for the output tool, and a read whose result the agent already\n" + + "has answers \"unchanged\": both save the agent's context.\n\n" + "apply is not offered unless --allow-apply is given: forks exist so a\n" + "human reviews an agent's work before it reaches real data.\n\n" + " claude mcp add box -- box mcp", Args: cobra.NoArgs, ValidArgsFunction: completeNothing, RunE: func(*cobra.Command, []string) error { - return mcp.New(server.NewHandler(Version), Version, allowApply).Serve(os.Stdin, os.Stdout) + api, closeAll, err := apiFor(hostsFile, localOnly) + if err != nil { + return err + } + defer closeAll() + return mcp.New(api, Version, mcp.Options{AllowApply: allowApply, OutputBudget: budget}).Serve(os.Stdin, os.Stdout) }, } c.Flags().BoolVar(&allowApply, "allow-apply", false, "also let the agent apply forks to their parents") + c.Flags().StringVar(&hostsFile, "hosts", "", "host list (default ~/.box/hosts.yaml)") + c.Flags().BoolVar(&localOnly, "local-only", false, "ignore the host list") + c.Flags().IntVar(&budget, "output-budget", mcp.DefaultOutputBudget, "bytes of each command stream shown before the middle is elided") return c } diff --git a/internal/mcp/compact.go b/internal/mcp/compact.go new file mode 100644 index 0000000..670d4a4 --- /dev/null +++ b/internal/mcp/compact.go @@ -0,0 +1,271 @@ +package mcp + +import ( + "bytes" + "fmt" + "regexp" + "strings" + "sync" + "unicode/utf8" +) + +// DefaultOutputBudget is how much of one stream a tool result shows before +// the middle is elided. Command output is mostly read for its first and last +// lines (what ran, how it ended), so the cut keeps both ends. +const DefaultOutputBudget = 8 << 10 + +// spoolLimit bounds the full outputs a session keeps for the output tool. +// They stay in memory only: command output can hold secrets. +const spoolLimit = 64 << 20 + +// ansiRe matches terminal escapes: CSI sequences (colours, cursor moves), +// OSC sequences (titles, hyperlinks) and the remaining two-byte escapes. +var ansiRe = regexp.MustCompile(`\x1b\[[0-?]*[ -/]*[@-~]|\x1b\][^\x07\x1b]*(?:\x07|\x1b\\)|\x1b[@-Z\\-_]`) + +// clean turns raw command output into what a reader of a terminal would have +// seen: escapes removed, carriage-return redraws (progress bars) resolved to +// their final state, and runs of identical lines collapsed. +func clean(b []byte) string { + if bytes.IndexByte(b, 0) >= 0 { + return fmt.Sprintf("[%d bytes of binary output not shown]", len(b)) + } + s := ansiRe.ReplaceAllString(strings.ToValidUTF8(string(b), "�"), "") + lines := strings.Split(s, "\n") + out := make([]string, 0, len(lines)) + for i := 0; i < len(lines); { + line := redraw(lines[i]) + j := i + 1 + for j < len(lines) && redraw(lines[j]) == line { + j++ + } + // Collapsing costs a marker line, so only runs of three or more. + if n := j - i; n >= 3 { + out = append(out, line, fmt.Sprintf("[previous line repeated %d more times]", n-1)) + } else { + for k := i; k < j; k++ { + out = append(out, line) + } + } + i = j + } + return strings.Join(out, "\n") +} + +// redraw keeps what is left on a line after carriage returns redraw it. +func redraw(line string) string { + line = strings.TrimSuffix(line, "\r") + if i := strings.LastIndexByte(line, '\r'); i >= 0 { + return line[i+1:] + } + return line +} + +// clipped is a stream cut to a budget: the text shown and, when something +// was elided, which lines. +type clipped struct { + text string + elided bool + from, to, of int // 1-based lines elided, of the total + bytesNotShown int +} + +// clipEnds keeps about two thirds of budget from the start of s and one third +// from the end, cutting at line boundaries. +func clipEnds(s string, budget int) clipped { + if len(s) <= budget { + return clipped{text: s} + } + lines := strings.SplitAfter(s, "\n") + if lines[len(lines)-1] == "" { + lines = lines[:len(lines)-1] // s ended with a newline, not with an empty line + } + headBudget, tailBudget := budget*2/3, budget-budget*2/3 + var head, tail []string + used := 0 + for _, l := range lines { + if used+len(l) > headBudget { + break + } + head = append(head, l) + used += len(l) + } + used = 0 + for i := len(lines) - 1; i >= len(head); i-- { + if used+len(lines[i]) > tailBudget { + break + } + tail = append([]string{lines[i]}, tail...) + used += len(lines[i]) + } + // One enormous line: show its two ends rather than nothing. + if len(head) == 0 && len(tail) == 0 { + h, t := runeCut(s, headBudget, false), runeCut(s, tailBudget, true) + return clipped{text: h + "\n" + markerPlaceholder + "\n" + t, elided: true, + from: 1, to: 1, of: 1, bytesNotShown: len(s) - len(h) - len(t)} + } + shown := strings.Join(head, "") + strings.Join(tail, "") + c := clipped{elided: true, from: len(head) + 1, to: len(lines) - len(tail), of: len(lines), + bytesNotShown: len(s) - len(shown)} + c.text = strings.Join(head, "") + markerPlaceholder + "\n" + strings.Join(tail, "") + return c +} + +// markerPlaceholder stands where the elision marker goes; the marker names +// the spooled output, which is only known once the stream is stored. +const markerPlaceholder = "\x00elided\x00" + +// runeCut takes n bytes from the start (or end) of s without splitting a rune. +func runeCut(s string, n int, fromEnd bool) string { + if fromEnd { + i := len(s) - n + for i < len(s) && !utf8.RuneStart(s[i]) { + i++ + } + return s[i:] + } + for n > 0 && n < len(s) && !utf8.RuneStart(s[n]) { + n-- + } + return s[:n] +} + +// spool keeps the cleaned full text of outputs that were clipped, so the +// model can ask for a part it was not shown. Oldest entries go first. +type spool struct { + mu sync.Mutex + next int + order []string + entries map[string]map[string]string // id -> stream -> text + size int +} + +func newSpool() *spool { return &spool{entries: map[string]map[string]string{}} } + +func (p *spool) put(streams map[string]string) string { + p.mu.Lock() + defer p.mu.Unlock() + p.next++ + id := fmt.Sprintf("o%d", p.next) + p.entries[id] = streams + p.order = append(p.order, id) + for _, s := range streams { + p.size += len(s) + } + for p.size > spoolLimit && len(p.order) > 1 { + old := p.order[0] + p.order = p.order[1:] + for _, s := range p.entries[old] { + p.size -= len(s) + } + delete(p.entries, old) + } + return id +} + +func (p *spool) get(id, stream string) (string, bool) { + p.mu.Lock() + defer p.mu.Unlock() + e, ok := p.entries[id] + if !ok { + return "", false + } + s, ok := e[stream] + return s, ok +} + +// section is one rendered stream. +type section struct{ name, text string } + +// render cleans (when cleanIt) and clips each named stream to budget. When +// anything was cut, the full text of every stream is spooled and the cut is +// marked with the id that retrieves it. A file is not cleaned, since the +// model may edit what it is shown; command output is. +func (s *Server) render(names []string, raw map[string][]byte, cleanIt bool, budget int) []section { + full := map[string]string{} + clips := map[string]clipped{} + spoolIt := false + for _, n := range names { + if len(raw[n]) == 0 { + continue + } + if cleanIt { + full[n] = clean(raw[n]) + } else { + full[n] = strings.ToValidUTF8(string(raw[n]), "\uFFFD") + } + clips[n] = clipEnds(full[n], budget) + spoolIt = spoolIt || clips[n].elided + } + id := "" + if spoolIt { + id = s.spool.put(full) + } + var out []section + for _, n := range names { + c, ok := clips[n] + if !ok { + continue + } + text := c.text + if c.elided { + marker := fmt.Sprintf("[… lines %d-%d of %d (%d bytes) not shown; see them with output id=%s stream=%s]", + c.from, c.to, c.of, c.bytesNotShown, id, n) + text = strings.Replace(text, markerPlaceholder, marker, 1) + } + out = append(out, section{n, text}) + } + return out +} + +// output answers the output tool: numbered lines of a spooled text, a range +// or those matching a pattern, clipped like anything else. +func (s *Server) output(a map[string]any) (string, bool) { + stream := str(a, "stream") + if stream == "" { + stream = "stdout" + } + text, ok := s.spool.get(str(a, "id"), stream) + if !ok { + return fmt.Sprintf("no output id=%s stream=%s: it was never elided, or is too old and was dropped", str(a, "id"), stream), true + } + lines := strings.SplitAfter(text, "\n") + if lines[len(lines)-1] == "" { + lines = lines[:len(lines)-1] + } + from, to := num(a, "from_line"), num(a, "to_line") + if from < 1 { + from = 1 + } + if to < 1 || to > len(lines) { + to = len(lines) + } + var re *regexp.Regexp + if g := str(a, "grep"); g != "" { + var err error + if re, err = regexp.Compile(g); err != nil { + return "bad grep pattern: " + err.Error(), true + } + } + var b strings.Builder + matched := 0 + for i := from; i <= to; i++ { + l := lines[i-1] + if re != nil && !re.MatchString(strings.TrimSuffix(l, "\n")) { + continue + } + matched++ + fmt.Fprintf(&b, "%d: %s", i, l) + if !strings.HasSuffix(l, "\n") { + b.WriteString("\n") + } + } + if matched == 0 { + return fmt.Sprintf("no lines (of %d) match", len(lines)), false + } + c := clipEnds(b.String(), maxText) + if !c.elided { + return c.text, false + } + return strings.Replace(c.text, markerPlaceholder, + fmt.Sprintf("[… %d of these lines not shown; narrow the range or the pattern]", c.to-c.from+1), 1), false +} diff --git a/internal/mcp/compact_test.go b/internal/mcp/compact_test.go new file mode 100644 index 0000000..3b00662 --- /dev/null +++ b/internal/mcp/compact_test.go @@ -0,0 +1,211 @@ +package mcp + +import ( + "bytes" + "encoding/base64" + "encoding/json" + "fmt" + "io" + "net/http" + "strings" + "testing" +) + +func TestClean(t *testing.T) { + for in, want := range map[string]string{ + "\x1b[31mred\x1b[0m": "red", + "\x1b]0;title\x07text": "text", + "10%\r50%\r100%\ndone": "100%\ndone", + "crlf\r\nline": "crlf\nline", + "a\na\nb": "a\na\nb", + "x\nx\nx\nx\ny": "x\n[previous line repeated 3 more times]\ny", + "\x1b[2K\rDownloading 1/3\r\x1b[2K\rDownloading 3/3\n": "Downloading 3/3\n", + } { + if got := clean([]byte(in)); got != want { + t.Errorf("clean(%q) = %q, want %q", in, got, want) + } + } + if got := clean([]byte("ELF\x00\x01")); !strings.Contains(got, "binary") { + t.Errorf("binary output should be summarised: %q", got) + } +} + +func TestClipEndsKeepsBothEnds(t *testing.T) { + var b strings.Builder + for i := 1; i <= 1000; i++ { + fmt.Fprintf(&b, "line %d\n", i) + } + c := clipEnds(b.String(), 600) + if !c.elided || !strings.HasPrefix(c.text, "line 1\n") || !strings.HasSuffix(c.text, "line 1000\n") { + t.Fatalf("head and tail should survive: %q", c.text) + } + if c.of != 1000 || c.from < 2 || c.to > 999 || c.from > c.to { + t.Errorf("elided range %d-%d of %d", c.from, c.to, c.of) + } + if len(c.text) > 600+len(markerPlaceholder)+1 { + t.Errorf("over budget: %d bytes", len(c.text)) + } + if c := clipEnds("short\n", 600); c.elided || c.text != "short\n" { + t.Errorf("short text must pass through: %+v", c) + } + // A single huge line still shows both ends, without splitting a rune. + huge := strings.Repeat("é", 2000) + c = clipEnds(huge, 100) + if !c.elided || !strings.HasPrefix(c.text, "é") || !strings.HasSuffix(c.text, "é") { + t.Errorf("one long line: %q", c.text) + } +} + +// execAPI answers exec with whatever stdout the test sets, and remembers +// what write_file sent. +type execAPI struct { + stdout map[string]string // by argv[0], e.g. "cat", "sh" + written string + calls int +} + +func (f *execAPI) ServeHTTP(w http.ResponseWriter, r *http.Request) { + f.calls++ + var body struct { + Argv []string `json:"argv"` + Stdin string `json:"stdin"` + } + json.NewDecoder(r.Body).Decode(&body) + if r.URL.Path == "/v1/sandboxes" { + w.Write([]byte(`[{"name":"ok","up":true,"warm":0,"warm_target":0,"network":"none","base":"alpine"}]`)) + return + } + if r.URL.Path == "/v1/hosts" { + w.Write([]byte(`[{"name":"local"},{"name":"gpu1","error":"ssh: connection refused"}]`)) + return + } + if len(body.Argv) > 0 && body.Argv[0] == "sh" && strings.Contains(body.Argv[2], "cat >") { + in, _ := base64.StdEncoding.DecodeString(body.Stdin) + f.written = string(in) + } + out := "" + if len(body.Argv) > 0 { + out = f.stdout[body.Argv[0]] + } + fmt.Fprintf(w, `{"exit_code":0,"stdout":%q,"stderr":"","duration_ms":1}`, base64.StdEncoding.EncodeToString([]byte(out))) +} + +// calls runs tools/call requests against api and returns each result's text. +func calls(t *testing.T, api http.Handler, opt Options, reqs ...string) []string { + t.Helper() + var lines []string + for i, r := range reqs { + lines = append(lines, fmt.Sprintf(`{"jsonrpc":"2.0","id":%d,"method":"tools/call","params":%s}`, i+1, r)) + } + var out bytes.Buffer + if err := New(api, "test", opt).Serve(strings.NewReader(strings.Join(lines, "\n")), &out); err != nil { + t.Fatal(err) + } + var texts []string + dec := json.NewDecoder(&out) + for dec.More() { + var m struct { + Result struct { + Content []struct{ Text string } `json:"content"` + } `json:"result"` + } + if err := dec.Decode(&m); err != nil && err != io.EOF { + t.Fatal(err) + } + texts = append(texts, m.Result.Content[0].Text) + } + return texts +} + +func TestLongOutputIsElidedAndRetrievable(t *testing.T) { + var b strings.Builder + for i := 1; i <= 5000; i++ { + fmt.Fprintf(&b, "row %d\n", i) + } + api := &execAPI{stdout: map[string]string{"sh": b.String()}} + got := calls(t, api, Options{OutputBudget: 1024}, + `{"name":"exec","arguments":{"sandbox":"ok","command":"seq"}}`, + `{"name":"output","arguments":{"id":"o1","grep":"^row 4242$"}}`, + `{"name":"output","arguments":{"id":"o1","from_line":10,"to_line":11}}`, + `{"name":"output","arguments":{"id":"o9"}}`, + ) + if !strings.Contains(got[0], "row 1\n") || !strings.Contains(got[0], "row 5000") || + !strings.Contains(got[0], "output id=o1 stream=stdout") || len(got[0]) > 1400 { + t.Errorf("exec result should keep both ends and name the spool (%d bytes):\n%s", len(got[0]), got[0]) + } + if got[1] != "4242: row 4242\n" { + t.Errorf("grep = %q", got[1]) + } + if got[2] != "10: row 10\n11: row 11\n" { + t.Errorf("range = %q", got[2]) + } + if !strings.Contains(got[3], "no output id=o9") { + t.Errorf("unknown id = %q", got[3]) + } +} + +func TestRereadSaysUnchanged(t *testing.T) { + api := &execAPI{stdout: map[string]string{"cat": "print('hi')\n"}} + read := `{"name":"read_file","arguments":{"sandbox":"ok","path":"/data/a.py"}}` + got := calls(t, api, Options{}, + read, + read, + `{"name":"read_file","arguments":{"sandbox":"ok","path":"/data/a.py","full":true}}`, + `{"name":"read_file","arguments":{"sandbox":"ok","path":"/data/b.py"}}`, + ) + if got[0] != "print('hi')\n" { + t.Errorf("first read = %q", got[0]) + } + if !strings.HasPrefix(got[1], "unchanged:") { + t.Errorf("second read should be unchanged: %q", got[1]) + } + if got[2] != "print('hi')\n" { + t.Errorf("full: true must show the content: %q", got[2]) + } + if got[3] != "print('hi')\n" { + t.Errorf("another path is not the same file: %q", got[3]) + } + if api.calls != 4 { + t.Errorf("every read must reach the sandbox; %d of 4 did", api.calls) + } +} + +// After the model writes a file, a read that finds what it wrote is +// unchanged; one that finds anything else is shown. +func TestWriteThenRead(t *testing.T) { + api := &execAPI{stdout: map[string]string{"cat": "v1\n"}} + got := calls(t, api, Options{}, + `{"name":"read_file","arguments":{"sandbox":"ok","path":"/data/f"}}`, + `{"name":"write_file","arguments":{"sandbox":"ok","path":"/data/f","content":"v2\n"}}`, + // The fake still serves v1: the file was reverted behind the model's back. + `{"name":"read_file","arguments":{"sandbox":"ok","path":"/data/f"}}`, + ) + if api.written != "v2\n" { + t.Fatalf("write_file sent %q", api.written) + } + if got[2] != "v1\n" { + t.Errorf("the model last held v2, so v1 must be shown, got %q", got[2]) + } + api2 := &execAPI{stdout: map[string]string{"cat": "v2\n"}} + got = calls(t, api2, Options{}, + `{"name":"write_file","arguments":{"sandbox":"ok","path":"/data/f","content":"v2\n"}}`, + `{"name":"read_file","arguments":{"sandbox":"ok","path":"/data/f"}}`, + ) + if !strings.HasPrefix(got[1], "unchanged:") { + t.Errorf("reading back what was written should be unchanged: %q", got[1]) + } +} + +func TestListIsCompactAndNamesUnreachableHosts(t *testing.T) { + got := calls(t, &execAPI{}, Options{}, + `{"name":"list_sandboxes","arguments":{}}`, + `{"name":"list_sandboxes","arguments":{}}`, + ) + if !strings.Contains(got[0], "ok: up, warm 0/0, network none, base alpine") || + !strings.Contains(got[0], "host gpu1 unreachable") || strings.Contains(got[0], "local unreachable") { + t.Errorf("list = %q", got[0]) + } + if !strings.HasPrefix(got[1], "unchanged:") { + t.Errorf("an identical list should be unchanged: %q", got[1]) + } +} diff --git a/internal/mcp/mcp.go b/internal/mcp/mcp.go index fa2a8cb..94dea22 100644 --- a/internal/mcp/mcp.go +++ b/internal/mcp/mcp.go @@ -23,7 +23,9 @@ import ( // ProtocolVersion is the MCP revision this server speaks. const ProtocolVersion = "2025-06-18" -// maxText bounds what one tool result hands back to the model. +// maxText bounds how much of a file read_file shows; past it the middle is +// elided and can be fetched with the output tool. Files get more room than +// command output, since the model reads them to edit them. const maxText = 64 << 10 type request struct { @@ -50,10 +52,22 @@ type Server struct { api http.Handler version string allowApply bool + budget int + spool *spool + seen *seen tools []tool toolsByName map[string]tool } +// Options configures a session. +type Options struct { + // AllowApply offers the apply tool. + AllowApply bool + // OutputBudget is how many bytes of each command stream a result shows + // before eliding the middle; 0 means DefaultOutputBudget. + OutputBudget int +} + type tool struct { Name string `json:"name"` Description string `json:"description"` @@ -61,11 +75,16 @@ type tool struct { call func(s *Server, args map[string]any) (string, bool) } -// New returns a server answering through api, the local API's handler. -func New(api http.Handler, version string, allowApply bool) *Server { - s := &Server{api: api, version: version, allowApply: allowApply, toolsByName: map[string]tool{}} +// New returns a server answering through api: the local API's handler, or a +// router in front of several hosts' APIs. +func New(api http.Handler, version string, opt Options) *Server { + if opt.OutputBudget <= 0 { + opt.OutputBudget = DefaultOutputBudget + } + s := &Server{api: api, version: version, allowApply: opt.AllowApply, budget: opt.OutputBudget, + spool: newSpool(), seen: newSeen(), toolsByName: map[string]tool{}} for _, t := range allTools() { - if t.Name == "apply" && !allowApply { + if t.Name == "apply" && !opt.AllowApply { continue } s.tools = append(s.tools, t) @@ -111,7 +130,10 @@ func (s *Server) handle(req request) (any, *rpcError) { "instructions": "box runs commands in disposable microVM sandboxes, each with its own kernel. " + "Sandboxes are created and built by a human with the box CLI; list them with list_sandboxes. " + "`run` uses a fresh VM per command; `up` then `exec` keeps one VM and its state between commands. " + - "To try changes safely, `fork` a sandbox and work in the fork; a human reviews `diff` and applies it.", + "To try changes safely, `fork` a sandbox and work in the fork; a human reviews `diff` and applies it. " + + "A sandbox on another host is named host/name. " + + "Long command output is shown by its first and last lines; fetch the rest with `output`. " + + "A read whose result you already hold says `unchanged`; pass full: true to see it again.", }, nil case "ping": return map[string]any{}, nil @@ -143,7 +165,7 @@ func (s *Server) handle(req request) (any, *rpcError) { // call makes one request to the API in-process and decodes its JSON. func (s *Server) call(method, path string, body any) (map[string]any, []any, error) { - var rd io.Reader + var rd io.Reader = http.NoBody // as a real server hands a handler: never nil if body != nil { b, _ := json.Marshal(body) rd = bytes.NewReader(b) @@ -191,31 +213,33 @@ func num(args map[string]any, k string) int { return 0 } -func clip(b []byte) string { - if len(b) > maxText { - return string(b[:maxText]) + fmt.Sprintf("\n[... %d more bytes not shown]", len(b)-maxText) - } - return string(b) +func flag(args map[string]any, k string) bool { + v, _ := args[k].(bool) + return v } -// commandResult renders an exec/run response for the model. A non-zero exit -// is a result, not a tool failure, so it is not marked as an error. -func commandResult(r map[string]any) string { - dec := func(k string) []byte { - s, _ := r[k].(string) - b, _ := base64.StdEncoding.DecodeString(s) - return b - } +func b64(r map[string]any, k string) []byte { + s, _ := r[k].(string) + b, _ := base64.StdEncoding.DecodeString(s) + return b +} + +// commandResult renders an exec/run response for the model: escapes and +// progress redraws removed, repeats collapsed, each stream cut to the budget +// with its full text kept for the output tool. A non-zero exit is a result, +// not a tool failure, so it is not marked as an error. +func (s *Server) commandResult(r map[string]any) string { var b strings.Builder fmt.Fprintf(&b, "exit code %v in %v ms", r["exit_code"], r["duration_ms"]) if t, _ := r["timed_out"].(bool); t { b.WriteString(" (timed out)") } - if out := dec("stdout"); len(out) > 0 { - b.WriteString("\n\nstdout:\n" + clip(out)) + if t, _ := r["truncated"].(bool); t { + b.WriteString(" (output over the server's limit was dropped)") } - if errb := dec("stderr"); len(errb) > 0 { - b.WriteString("\n\nstderr:\n" + clip(errb)) + raw := map[string][]byte{"stdout": b64(r, "stdout"), "stderr": b64(r, "stderr")} + for _, sec := range s.render([]string{"stdout", "stderr"}, raw, true, s.budget) { + b.WriteString("\n\n" + sec.name + ":\n" + sec.text) } return b.String() } @@ -225,7 +249,8 @@ func schema(props map[string]any, required ...string) map[string]any { } var ( - pSandbox = map[string]any{"type": "string", "description": "sandbox name, from list_sandboxes"} + pSandbox = map[string]any{"type": "string", "description": "sandbox name from list_sandboxes (host/name for another host)"} + pFull = map[string]any{"type": "boolean", "description": "show the content even if you already saw it unchanged"} pCommand = map[string]any{"type": "string", "description": "a shell command line, run with sh -c"} pTimeout = map[string]any{"type": "integer", "description": "seconds; the command exits 124 when it is hit"} ) @@ -244,7 +269,7 @@ func allTools() []tool { if err != nil { return err.Error(), true } - return commandResult(r), false + return s.commandResult(r), false } } simple := func(method, verb, done string) func(*Server, map[string]any) (string, bool) { @@ -262,15 +287,23 @@ func allTools() []tool { return []tool{ { Name: "list_sandboxes", - Description: "List the sandboxes you can use, whether each is up, and its warm pool.", - InputSchema: schema(map[string]any{}), - call: func(s *Server, _ map[string]any) (string, bool) { + Description: "List the sandboxes you can use, on every host, whether each is up, and its warm pool.", + InputSchema: schema(map[string]any{"full": pFull}), + call: func(s *Server, a map[string]any) (string, bool) { _, arr, err := s.call("GET", "/v1/sandboxes", nil) if err != nil { return err.Error(), true } - b, _ := json.MarshalIndent(arr, "", " ") - return string(b), false + text := listText(arr) + // A router also reports hosts it could not reach; a single + // host's API has no such route, and the error is ignored. + if _, hosts, err := s.call("GET", "/v1/hosts", nil); err == nil { + text += hostsText(hosts) + } + if s.seen.check(listKey, []byte(text), flag(a, "full")) { + return unchangedNote("the sandbox list", len(text)), false + } + return text, false }, }, { @@ -316,8 +349,9 @@ func allTools() []tool { return err.Error(), true } if code, _ := r["exit_code"].(float64); code != 0 { - return commandResult(r), true + return s.commandResult(r), true } + s.seen.know(fileKey(str(a, "sandbox"), str(a, "path")), []byte(str(a, "content"))) return "written " + str(a, "path"), false }, }, @@ -325,7 +359,8 @@ func allTools() []tool { Name: "read_file", Description: "Read a file from inside a running sandbox (see up).", InputSchema: schema(map[string]any{"sandbox": pSandbox, - "path": map[string]any{"type": "string", "description": "absolute path in the guest"}}, "sandbox", "path"), + "path": map[string]any{"type": "string", "description": "absolute path in the guest"}, + "full": pFull}, "sandbox", "path"), call: func(s *Server, a map[string]any) (string, bool) { r, _, err := s.call("POST", "/v1/sandboxes/"+str(a, "sandbox")+"/exec", map[string]any{"argv": []string{"cat", "--", str(a, "path")}}) @@ -333,10 +368,17 @@ func allTools() []tool { return err.Error(), true } if code, _ := r["exit_code"].(float64); code != 0 { - return commandResult(r), true + return s.commandResult(r), true + } + out := b64(r, "stdout") + if s.seen.check(fileKey(str(a, "sandbox"), str(a, "path")), out, flag(a, "full")) { + return unchangedNote(str(a, "path"), len(out)), false + } + secs := s.render([]string{"content"}, map[string][]byte{"content": out}, false, maxText) + if len(secs) == 0 { + return "(empty file)", false } - out, _ := base64.StdEncoding.DecodeString(r["stdout"].(string)) - return clip(out), false + return secs[0].text, false }, }, { @@ -346,17 +388,22 @@ func allTools() []tool { InputSchema: schema(map[string]any{"sandbox": pSandbox, "name": map[string]any{"type": "string", "description": "name for the fork"}}, "sandbox", "name"), call: func(s *Server, a map[string]any) (string, bool) { - _, _, err := s.call("POST", "/v1/sandboxes/"+str(a, "sandbox")+"/fork", map[string]any{"as": str(a, "name")}) + r, _, err := s.call("POST", "/v1/sandboxes/"+str(a, "sandbox")+"/fork", map[string]any{"as": str(a, "name")}) if err != nil { return err.Error(), true } - return "forked " + str(a, "sandbox") + " to " + str(a, "name"), false + // A router names a fork on another host host/name. + fork, _ := r["fork"].(string) + if fork == "" { + fork = str(a, "name") + } + return "forked " + str(a, "sandbox") + " to " + fork, false }, }, { Name: "diff", Description: "List a fork's changes to /data: A added, M modified, D deleted.", - InputSchema: schema(map[string]any{"sandbox": map[string]any{"type": "string", "description": "the fork"}}, "sandbox"), + InputSchema: schema(map[string]any{"sandbox": map[string]any{"type": "string", "description": "the fork"}, "full": pFull}, "sandbox"), call: func(s *Server, a map[string]any) (string, bool) { r, _, err := s.call("GET", "/v1/sandboxes/"+str(a, "sandbox")+"/diff", nil) if err != nil { @@ -371,7 +418,11 @@ func allTools() []tool { m, _ := c.(map[string]any) fmt.Fprintf(&b, "%v %v\n", m["kind"], m["path"]) } - return b.String(), false + if s.seen.check(diffKey(str(a, "sandbox")), []byte(b.String()), flag(a, "full")) { + return unchangedNote("the diff of "+str(a, "sandbox"), b.Len()), false + } + secs := s.render([]string{"changes"}, map[string][]byte{"changes": []byte(b.String())}, false, maxText) + return secs[0].text, false }, }, { @@ -380,6 +431,19 @@ func allTools() []tool { InputSchema: schema(map[string]any{"sandbox": map[string]any{"type": "string", "description": "the fork"}}, "sandbox"), call: simple("POST", "discard", "discarded"), }, + { + Name: "output", + Description: "Show part of a command output or file that an earlier result elided, by its id: " + + "a range of lines, or the lines matching a regular expression. Lines are numbered.", + InputSchema: schema(map[string]any{ + "id": map[string]any{"type": "string", "description": "the id from the elision marker, e.g. o3"}, + "stream": map[string]any{"type": "string", "description": "stdout, stderr, content or changes, as the marker says; default stdout"}, + "from_line": map[string]any{"type": "integer", "description": "first line, 1-based"}, + "to_line": map[string]any{"type": "integer", "description": "last line, inclusive"}, + "grep": map[string]any{"type": "string", "description": "only lines matching this regular expression (RE2)"}, + }, "id"), + call: func(s *Server, a map[string]any) (string, bool) { return s.output(a) }, + }, { Name: "apply", Description: "Merge a fork's changes into its parent's /data. Refused while either is up.", @@ -388,3 +452,39 @@ func allTools() []tool { }, } } + +// listText renders the sandbox list one line per sandbox: the same facts as +// the API's JSON at a fraction of the tokens. +func listText(arr []any) string { + if len(arr) == 0 { + return "no sandboxes; a human creates them with: box new " + } + var b strings.Builder + for _, it := range arr { + m, _ := it.(map[string]any) + if e, _ := m["error"].(string); e != "" { + fmt.Fprintf(&b, "%v: error: %s\n", m["name"], e) + continue + } + state := "down" + if up, _ := m["up"].(bool); up { + state = "up" + } + fmt.Fprintf(&b, "%v: %s, warm %v/%v, network %v, base %v\n", + m["name"], state, m["warm"], m["warm_target"], m["network"], m["base"]) + } + return b.String() +} + +// hostsText names the hosts a router could not reach, whose sandboxes are +// missing from the list. +func hostsText(arr []any) string { + var b strings.Builder + for _, it := range arr { + m, _ := it.(map[string]any) + if e, _ := m["error"].(string); e != "" { + fmt.Fprintf(&b, "host %v unreachable, its sandboxes are not listed: %s\n", m["name"], e) + } + } + return b.String() +} diff --git a/internal/mcp/mcp_test.go b/internal/mcp/mcp_test.go index afeb859..771bd7b 100644 --- a/internal/mcp/mcp_test.go +++ b/internal/mcp/mcp_test.go @@ -28,7 +28,7 @@ func fakeAPI() http.Handler { func session(t *testing.T, allowApply bool, lines ...string) []map[string]any { t.Helper() var out bytes.Buffer - if err := New(fakeAPI(), "test", allowApply).Serve(strings.NewReader(strings.Join(lines, "\n")), &out); err != nil { + if err := New(fakeAPI(), "test", Options{AllowApply: allowApply}).Serve(strings.NewReader(strings.Join(lines, "\n")), &out); err != nil { t.Fatal(err) } var res []map[string]any diff --git a/internal/mcp/seen.go b/internal/mcp/seen.go new file mode 100644 index 0000000..9f8631a --- /dev/null +++ b/internal/mcp/seen.go @@ -0,0 +1,48 @@ +package mcp + +import ( + "crypto/sha256" + "fmt" + "sync" +) + +// seen remembers, per session, a hash of the last content the model holds +// for a key (a file, a fork's diff, the sandbox list). It caches no data: +// every read still goes to the sandbox, and only when the fresh result is +// identical to what the model already holds is it answered with a short note +// in place of the content. A stale answer is impossible by construction. +type seen struct { + mu sync.Mutex + last map[string][32]byte +} + +func newSeen() *seen { return &seen{last: map[string][32]byte{}} } + +// check reports whether content is what the model last held for key, and +// records it as held. full asks for the content regardless. +func (s *seen) check(key string, content []byte, full bool) (unchanged bool) { + sum := sha256.Sum256(content) + s.mu.Lock() + defer s.mu.Unlock() + prev, ok := s.last[key] + s.last[key] = sum + return ok && prev == sum && !full +} + +// know records content the model supplied itself, such as a file it wrote: +// it holds that content without having been shown it. +func (s *seen) know(key string, content []byte) { + s.mu.Lock() + defer s.mu.Unlock() + s.last[key] = sha256.Sum256(content) +} + +func fileKey(sandbox, path string) string { return "file\x00" + sandbox + "\x00" + path } +func diffKey(sandbox string) string { return "diff\x00" + sandbox } + +const listKey = "list" + +func unchangedNote(what string, n int) string { + return fmt.Sprintf("unchanged: %s is identical to the last time you saw it (%d bytes). "+ + "Pass full: true to see it again.", what, n) +} diff --git a/internal/router/hosts.go b/internal/router/hosts.go new file mode 100644 index 0000000..65be842 --- /dev/null +++ b/internal/router/hosts.go @@ -0,0 +1,396 @@ +package router + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "fmt" + "io" + "net" + "net/http" + "os" + "os/exec" + "path/filepath" + "regexp" + "strings" + "sync" + "time" + + "github.com/goccy/go-yaml" + + "box/internal/sandbox" +) + +// Config is ~/.box/hosts.yaml: the remote hosts a router fronts besides the +// local one. +type Config struct { + Hosts []HostConfig `yaml:"hosts"` +} + +// HostConfig is one remote host. It is reached over SSH, which is the whole +// access control: whoever can log in can run commands in its sandboxes, as +// whoever can open the local socket can here. Socket is for a box API whose +// socket is already reachable from this machine (forwarded by hand, or a +// second box home in tests). +type HostConfig struct { + Name string `yaml:"name"` + SSH string `yaml:"ssh"` + Socket string `yaml:"socket"` + Box string `yaml:"box"` // the box binary on the remote host; default "box" + MaxUp int `yaml:"max_up"` +} + +// ConfigPath is where the host list lives. +func ConfigPath() (string, error) { + h, err := sandbox.Home() + if err != nil { + return "", err + } + return filepath.Join(h, "hosts.yaml"), nil +} + +var ( + // An ssh destination: a user@host or an ~/.ssh/config alias. Never an + // option, which ssh would otherwise parse ahead of the "--". + sshTargetRe = regexp.MustCompile(`^[A-Za-z0-9_][A-Za-z0-9._@-]*$`) + // ssh runs the remote command through the remote user's shell, so the + // binary path is restricted to characters no shell gives meaning to. + remotePathRe = regexp.MustCompile(`^/?[A-Za-z0-9._/-]+$`) + // A socket path also appears in -L local:remote, so no colon either. + socketPathRe = regexp.MustCompile(`^/[A-Za-z0-9._/-]+$`) +) + +// LoadConfig reads and checks a host list. +func LoadConfig(path string) (Config, error) { + b, err := os.ReadFile(path) + if err != nil { + return Config{}, err + } + var c Config + if err := decodeYAML(b, &c); err != nil { + return Config{}, fmt.Errorf("%s: %w", path, err) + } + if err := c.validate(); err != nil { + return Config{}, fmt.Errorf("%s: %w", path, err) + } + return c, nil +} + +// decodeYAML is strict, and turns the YAML library's panics on malformed +// input into errors, as boxfile.Parse does. +func decodeYAML(b []byte, v any) (err error) { + defer func() { + if r := recover(); r != nil { + err = fmt.Errorf("cannot parse: %v", r) + } + }() + if len(bytes.TrimSpace(b)) == 0 { + return nil + } + dec := yaml.NewDecoder(bytes.NewReader(b), yaml.DisallowUnknownField()) + if err := dec.Decode(v); err != nil && !errors.Is(err, io.EOF) { + return err + } + return nil +} + +func (c Config) validate() error { + seen := map[string]bool{Local: true} + for i, h := range c.Hosts { + if err := sandbox.ValidName(h.Name); err != nil { + return fmt.Errorf("hosts[%d]: host name: %w", i, err) + } + if seen[h.Name] { + if h.Name == Local { + return fmt.Errorf("hosts[%d]: %q is the name of this machine", i, Local) + } + return fmt.Errorf("hosts[%d]: %q is listed twice", i, h.Name) + } + seen[h.Name] = true + switch { + case (h.SSH == "") == (h.Socket == ""): + return fmt.Errorf("host %s: give exactly one of ssh and socket", h.Name) + case h.SSH != "" && !sshTargetRe.MatchString(h.SSH): + return fmt.Errorf("host %s: ssh must be user@host or an ssh config alias, got %q", h.Name, h.SSH) + case h.Socket != "" && !socketPathRe.MatchString(h.Socket): + return fmt.Errorf("host %s: socket must be an absolute path, got %q", h.Name, h.Socket) + case h.Box != "" && !remotePathRe.MatchString(h.Box): + return fmt.Errorf("host %s: box must be a plain path, got %q", h.Name, h.Box) + case h.MaxUp < 0: + return fmt.Errorf("host %s: max_up must not be negative", h.Name) + } + } + return nil +} + +// sshCommand is the ssh client; a variable so tests can stand in for it. +var sshCommand = "ssh" + +// sshOpts apply to every ssh the router starts. BatchMode: never prompt, +// since stdin and stdout belong to the MCP client. Host keys are checked as +// the user's ssh config says; nothing here relaxes that. +var sshOpts = []string{"-o", "BatchMode=yes", "-o", "ConnectTimeout=10"} + +func ensureArgs(h HostConfig) []string { + box := h.Box + if box == "" { + box = "box" + } + return append(append([]string{}, sshOpts...), "--", h.SSH, box, "__ensure-serve") +} + +func forwardArgs(h HostConfig, local, remote string) []string { + return append(append([]string{}, sshOpts...), + "-N", "-o", "ExitOnForwardFailure=yes", "-o", "ServerAliveInterval=30", + "-o", "StreamLocalBindUnlink=yes", + "-L", local+":"+remote, "--", h.SSH) +} + +// Remote is the API of one configured host. It connects on first use, keeps +// one ssh forward per host, and reconnects when the forward or the remote +// server has gone away. +type Remote struct { + cfg HostConfig + dir string // owner-only directory for the forwarded socket + + mu sync.Mutex + client *http.Client + fwd *exec.Cmd + fwdDone chan struct{} + lastOK time.Time + retryAt time.Time + lastErr error + closed bool +} + +// pingAfter is how long a connection is trusted without a check. A remote +// box serve exits after 30 idle minutes, or sooner when its binary is +// upgraded, so a request after a quiet spell checks first. +const pingAfter = 60 * time.Second + +// retryAfter spaces out reconnection attempts to a host that is down. +const retryAfter = 15 * time.Second + +// NewRemote returns the API for h. dir holds forwarded sockets and must be +// private to this user. +func NewRemote(h HostConfig, dir string) *Remote { return &Remote{cfg: h, dir: dir} } + +func (rm *Remote) ServeHTTP(w http.ResponseWriter, r *http.Request) { + client, err := rm.ready(r.Context()) + if err != nil { + fail(w, http.StatusBadGateway, "unreachable", "host %s: %v", rm.cfg.Name, err) + return + } + out, err := http.NewRequestWithContext(r.Context(), r.Method, "http://box"+r.URL.Path, r.Body) + if err != nil { + fail(w, http.StatusBadRequest, "bad_request", "%v", err) + return + } + out.Header.Set("Content-Type", "application/json") + resp, err := client.Do(out) + if err != nil { + rm.mu.Lock() + rm.lastOK = time.Time{} // check before trusting it again + rm.mu.Unlock() + fail(w, http.StatusBadGateway, "unreachable", "host %s: %v", rm.cfg.Name, err) + return + } + defer resp.Body.Close() + for k, v := range resp.Header { + w.Header()[k] = v + } + w.WriteHeader(resp.StatusCode) + io.Copy(w, resp.Body) + rm.mu.Lock() + rm.lastOK = time.Now() + rm.mu.Unlock() +} + +// ready returns a client known to reach the host's box serve recently. +func (rm *Remote) ready(ctx context.Context) (*http.Client, error) { + rm.mu.Lock() + defer rm.mu.Unlock() + if rm.closed { + return nil, errors.New("router closed") + } + if rm.client != nil && rm.alive() { + if time.Since(rm.lastOK) < pingAfter { + return rm.client, nil + } + if rm.ping(ctx) == nil { + return rm.client, nil + } + // The forward is up but nothing answers behind it: the remote server + // exited. Start it again; the forward reaches the same socket path. + if rm.cfg.SSH != "" { + if _, err := rm.ensure(ctx); err == nil && rm.ping(ctx) == nil { + return rm.client, nil + } + } + rm.teardown() + } + if time.Now().Before(rm.retryAt) { + return nil, rm.lastErr + } + if err := rm.connect(ctx); err != nil { + rm.teardown() + rm.lastErr, rm.retryAt = err, time.Now().Add(retryAfter) + return nil, err + } + return rm.client, nil +} + +func (rm *Remote) alive() bool { + if rm.fwd == nil { + return true // a socket host has no process to lose + } + select { + case <-rm.fwdDone: + return false + default: + return true + } +} + +func unixClient(sock string) *http.Client { + return &http.Client{Transport: &http.Transport{ + DialContext: func(ctx context.Context, _, _ string) (net.Conn, error) { + var d net.Dialer + return d.DialContext(ctx, "unix", sock) + }, + }} +} + +func (rm *Remote) connect(ctx context.Context) error { + if rm.cfg.Socket != "" { + rm.client = unixClient(rm.cfg.Socket) + return rm.ping(ctx) + } + remote, err := rm.ensure(ctx) + if err != nil { + return err + } + local := filepath.Join(rm.dir, rm.cfg.Name+".sock") + os.Remove(local) + cmd := exec.Command(sshCommand, forwardArgs(rm.cfg, local, remote)...) + var stderr bytes.Buffer + cmd.Stderr = &stderr + if err := cmd.Start(); err != nil { + return fmt.Errorf("cannot start ssh: %w", err) + } + done := make(chan struct{}) + go func() { cmd.Wait(); close(done) }() + rm.fwd, rm.fwdDone = cmd, done + deadline := time.Now().Add(15 * time.Second) + for { + if c, err := net.Dial("unix", local); err == nil { + c.Close() + break + } + select { + case <-done: + return fmt.Errorf("ssh forward to %s exited: %s", rm.cfg.SSH, strings.TrimSpace(stderr.String())) + case <-ctx.Done(): + return ctx.Err() + case <-time.After(25 * time.Millisecond): + } + if time.Now().After(deadline) { + return fmt.Errorf("ssh forward to %s never came up", rm.cfg.SSH) + } + } + rm.client = unixClient(local) + return rm.ping(ctx) +} + +// ensure starts box serve on the host if it is not running and returns the +// path of its socket there. +func (rm *Remote) ensure(ctx context.Context) (string, error) { + cctx, cancel := context.WithTimeout(ctx, 30*time.Second) + defer cancel() + var stdout, stderr bytes.Buffer + cmd := exec.CommandContext(cctx, sshCommand, ensureArgs(rm.cfg)...) + cmd.Stdout, cmd.Stderr = &stdout, &stderr + if err := cmd.Run(); err != nil { + msg := strings.TrimSpace(stderr.String()) + if msg == "" { + msg = err.Error() + } + return "", fmt.Errorf("ssh %s: %s", rm.cfg.SSH, msg) + } + sock := strings.TrimSpace(stdout.String()) + if !socketPathRe.MatchString(sock) { + return "", fmt.Errorf("%s: box __ensure-serve answered %q, not a socket path (is box on the host up to date?)", rm.cfg.SSH, sock) + } + return sock, nil +} + +// ping checks the host answers as a box API of the version this router speaks. +func (rm *Remote) ping(ctx context.Context) error { + cctx, cancel := context.WithTimeout(ctx, hostTimeout) + defer cancel() + req, _ := http.NewRequestWithContext(cctx, "GET", "http://box/v1/version", nil) + resp, err := rm.client.Do(req) + if err != nil { + return err + } + defer resp.Body.Close() + var v struct{ API string } + if err := decodeJSON(resp.Body, &v); err != nil || v.API != "v1" { + return fmt.Errorf("host %s does not speak API v1 (got %q)", rm.cfg.Name, v.API) + } + rm.lastOK = time.Now() + return nil +} + +func (rm *Remote) teardown() { + if rm.fwd != nil { + rm.fwd.Process.Kill() + <-rm.fwdDone + rm.fwd, rm.fwdDone = nil, nil + } + if rm.client != nil { + rm.client.CloseIdleConnections() + rm.client = nil + } +} + +// Close stops the host's ssh forward. +func (rm *Remote) Close() error { + rm.mu.Lock() + defer rm.mu.Unlock() + rm.closed = true + rm.teardown() + return nil +} + +func decodeJSON(r io.Reader, v any) error { return json.NewDecoder(r).Decode(v) } + +// Open builds a router over the local API and the configured hosts. The +// returned close stops every ssh forward and removes their sockets. +func Open(version string, local http.Handler, cfg Config) (*Router, func(), error) { + base := os.Getenv("XDG_RUNTIME_DIR") + if base == "" { + base = os.TempDir() + } + // MkdirTemp makes it 0700: only this user can reach the forwarded + // sockets, as only this user can reach the local one. + dir, err := os.MkdirTemp(base, "box-router-") + if err != nil { + return nil, nil, err + } + hosts := []Host{{Name: Local, API: local}} + var remotes []*Remote + for _, hc := range cfg.Hosts { + rm := NewRemote(hc, dir) + remotes = append(remotes, rm) + hosts = append(hosts, Host{Name: hc.Name, API: rm, MaxUp: hc.MaxUp}) + } + closeAll := func() { + for _, rm := range remotes { + rm.Close() + } + os.RemoveAll(dir) + } + return New(version, hosts), closeAll, nil +} diff --git a/internal/router/hosts_test.go b/internal/router/hosts_test.go new file mode 100644 index 0000000..6304cff --- /dev/null +++ b/internal/router/hosts_test.go @@ -0,0 +1,253 @@ +package router + +import ( + "io" + "net" + "net/http" + "os" + "path/filepath" + "strings" + "testing" + "time" + + "box/internal/sandbox" + "box/internal/server" +) + +// The test binary stands in for ssh when BOX_FAKE_SSH_REMOTE is set: it +// answers __ensure-serve with that socket path, and for -L local:remote it +// forwards local to remote in-process, as ssh would over the network. +func TestMain(m *testing.M) { + if remote := os.Getenv("BOX_FAKE_SSH_REMOTE"); remote != "" { + fakeSSH(os.Args[1:], remote) + return + } + os.Exit(m.Run()) +} + +func fakeSSH(args []string, remote string) { + if log := os.Getenv("BOX_FAKE_SSH_LOG"); log != "" { + f, _ := os.OpenFile(log, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0o600) + f.WriteString(strings.Join(args, " ") + "\n") + f.Close() + } + if os.Getenv("BOX_FAKE_SSH_FAIL") != "" { + os.Stderr.WriteString("Permission denied (publickey).\n") + os.Exit(255) + } + if args[len(args)-1] == "__ensure-serve" { + os.Stdout.WriteString(remote + "\n") + os.Exit(0) + } + for i, a := range args { + if a != "-L" { + continue + } + local, to, _ := strings.Cut(args[i+1], ":") + ln, err := net.Listen("unix", local) + if err != nil { + os.Exit(255) + } + for { + c, err := ln.Accept() + if err != nil { + os.Exit(255) + } + go func() { + defer c.Close() + r, err := net.Dial("unix", to) + if err != nil { + return // as ssh does: the channel opens and closes + } + defer r.Close() + go io.Copy(r, c) + io.Copy(c, r) + }() + } + } + os.Exit(2) +} + +// serveAPI serves h on a fresh Unix socket until the test ends, or until the +// returned stop is called. +func serveAPI(t *testing.T, sock string, h http.Handler) (stop func()) { + t.Helper() + ln, err := net.Listen("unix", sock) + if err != nil { + t.Fatal(err) + } + srv := &http.Server{Handler: h} + go srv.Serve(ln) + stop = func() { srv.Close() } + t.Cleanup(stop) + return stop +} + +func fakeSSHEnv(t *testing.T, remote string) string { + t.Helper() + exe, err := os.Executable() + if err != nil { + t.Fatal(err) + } + old := sshCommand + sshCommand = exe + t.Cleanup(func() { sshCommand = old }) + log := filepath.Join(t.TempDir(), "ssh.log") + t.Setenv("BOX_FAKE_SSH_REMOTE", remote) + t.Setenv("BOX_FAKE_SSH_LOG", log) + return log +} + +func TestConfigValidation(t *testing.T) { + good := `hosts: + - name: gpu1 + ssh: me@gpu1.example.com + max_up: 4 + - name: lab + ssh: lab-alias + box: /opt/box/bin/box + - name: fwd + socket: /run/user/1000/fwd.sock +` + p := filepath.Join(t.TempDir(), "hosts.yaml") + os.WriteFile(p, []byte(good), 0o600) + c, err := LoadConfig(p) + if err != nil || len(c.Hosts) != 3 || c.Hosts[0].MaxUp != 4 { + t.Fatalf("good config: %v %+v", err, c) + } + for name, body := range map[string]string{ + "option as target": "hosts:\n - name: a\n ssh: -oProxyCommand=touch /tmp/x\n", + "space in target": "hosts:\n - name: a\n ssh: me@h ls\n", + "shell in box path": "hosts:\n - name: a\n ssh: h\n box: box;id\n", + "relative socket": "hosts:\n - name: a\n socket: x.sock\n", + "colon in socket": "hosts:\n - name: a\n socket: /a:b\n", + "both transports": "hosts:\n - name: a\n ssh: h\n socket: /s\n", + "neither transport": "hosts:\n - name: a\n", + "named local": "hosts:\n - name: local\n ssh: h\n", + "duplicate": "hosts:\n - name: a\n ssh: h\n - name: a\n ssh: g\n", + "bad name": "hosts:\n - name: ../a\n ssh: h\n", + "unknown key": "hosts:\n - name: a\n ssh: h\n port: 22\n", + "negative max_up": "hosts:\n - name: a\n ssh: h\n max_up: -1\n", + "yaml library panic": "hosts: !\n", + } { + os.WriteFile(p, []byte(body), 0o600) + if _, err := LoadConfig(p); err == nil { + t.Errorf("%s: should be refused", name) + } + } +} + +// ssh never sees a destination it could read as an option, and the remote +// command is nothing but the box binary and its fixed argument. +func TestSSHArgs(t *testing.T) { + h := HostConfig{Name: "a", SSH: "me@host"} + got := strings.Join(ensureArgs(h), " ") + if !strings.HasSuffix(got, "-- me@host box __ensure-serve") || !strings.Contains(got, "BatchMode=yes") { + t.Errorf("ensure: %s", got) + } + got = strings.Join(forwardArgs(h, "/l.sock", "/r.sock"), " ") + if !strings.HasSuffix(got, "-L /l.sock:/r.sock -- me@host") || !strings.Contains(got, "ExitOnForwardFailure=yes") { + t.Errorf("forward: %s", got) + } + if strings.Contains(got, "StrictHostKeyChecking") { + t.Error("host key checking must be left to the user's ssh config") + } +} + +// A remote host reached over (fake) ssh: the router lists its sandboxes, +// qualified, through a real box API. +func TestRemoteOverSSH(t *testing.T) { + home := t.TempDir() + t.Setenv("BOX_HOME", home) + if _, err := sandbox.Create("dev"); err != nil { + t.Fatal(err) + } + bf, _ := sandbox.BoxfilePath("dev") + if err := os.WriteFile(bf, []byte("base: docker.io/library/alpine\n"), 0o644); err != nil { + t.Fatal(err) + } + remote := filepath.Join(t.TempDir(), "remote.sock") + serveAPI(t, remote, server.NewHandler("remote-v")) + log := fakeSSHEnv(t, remote) + + rt, closeAll, err := Open("v", &fakeHost{}, Config{Hosts: []HostConfig{{Name: "gpu1", SSH: "me@gpu1"}}}) + if err != nil { + t.Fatal(err) + } + defer closeAll() + code, _, arr := do(t, rt, "GET", "/v1/sandboxes", "") + if code != 200 || len(arr) != 1 || arr[0]["name"] != "gpu1/dev" { + t.Fatalf("list = %d %v", code, arr) + } + _, _, hosts := do(t, rt, "GET", "/v1/hosts", "") + if len(hosts) != 2 || hosts[1]["version"] != "remote-v" || hosts[1]["error"] != nil { + t.Errorf("hosts = %v", hosts) + } + b, _ := os.ReadFile(log) + if n := strings.Count(string(b), "__ensure-serve"); n != 1 { + t.Errorf("ensure ran %d times, want once:\n%s", n, b) + } + if strings.Count(string(b), "-L ") != 1 { + t.Errorf("want one forward:\n%s", b) + } +} + +// A host whose ssh fails is reported, not fatal, and is not retried on +// every request. +func TestUnreachableHost(t *testing.T) { + log := fakeSSHEnv(t, "/nonexistent.sock") + t.Setenv("BOX_FAKE_SSH_FAIL", "1") + rt, closeAll, err := Open("v", &fakeHost{sandboxes: []map[string]any{sb("dev", false, 0)}}, + Config{Hosts: []HostConfig{{Name: "gpu1", SSH: "me@gpu1"}}}) + if err != nil { + t.Fatal(err) + } + defer closeAll() + _, _, arr := do(t, rt, "GET", "/v1/sandboxes", "") + if len(arr) != 1 || arr[0]["name"] != "dev" { + t.Errorf("the reachable host's sandboxes should still list: %v", arr) + } + _, _, hosts := do(t, rt, "GET", "/v1/hosts", "") + if e, _ := hosts[1]["error"].(string); !strings.Contains(e, "Permission denied") { + t.Errorf("the ssh error should be reported: %v", hosts[1]) + } + code, obj, _ := do(t, rt, "POST", "/v1/sandboxes/gpu1/dev/exec", "{}") + if code != 502 || obj["code"] != "unreachable" { + t.Errorf("exec on an unreachable host = %d %v", code, obj) + } + b, _ := os.ReadFile(log) + if n := strings.Count(string(b), "__ensure-serve"); n != 1 { + t.Errorf("a failing host should be retried after a pause, not per request: %d attempts", n) + } +} + +// When the remote server goes away and comes back (an idle exit, an +// upgrade), the next request after a quiet spell finds it again. +func TestRemoteServerComesBack(t *testing.T) { + remote := filepath.Join(t.TempDir(), "remote.sock") + stop := serveAPI(t, remote, &fakeHost{sandboxes: []map[string]any{sb("dev", false, 0)}}) + fakeSSHEnv(t, remote) + rm := NewRemote(HostConfig{Name: "gpu1", SSH: "me@gpu1"}, t.TempDir()) + defer rm.Close() + h := &Host{Name: "gpu1", API: rm} + rt := New("v", []Host{*h}) + + if code, _, _ := do(t, rt, "POST", "/v1/sandboxes/gpu1/dev/up", ""); code != 200 { + t.Fatalf("first request: %d", code) + } + stop() + os.Remove(remote) + if code, obj, _ := do(t, rt, "POST", "/v1/sandboxes/gpu1/dev/up", ""); code != 502 { + t.Errorf("with the server gone: %d %v", code, obj) + } + serveAPI(t, remote, &fakeHost{sandboxes: []map[string]any{sb("dev", false, 0)}}) + if code, obj, _ := do(t, rt, "POST", "/v1/sandboxes/gpu1/dev/up", ""); code != 200 { + t.Errorf("after it came back: %d %v", code, obj) + } + rm.mu.Lock() + rm.lastOK = time.Now().Add(-2 * pingAfter) + rm.mu.Unlock() + if code, _, _ := do(t, rt, "POST", "/v1/sandboxes/gpu1/dev/up", ""); code != 200 { + t.Errorf("after a quiet spell: %d", code) + } +} diff --git a/internal/router/router.go b/internal/router/router.go new file mode 100644 index 0000000..38459ef --- /dev/null +++ b/internal/router/router.go @@ -0,0 +1,453 @@ +// Package router puts several hosts' box APIs behind one: the local box and +// any number of remote machines running box. It speaks the same /v1 API as +// box serve, so the MCP server works through it unchanged. +// +// A sandbox on another host is named host/name; a bare name means the local +// sandbox of that name, or the only host that has it. When a bare name exists +// on several remote hosts, up and run are placed (where it is already up or +// warm, then the least loaded host); anything that acts on a particular +// sandbox's state must name the host. +// +// GET /v1/version +// GET /v1/sandboxes every host's, names qualified, plus "host" +// GET /v1/hosts reachability and load per host +// * /v1/sandboxes/[host/]{name}/{verb} forwarded to that host +package router + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "net/http" + "strings" + "sync" + "time" + + "box/internal/sandbox" +) + +// Local is the name of the host box itself runs on. +const Local = "local" + +const ( + listTTL = 2 * time.Second // absorbs a burst of lists and name lookups + hostTimeout = 3 * time.Second // per host, for lists and status +) + +// Host is one box API behind the router. +type Host struct { + Name string + API http.Handler + // MaxUp caps how many sandboxes placement keeps up on this host at + // once; 0 is no cap. Naming the host explicitly is not capped. + MaxUp int +} + +// Router is an http.Handler over several hosts. +type Router struct { + version string + hosts []Host + + mu sync.Mutex + snap *snapshot +} + +type snapshot struct { + at time.Time + hosts []hostList // in config order +} + +type hostList struct { + host *Host + items []map[string]any + err error +} + +// New returns a router over hosts, which are consulted in order. version is +// what /v1/version reports for the router itself. +func New(version string, hosts []Host) *Router { + return &Router{version: version, hosts: hosts} +} + +type apiError struct { + Error string `json:"error"` + Code string `json:"code"` +} + +func reply(w http.ResponseWriter, status int, v any) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(status) + json.NewEncoder(w).Encode(v) +} + +func fail(w http.ResponseWriter, status int, code, format string, args ...any) { + reply(w, status, apiError{Error: fmt.Sprintf(format, args...), Code: code}) +} + +// verbs are the per-sandbox routes of box serve, and the methods they take. +var verbs = map[string]string{ + "up": "POST", "down": "POST", "exec": "POST", "run": "POST", + "fork": "POST", "diff": "GET", "apply": "POST", "discard": "POST", +} + +func (rt *Router) ServeHTTP(w http.ResponseWriter, r *http.Request) { + path := r.URL.Path + switch { + case path == "/v1/version" && r.Method == "GET": + reply(w, http.StatusOK, map[string]any{"version": rt.version, "api": "v1", "router": true}) + case path == "/v1/sandboxes" && r.Method == "GET": + rt.handleList(w, r) + case path == "/v1/hosts" && r.Method == "GET": + rt.handleHosts(w, r) + case strings.HasPrefix(path, "/v1/sandboxes/"): + rt.handleSandbox(w, r, strings.Split(strings.TrimPrefix(path, "/v1/sandboxes/"), "/")) + default: + fail(w, http.StatusNotFound, "not_found", "no route %s %s", r.Method, path) + } +} + +// qualify names a sandbox as the list shows it: bare for local, host/name +// for anywhere else. +func qualify(host, name string) string { + if host == Local { + return name + } + return host + "/" + name +} + +func (rt *Router) handleList(w http.ResponseWriter, r *http.Request) { + out := []map[string]any{} + for _, hl := range rt.lists(r.Context()).hosts { + for _, it := range hl.items { + m := make(map[string]any, len(it)+1) + for k, v := range it { + m[k] = v + } + name, _ := it["name"].(string) + m["name"] = qualify(hl.host.Name, name) + m["host"] = hl.host.Name + out = append(out, m) + } + } + reply(w, http.StatusOK, out) +} + +type hostStatus struct { + Name string `json:"name"` + Version string `json:"version,omitempty"` + Sandboxes int `json:"sandboxes"` + Up int `json:"up"` + Warm int `json:"warm"` + MaxUp int `json:"max_up,omitempty"` + Error string `json:"error,omitempty"` +} + +func (rt *Router) handleHosts(w http.ResponseWriter, r *http.Request) { + snap := rt.lists(r.Context()) + out := make([]hostStatus, len(snap.hosts)) + var wg sync.WaitGroup + for i, hl := range snap.hosts { + st := hostStatus{Name: hl.host.Name, Sandboxes: len(hl.items), MaxUp: hl.host.MaxUp} + st.Up, st.Warm = load(hl.items) + if hl.err != nil { + st.Error = hl.err.Error() + out[i] = st + continue + } + wg.Add(1) + go func() { + defer wg.Done() + ctx, cancel := context.WithTimeout(r.Context(), hostTimeout) + defer cancel() + var v struct{ Version string } + if _, err := call(ctx, hl.host, "GET", "/v1/version", &v); err != nil { + st.Error = err.Error() + } + st.Version = v.Version + out[i] = st + }() + } + wg.Wait() + reply(w, http.StatusOK, out) +} + +// load counts a host's sandboxes that are up and its warm VMs. +func load(items []map[string]any) (up, warm int) { + for _, it := range items { + if u, _ := it["up"].(bool); u { + up++ + } + if n, _ := it["warm"].(float64); n > 0 { + warm += int(n) + } + } + return up, warm +} + +func (rt *Router) handleSandbox(w http.ResponseWriter, r *http.Request, parts []string) { + var hostName, name, verb string + switch len(parts) { + case 2: + name, verb = parts[0], parts[1] + case 3: + hostName, name, verb = parts[0], parts[1], parts[2] + default: + fail(w, http.StatusNotFound, "not_found", "no route %s %s", r.Method, r.URL.Path) + return + } + method, ok := verbs[verb] + if !ok { + fail(w, http.StatusNotFound, "not_found", "no route %s %s", r.Method, r.URL.Path) + return + } + if r.Method != method { + fail(w, http.StatusMethodNotAllowed, "bad_request", "%s takes %s", verb, method) + return + } + // Both become path segments on the host's API: neither may be anything + // but a single safe component. + names := []string{name} + if hostName != "" { + names = append(names, hostName) + } + for _, n := range names { + if err := sandbox.ValidName(n); err != nil { + fail(w, http.StatusBadRequest, "bad_name", "%v", err) + return + } + } + var h *Host + if hostName != "" { + if h = rt.host(hostName); h == nil { + fail(w, http.StatusNotFound, "no_host", "no host %q; configured: %s", hostName, rt.hostNames()) + return + } + } else { + var status int + var code, msg string + if h, status, code, msg = rt.resolve(r.Context(), name, verb); h == nil { + fail(w, status, code, "%s", msg) + return + } + } + rt.forward(w, r, h, name, verb) +} + +func (rt *Router) host(name string) *Host { + for i := range rt.hosts { + if rt.hosts[i].Name == name { + return &rt.hosts[i] + } + } + return nil +} + +func (rt *Router) hostNames() string { + var names []string + for _, h := range rt.hosts { + names = append(names, h.Name) + } + return strings.Join(names, ", ") +} + +// resolve picks the host for a bare name, or says why it cannot. +func (rt *Router) resolve(ctx context.Context, name, verb string) (*Host, int, string, string) { + type cand struct { + hl hostList + up bool + warm int + hostUp int + } + var cands []cand + var unreachable []string + for _, hl := range rt.lists(ctx).hosts { + if hl.err != nil { + unreachable = append(unreachable, hl.host.Name) + continue + } + for _, it := range hl.items { + if it["name"] != name { + continue + } + c := cand{hl: hl} + c.up, _ = it["up"].(bool) + if n, ok := it["warm"].(float64); ok { + c.warm = int(n) + } + c.hostUp, _ = load(hl.items) + if hl.host.Name == Local { + return hl.host, 0, "", "" // a bare name is the local one first + } + cands = append(cands, c) + } + } + switch len(cands) { + case 0: + msg := fmt.Sprintf("no sandbox %q on any host", name) + if len(unreachable) > 0 { + msg += fmt.Sprintf(" (unreachable: %s)", strings.Join(unreachable, ", ")) + } + return nil, http.StatusNotFound, "no_sandbox", msg + case 1: + return cands[0].hl.host, 0, "", "" + } + var where []string + for _, c := range cands { + where = append(where, qualify(c.hl.host.Name, name)) + } + ambiguous := fmt.Sprintf("%q is on several hosts (%s); name one", name, strings.Join(where, ", ")) + + // Where it is already up is where exec, down and up mean. + var up []cand + for _, c := range cands { + if c.up { + up = append(up, c) + } + } + switch { + case len(up) == 1 && (verb == "exec" || verb == "down" || verb == "up"): + return up[0].hl.host, 0, "", "" + case verb != "up" && verb != "run": + // Each copy's state is its own: acting on one must be a choice. + return nil, http.StatusConflict, "ambiguous", ambiguous + } + if verb == "run" { + best := -1 + for i, c := range cands { + if c.warm > 0 && (best < 0 || c.warm > cands[best].warm) { + best = i + } + } + if best >= 0 { + return cands[best].hl.host, 0, "", "" + } + } + best := -1 + for i, c := range cands { + if c.hl.host.MaxUp > 0 && c.hostUp >= c.hl.host.MaxUp { + continue + } + if best < 0 || c.hostUp < cands[best].hostUp { + best = i + } + } + if best < 0 { + return nil, http.StatusConflict, "capacity", fmt.Sprintf("every host with %q is at its max_up; name one (%s)", name, strings.Join(where, ", ")) + } + return cands[best].hl.host, 0, "", "" +} + +// forward hands the request to h with the path rewritten to the bare name. +// A fork's answer names the fork and its parent as the list does. +func (rt *Router) forward(w http.ResponseWriter, r *http.Request, h *Host, name, verb string) { + out := r.Clone(r.Context()) + out.URL.Path = "/v1/sandboxes/" + name + "/" + verb + out.URL.RawPath = "" + out.RequestURI = "" + if r.Method != "GET" { + defer rt.invalidate() // up, down, fork, run... change what a list shows + } + if verb != "fork" || h.Name == Local { + h.API.ServeHTTP(w, out) + return + } + rec := newRecorder() + h.API.ServeHTTP(rec, out) + var res map[string]any + if rec.code < 300 && json.Unmarshal(rec.body.Bytes(), &res) == nil { + for _, k := range []string{"fork", "parent"} { + if v, ok := res[k].(string); ok { + res[k] = qualify(h.Name, v) + } + } + reply(w, rec.code, res) + return + } + for k, v := range rec.header { + w.Header()[k] = v + } + w.WriteHeader(rec.code) + w.Write(rec.body.Bytes()) +} + +func (rt *Router) invalidate() { + rt.mu.Lock() + rt.snap = nil + rt.mu.Unlock() +} + +// lists is every host's sandbox list, fetched in parallel and kept briefly. +// A host that fails is reported in its entry, not as a failure of the whole. +func (rt *Router) lists(ctx context.Context) *snapshot { + rt.mu.Lock() + defer rt.mu.Unlock() + if rt.snap != nil && time.Since(rt.snap.at) < listTTL { + return rt.snap + } + snap := &snapshot{at: time.Now(), hosts: make([]hostList, len(rt.hosts))} + var wg sync.WaitGroup + for i := range rt.hosts { + wg.Add(1) + go func() { + defer wg.Done() + hctx, cancel := context.WithTimeout(ctx, hostTimeout) + defer cancel() + hl := hostList{host: &rt.hosts[i]} + _, hl.err = call(hctx, hl.host, "GET", "/v1/sandboxes", &hl.items) + snap.hosts[i] = hl + }() + } + wg.Wait() + rt.snap = snap + return snap +} + +// call makes one request to a host and decodes its JSON answer into v. +func call(ctx context.Context, h *Host, method, path string, v any) (int, error) { + req, err := http.NewRequestWithContext(ctx, method, "http://box"+path, http.NoBody) + if err != nil { + return 0, err + } + rec := newRecorder() + done := make(chan struct{}) + go func() { + defer close(done) + h.API.ServeHTTP(rec, req) + }() + select { + case <-done: + case <-ctx.Done(): + return 0, fmt.Errorf("host %s: %w", h.Name, ctx.Err()) + } + if rec.code >= 400 { + var e apiError + json.Unmarshal(rec.body.Bytes(), &e) + if e.Error == "" { + e.Error = http.StatusText(rec.code) + } + return rec.code, fmt.Errorf("%s", e.Error) + } + if err := json.Unmarshal(rec.body.Bytes(), v); err != nil { + return rec.code, fmt.Errorf("host %s: unreadable answer to %s", h.Name, path) + } + return rec.code, nil +} + +// recorder is an in-memory http.ResponseWriter. +type recorder struct { + mu sync.Mutex + header http.Header + body bytes.Buffer + code int +} + +func newRecorder() *recorder { return &recorder{header: http.Header{}, code: http.StatusOK} } + +func (r *recorder) Header() http.Header { return r.header } +func (r *recorder) Write(b []byte) (int, error) { + r.mu.Lock() + defer r.mu.Unlock() + return r.body.Write(b) +} +func (r *recorder) WriteHeader(c int) { r.code = c } diff --git a/internal/router/router_test.go b/internal/router/router_test.go new file mode 100644 index 0000000..03b91ef --- /dev/null +++ b/internal/router/router_test.go @@ -0,0 +1,194 @@ +package router + +import ( + "encoding/json" + "fmt" + "io" + "net/http" + "net/http/httptest" + "strings" + "sync" + "testing" +) + +// fakeHost answers like one box serve with the given sandboxes, and records +// which per-sandbox requests reached it. +type fakeHost struct { + sandboxes []map[string]any + down bool // answers every request 502, like an unreachable remote + + mu sync.Mutex + hits []string +} + +func (f *fakeHost) ServeHTTP(w http.ResponseWriter, r *http.Request) { + if f.down { + w.WriteHeader(http.StatusBadGateway) + w.Write([]byte(`{"error":"connection refused","code":"unreachable"}`)) + return + } + switch r.URL.Path { + case "/v1/sandboxes": + json.NewEncoder(w).Encode(f.sandboxes) + return + case "/v1/version": + w.Write([]byte(`{"version":"v-test","api":"v1"}`)) + return + } + f.mu.Lock() + f.hits = append(f.hits, r.Method+" "+r.URL.Path) + f.mu.Unlock() + if strings.HasSuffix(r.URL.Path, "/fork") { + var b struct{ As string } + json.NewDecoder(r.Body).Decode(&b) + parent := strings.Split(r.URL.Path, "/")[3] + fmt.Fprintf(w, `{"fork":%q,"parent":%q}`, b.As, parent) + return + } + w.Write([]byte(`{"ok":true}`)) +} + +func (f *fakeHost) got() []string { + f.mu.Lock() + defer f.mu.Unlock() + return append([]string(nil), f.hits...) +} + +func sb(name string, up bool, warm int) map[string]any { + return map[string]any{"name": name, "up": up, "warm": warm} +} + +func do(t *testing.T, h http.Handler, method, path, body string) (int, map[string]any, []map[string]any) { + t.Helper() + req := httptest.NewRequest(method, path, strings.NewReader(body)) + rec := httptest.NewRecorder() + h.ServeHTTP(rec, req) + b, _ := io.ReadAll(rec.Body) + var obj map[string]any + var arr []map[string]any + if json.Unmarshal(b, &obj) != nil { + json.Unmarshal(b, &arr) + } + return rec.Code, obj, arr +} + +func TestListQualifiesNamesAndSurvivesADownHost(t *testing.T) { + local := &fakeHost{sandboxes: []map[string]any{sb("dev", true, 0)}} + gpu := &fakeHost{sandboxes: []map[string]any{sb("train", false, 2), sb("dev", false, 0)}} + gone := &fakeHost{down: true} + rt := New("v", []Host{{Name: Local, API: local}, {Name: "gpu1", API: gpu}, {Name: "gone", API: gone}}) + + code, _, arr := do(t, rt, "GET", "/v1/sandboxes", "") + if code != 200 { + t.Fatalf("list: %d", code) + } + var names []string + for _, it := range arr { + names = append(names, it["name"].(string)+"@"+it["host"].(string)) + } + if got := strings.Join(names, " "); got != "dev@local gpu1/train@gpu1 gpu1/dev@gpu1" { + t.Errorf("names = %s", got) + } + + _, _, hosts := do(t, rt, "GET", "/v1/hosts", "") + if len(hosts) != 3 || hosts[2]["error"] == nil || hosts[0]["error"] != nil || hosts[1]["warm"].(float64) != 2 { + t.Errorf("hosts = %v", hosts) + } +} + +func TestResolveAndForward(t *testing.T) { + local := &fakeHost{sandboxes: []map[string]any{sb("dev", false, 0)}} + a := &fakeHost{sandboxes: []map[string]any{sb("dev", false, 0), sb("only-a", false, 0), sb("twin", true, 0)}} + b := &fakeHost{sandboxes: []map[string]any{sb("twin", false, 3)}} + rt := New("v", []Host{{Name: Local, API: local}, {Name: "a", API: a}, {Name: "b", API: b}}) + + for _, c := range []struct { + method, path string + want int + host *fakeHost + hit string + }{ + {"POST", "/v1/sandboxes/dev/exec", 200, local, "POST /v1/sandboxes/dev/exec"}, // local first + {"POST", "/v1/sandboxes/a/dev/exec", 200, a, "POST /v1/sandboxes/dev/exec"}, // named host + {"POST", "/v1/sandboxes/only-a/up", 200, a, "POST /v1/sandboxes/only-a/up"}, // the only host with it + {"POST", "/v1/sandboxes/twin/exec", 200, a, "POST /v1/sandboxes/twin/exec"}, // where it is up + {"POST", "/v1/sandboxes/twin/run", 200, b, "POST /v1/sandboxes/twin/run"}, // where it is warm + {"GET", "/v1/sandboxes/b/twin/diff", 200, b, "GET /v1/sandboxes/twin/diff"}, + } { + code, obj, _ := do(t, rt, c.method, c.path, `{}`) + if code != c.want { + t.Errorf("%s %s = %d %v", c.method, c.path, code, obj) + continue + } + if hits := c.host.got(); len(hits) == 0 || hits[len(hits)-1] != c.hit { + t.Errorf("%s %s: host got %v, want %s", c.method, c.path, hits, c.hit) + } + } + + for _, c := range []struct { + method, path, code string + status int + }{ + {"GET", "/v1/sandboxes/twin/diff", "ambiguous", 409}, // state of one copy: must be named + {"POST", "/v1/sandboxes/ghost/exec", "no_sandbox", 404}, + {"POST", "/v1/sandboxes/nohost/dev/exec", "no_host", 404}, + {"POST", "/v1/sandboxes/..%2Fx/exec", "bad_name", 400}, + {"POST", "/v1/sandboxes/-x/dev/exec", "bad_name", 400}, + {"POST", "/v1/sandboxes/dev/nope", "not_found", 404}, + {"GET", "/v1/sandboxes/dev/exec", "bad_request", 405}, + {"POST", "/v1/sandboxes/a/b/c/exec", "not_found", 404}, + } { + status, obj, _ := do(t, rt, c.method, c.path, `{}`) + if status != c.status || obj["code"] != c.code { + t.Errorf("%s %s = %d %v, want %d %s", c.method, c.path, status, obj, c.status, c.code) + } + } +} + +// up on a name several remotes share goes to the least loaded host under +// its cap, and is refused when every host is at its cap. +func TestPlacement(t *testing.T) { + busy := &fakeHost{sandboxes: []map[string]any{sb("w", false, 0), sb("x", true, 0), sb("y", true, 0)}} + idle := &fakeHost{sandboxes: []map[string]any{sb("w", false, 0), sb("z", true, 0)}} + rt := New("v", []Host{{Name: "busy", API: busy}, {Name: "idle", API: idle}}) + if code, obj, _ := do(t, rt, "POST", "/v1/sandboxes/w/up", ""); code != 200 || len(idle.got()) != 1 { + t.Fatalf("up should go to the less loaded host: %d %v, idle got %v", code, obj, idle.got()) + } + + rt = New("v", []Host{{Name: "busy", API: busy, MaxUp: 2}, {Name: "idle", API: idle, MaxUp: 1}}) + if code, obj, _ := do(t, rt, "POST", "/v1/sandboxes/w/up", ""); code != 409 || obj["code"] != "capacity" { + t.Errorf("every host at max_up: %d %v", code, obj) + } + // Naming the host is the human's choice and is not capped. + if code, _, _ := do(t, rt, "POST", "/v1/sandboxes/busy/w/up", ""); code != 200 { + t.Errorf("an explicit host should not be capped: %d", code) + } +} + +func TestForkIsQualifiedOnItsHost(t *testing.T) { + a := &fakeHost{sandboxes: []map[string]any{sb("dev", false, 0)}} + rt := New("v", []Host{{Name: Local, API: &fakeHost{}}, {Name: "a", API: a}}) + code, obj, _ := do(t, rt, "POST", "/v1/sandboxes/a/dev/fork", `{"as":"try1"}`) + if code != 200 || obj["fork"] != "a/try1" || obj["parent"] != "a/dev" { + t.Errorf("fork = %d %v", code, obj) + } + if hits := a.got(); len(hits) != 1 || hits[0] != "POST /v1/sandboxes/dev/fork" { + t.Errorf("fork reached %v", hits) + } +} + +// A request that changes state drops the cached lists, so the next lookup +// sees, for instance, a sandbox that just came up. +func TestMutationInvalidatesLists(t *testing.T) { + a := &fakeHost{sandboxes: []map[string]any{sb("dev", false, 0)}} + rt := New("v", []Host{{Name: "a", API: a}}) + do(t, rt, "GET", "/v1/sandboxes", "") + a.sandboxes = append(a.sandboxes, sb("new", false, 0)) + if _, _, arr := do(t, rt, "GET", "/v1/sandboxes", ""); len(arr) != 1 { + t.Fatalf("within the TTL the list should be cached, got %d", len(arr)) + } + do(t, rt, "POST", "/v1/sandboxes/dev/up", "") + if _, _, arr := do(t, rt, "GET", "/v1/sandboxes", ""); len(arr) != 2 { + t.Errorf("after a POST the list should be fresh, got %d", len(arr)) + } +}