byre/internal/agents/bash.go

package agents

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"os"
	"os/exec"
	"strings"
	"time"

	"g.fc.run/ai/byre/internal/mcp"
	"g.fc.run/ai/byre/internal/procgroup"
)

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 := procgroup.CommandContext(cctx, shell, "-l", "-c", args.Command)
	cmd.Dir = l.dir
	cmd.Env = append(os.Environ(), portEnv(task)...)
	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
}