exec tools
This commit is contained in:
@@ -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
|
||||
}
|
||||
@@ -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)
|
||||
|
||||
@@ -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]
|
||||
|
||||
Reference in New Issue
Block a user