exec tools

This commit is contained in:
Ian Gulliver
2026-10-19 14:08:14 -07:00
parent aab15c7d79
commit 598ee788f8
4 changed files with 147 additions and 8 deletions
+111
View File
@@ -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
}
+1 -1
View File
@@ -738,7 +738,7 @@ func (m *Manager) startLive(ctx context.Context, sess api.AgentSession, task api
if err != nil {
return nil, err
}
tools := append([]mcp.Tool{m.askTool(l), m.cancelQuestionTool(l), m.setMergeTitleTool(l), m.stopMergeTool(l), m.mergeToMainTool(l), m.setTaskSummaryTool(l), m.searchClosedTool(l), m.exposePortTool(l), m.takeSnapshotTool(l), m.terminalReadTool(l), m.terminalTypeTool(l)}, append(m.subtaskTools(l), m.clientTools(l)...)...)
tools := append([]mcp.Tool{m.askTool(l), m.cancelQuestionTool(l), m.setMergeTitleTool(l), m.stopMergeTool(l), m.mergeToMainTool(l), m.setTaskSummaryTool(l), m.searchClosedTool(l), m.exposePortTool(l), m.takeSnapshotTool(l), m.terminalReadTool(l), m.terminalTypeTool(l), m.bashTool(l)}, append(m.subtaskTools(l), m.clientTools(l)...)...)
lead := true
if sess.Role == api.Implementer {
g += longTitleGuidance(task)
+33 -7
View File
@@ -206,6 +206,11 @@ type Session struct {
pendingMu sync.Mutex
pending map[string]chan controlResult
ctx context.Context
stop context.CancelFunc
inflightMu sync.Mutex
inflight map[string]context.CancelFunc
hooks map[string]HookFunc
servers map[string]*mcp.Server
@@ -236,12 +241,14 @@ func Start(ctx context.Context, opts Options) (*Session, error) {
done: make(chan struct{}),
transcript: opts.Transcript,
pending: map[string]chan controlResult{},
inflight: map[string]context.CancelFunc{},
hooks: map[string]HookFunc{},
servers: map[string]*mcp.Server{},
id: initialID(opts),
onBadLine: opts.OnBadLine,
}
s.lastLine.Store(time.Now().UnixNano())
s.ctx, s.stop = context.WithCancel(context.Background())
args, err := s.buildArgs(opts)
if err != nil {
@@ -556,6 +563,7 @@ func (s *Session) record(dir string, line []byte) {
func (s *Session) readLoop(stdout io.Reader) {
defer close(s.done)
defer close(s.events)
defer s.stop()
r := bufio.NewReader(stdout)
for {
@@ -606,10 +614,23 @@ func (s *Session) handleLine(line []byte) {
return
case "control_request":
go s.handleControlRequest(env.RequestID, env.Request)
ctx, cancel := context.WithCancel(s.ctx)
s.inflightMu.Lock()
s.inflight[env.RequestID] = cancel
s.inflightMu.Unlock()
go s.handleControlRequest(ctx, cancel, env.RequestID, env.Request)
return
case "control_cancel_request", "keep_alive":
case "control_cancel_request":
s.inflightMu.Lock()
cancel := s.inflight[env.RequestID]
s.inflightMu.Unlock()
if cancel != nil {
cancel()
}
return
case "keep_alive":
return
}
@@ -674,8 +695,15 @@ func (s *Session) handleControlResponse(raw json.RawMessage) {
ch <- controlResult{response: resp.Response}
}
func (s *Session) handleControlRequest(id string, raw json.RawMessage) {
result, err := s.dispatchControl(raw)
func (s *Session) handleControlRequest(ctx context.Context, cancel context.CancelFunc, id string, raw json.RawMessage) {
defer func() {
s.inflightMu.Lock()
delete(s.inflight, id)
s.inflightMu.Unlock()
cancel()
}()
result, err := s.dispatchControl(ctx, raw)
if err != nil {
slog.Error("control request failed", "request_id", id, "err", err)
s.write(map[string]any{
@@ -699,7 +727,7 @@ func (s *Session) handleControlRequest(id string, raw json.RawMessage) {
})
}
func (s *Session) dispatchControl(raw json.RawMessage) (any, error) {
func (s *Session) dispatchControl(ctx context.Context, raw json.RawMessage) (any, error) {
req := struct {
Subtype string `json:"subtype"`
CallbackID string `json:"callback_id"`
@@ -712,8 +740,6 @@ func (s *Session) dispatchControl(raw json.RawMessage) (any, error) {
return nil, err
}
ctx := context.Background()
switch req.Subtype {
case "hook_callback":
fn := s.hooks[req.CallbackID]