Require a timeout on every agent tool

This commit is contained in:
Ian Gulliver
2026-10-22 15:35:33 -07:00
parent c59e1cb4c3
commit b2cb15dd2a
12 changed files with 196 additions and 111 deletions
+9 -19
View File
@@ -17,22 +17,14 @@ import (
const bashKeep = 64 << 10
type bashArgs struct {
Command string `json:"command"`
TimeoutMs *int `json:"timeout_ms"`
Command string `json:"command"`
}
func (a bashArgs) timeout() (time.Duration, error) {
func (a bashArgs) check() error {
if strings.TrimSpace(a.Command) == "" {
return 0, errors.New("command is empty")
return errors.New("command is empty")
}
timeout := defaultBashTimeout
if a.TimeoutMs != nil {
timeout = time.Duration(*a.TimeoutMs) * time.Millisecond
}
if timeout <= 0 || timeout > maxBashTimeout {
return 0, fmt.Errorf("timeout_ms must be between 1 and %d", maxBashTimeout.Milliseconds())
}
return timeout, nil
return nil
}
func (m *Manager) bashTool(l *live) mcp.Tool {
@@ -45,7 +37,7 @@ func (m *Manager) bashTool(l *live) mcp.Tool {
"type": "object",
"properties": {
"command": {"type": "string"},
"timeout_ms": {"type": "integer", "description": "Kill the command after this long (default 120000, at most 600000)."}
"timeout_ms": {"type": "integer", "description": "Kill the command after this long (at most 600000)."}
},
"required": ["command"]
}`),
@@ -61,7 +53,7 @@ func (m *Manager) bash(ctx context.Context, l *live, raw json.RawMessage) (strin
if err != nil {
return "", fmt.Errorf("malformed bash call: %w", err)
}
timeout, err := args.timeout()
err = args.check()
if err != nil {
return "", err
}
@@ -70,9 +62,7 @@ func (m *Manager) bash(ctx context.Context, l *live, raw json.RawMessage) (strin
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 := procgroup.CommandContext(ctx, shell, "-l", "-c", args.Command)
cmd.Dir = l.dir
cmd.WaitDelay = 5 * time.Second
out := &tailWriter{}
@@ -83,10 +73,10 @@ func (m *Manager) bash(ctx context.Context, l *live, raw json.RawMessage) (strin
status := "exit status 0"
exit := &exec.ExitError{}
switch {
case errors.Is(ctx.Err(), context.DeadlineExceeded):
status = fmt.Sprintf("killed after %s", mcp.Timeout(ctx))
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):