Require a timeout on every agent tool
This commit is contained in:
+9
-19
@@ -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):
|
||||
|
||||
Reference in New Issue
Block a user