exec tools

internal/agents/bash.go

@@ -0,0 +1,111 @@
+package agents
+
+import (
+	"context"
+	"encoding/json"
+	"errors"
+	"fmt"
+	"os"
+	"os/exec"
+	"strings"
+	"syscall"
+	"time"
+
+	"g.fc.run/ai/byre/internal/mcp"
+)
+
+const bashKeep = 64 << 10
+
+type bashArgs struct {
+	Command   string `json:"command"`
+	TimeoutMs *int   `json:"timeout_ms"`
+}
+
+func (m *Manager) bashTool(l *live) mcp.Tool {
+	return mcp.Tool{
+		Name: "bash",
+		Description: "Run a shell command on the Mac byre runs on, in your task's worktree, with $SHELL as a login shell, and get its exit status and output (stdout and stderr together, the last 64 KiB). " +
+			"byre runs it itself, outside any sandbox your own shell tool has, with byre's environment and $PORT set to this task's first port. " +
+			"Stdin is empty. The command runs in its own process group, which is killed at the timeout.",
+		InputSchema: json.RawMessage(`{
+			"type": "object",
+			"properties": {
+				"command": {"type": "string"},
+				"timeout_ms": {"type": "integer", "description": "Kill the command after this long (default 120000, at most 600000)."}
+			},
+			"required": ["command"]
+		}`),
+		Handler: func(ctx context.Context, raw json.RawMessage) (string, error) {
+			return m.bash(ctx, l, raw)
+		},
+	}
+}
+
+func (m *Manager) bash(ctx context.Context, l *live, raw json.RawMessage) (string, error) {
+	args := bashArgs{}
+	err := json.Unmarshal(raw, &args)
+	if err != nil {
+		return "", fmt.Errorf("malformed bash call: %w", err)
+	}
+	if strings.TrimSpace(args.Command) == "" {
+		return "", errors.New("command is empty")
+	}
+	timeout := defaultBashTimeout
+	if args.TimeoutMs != nil {
+		timeout = time.Duration(*args.TimeoutMs) * time.Millisecond
+	}
+	if timeout <= 0 || timeout > maxBashTimeout {
+		return "", fmt.Errorf("timeout_ms must be between 1 and %d", maxBashTimeout.Milliseconds())
+	}
+	task, err := m.store.GetTask(ctx, l.taskID)
+	if err != nil {
+		return "", err
+	}
+	shell := os.Getenv("SHELL")
+	if shell == "" {
+		return "", errors.New("SHELL is not set in byre's environment")
+	}
+
+	cctx, cancel := context.WithTimeout(ctx, timeout)
+	defer cancel()
+	cmd := exec.CommandContext(cctx, shell, "-l", "-c", args.Command)
+	cmd.Dir = l.dir
+	cmd.Env = append(os.Environ(), portEnv(task)...)
+	cmd.SysProcAttr = &syscall.SysProcAttr{Setpgid: true}
+	cmd.Cancel = func() error {
+		return syscall.Kill(-cmd.Process.Pid, syscall.SIGKILL)
+	}
+	cmd.WaitDelay = 5 * time.Second
+	out := &tailWriter{}
+	cmd.Stdout = out
+	cmd.Stderr = out
+	err = cmd.Run()
+
+	status := "exit status 0"
+	exit := &exec.ExitError{}
+	switch {
+	case ctx.Err() != nil:
+		return "", ctx.Err()
+	case cctx.Err() != nil:
+		status = fmt.Sprintf("killed after %s", timeout)
+	case errors.As(err, &exit):
+		status = exit.ProcessState.String()
+	case errors.Is(err, exec.ErrWaitDelay):
+		status = "exit status 0; output cut off after 5s: a background process kept it open"
+	case err != nil:
+		return "", err
+	}
+	return status + "\n" + string(out.b[max(0, len(out.b)-bashKeep):]), nil
+}
+
+type tailWriter struct {
+	b []byte
+}
+
+func (w *tailWriter) Write(p []byte) (int, error) {
+	w.b = append(w.b, p...)
+	if len(w.b) > 2*bashKeep {
+		w.b = append([]byte{}, w.b[len(w.b)-bashKeep:]...)
+	}
+	return len(p), nil
+}