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)) + } +}