ai/byre
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 14 Wiki Activity

Require a timeout on every agent tool

Browse Source
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
Download Patch File Download Diff File Expand all files Collapse all files
internal/agents/bash.go
+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):
internal/agents/clients.go
+24 -50
View File
@@ -19,14 +19,8 @@ import (
)
const (
defaultBashTimeout = 2 * time.Minute
maxBashTimeout = 10 * time.Minute
replyGrace = 30 * time.Second
copyTimeout = 10 * time.Minute
fileTimeout = 2 * time.Minute
handsTimeout = 2 * time.Minute
defaultReadLimit = 2000
clientByre = "/Applications/Byre.app/Contents/Helpers/byre"
defaultReadLimit = 2000
clientByre = "/Applications/Byre.app/Contents/Helpers/byre"
)
func (m *Manager) clientTools(l *live) []mcp.Tool {
@@ -53,7 +47,7 @@ func (m *Manager) clientTools(l *live) []mcp.Tool {
"properties": {
"client": {"type": "string", "description": "The client's name from list_clients, or its id when several are connected from that Mac."},
"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": ["client", "command"]
}`),
@@ -319,18 +313,18 @@ func (m *Manager) clientBash(ctx context.Context, raw json.RawMessage) (string,
if err != nil {
return "", fmt.Errorf("malformed client_bash call: %w", err)
}
timeout, err := args.timeout()
err = args.check()
if err != nil {
return "", err
}
exitCode, output, err := m.runClientBash(ctx, args.Client, args.Command, timeout)
exitCode, output, err := m.runClientBash(ctx, args.Client, args.Command)
if err != nil {
return "", err
}
status := fmt.Sprintf("exit status %d", exitCode)
if exitCode == -1 {
status = fmt.Sprintf("killed after %s", timeout)
status = fmt.Sprintf("killed after %s", mcp.Timeout(ctx))
}
return status + "\n" + output, nil
}
@@ -378,9 +372,7 @@ func (m *Manager) upload(ctx context.Context, client, full, dest string) error {
id := m.clients.AddTransfer(clients.Transfer{Dir: filepath.Dir(full), Name: filepath.Base(full)})
defer m.clients.RemoveTransfer(id)
cctx, cancel := context.WithTimeout(ctx, copyTimeout)
defer cancel()
_, err = m.clients.Call(cctx, client, api.ClientRequest{Op: api.Upload, Transfer: &id, Dest: &dest})
_, err = m.clients.Call(ctx, client, api.ClientRequest{Op: api.Upload, Transfer: &id, Dest: &dest})
return err
}
@@ -410,7 +402,7 @@ func (m *Manager) clientOpen(ctx context.Context, l *live, raw json.RawMessage)
}
command := "open " + dir + "/" + shellQuote(name)
exitCode, output, err := m.runClientBash(ctx, args.Client, command, fileTimeout)
exitCode, output, err := m.runClientBash(ctx, args.Client, command)
if err != nil {
return "", err
}
@@ -445,9 +437,7 @@ func (m *Manager) showPage(ctx context.Context, l *live, raw json.RawMessage) (s
if u.Scheme != "http" && u.Scheme != "https" {
return "", fmt.Errorf("url %q isn't http or https", args.URL)
}
cctx, cancel := context.WithTimeout(ctx, fileTimeout)
defer cancel()
reply, err := m.clients.Call(cctx, args.Client, api.ClientRequest{Op: api.OpenPage, Task: &l.taskID, Url: &args.URL})
reply, err := m.clients.Call(ctx, args.Client, api.ClientRequest{Op: api.OpenPage, Task: &l.taskID, Url: &args.URL})
if err != nil {
return "", err
}
@@ -492,9 +482,7 @@ func (m *Manager) drivePage(ctx context.Context, l *live, raw json.RawMessage) m
default:
return mcp.TextResult(fmt.Sprintf("unknown op %q", args.Op), true)
}
cctx, cancel := context.WithTimeout(ctx, handsTimeout)
defer cancel()
reply, err := m.clients.Call(cctx, args.Client, req)
reply, err := m.clients.Call(ctx, args.Client, req)
if err != nil {
return mcp.TextResult(err.Error(), true)
}
@@ -540,9 +528,7 @@ func (m *Manager) show(ctx context.Context, l *live, raw json.RawMessage, tool s
} else if fi.IsDir() {
shown += "/"
}
cctx, cancel := context.WithTimeout(ctx, fileTimeout)
defer cancel()
_, err = m.clients.Call(cctx, args.Client, api.ClientRequest{Op: op, Task: &l.taskID, Path: &shown})
_, err = m.clients.Call(ctx, args.Client, api.ClientRequest{Op: op, Task: &l.taskID, Path: &shown})
if err != nil {
return "", err
}
@@ -581,9 +567,7 @@ func (m *Manager) clientDownload(ctx context.Context, l *live, raw json.RawMessa
id := m.clients.AddTransfer(clients.Transfer{Dir: tmp, Name: "archive.tgz", Receive: true})
defer m.clients.RemoveTransfer(id)
cctx, cancel := context.WithTimeout(ctx, copyTimeout)
defer cancel()
_, err = m.clients.Call(cctx, args.Client, api.ClientRequest{Op: api.Download, Path: &args.Path, Transfer: &id})
_, err = m.clients.Call(ctx, args.Client, api.ClientRequest{Op: api.Download, Path: &args.Path, Transfer: &id})
if err != nil {
return "", err
}
@@ -621,9 +605,7 @@ func (m *Manager) readClientFile(ctx context.Context, client, path string) ([]by
if err != nil {
return nil, err
}
cctx, cancel := context.WithTimeout(ctx, fileTimeout)
defer cancel()
reply, err := m.clients.Call(cctx, client, api.ClientRequest{Op: api.ReadFile, Path: &path})
reply, err := m.clients.Call(ctx, client, api.ClientRequest{Op: api.ReadFile, Path: &path})
if err != nil {
return nil, err
}
@@ -638,9 +620,7 @@ func (m *Manager) writeClientFile(ctx context.Context, client, path string, data
if err != nil {
return err
}
cctx, cancel := context.WithTimeout(ctx, fileTimeout)
defer cancel()
_, err = m.clients.Call(cctx, client, api.ClientRequest{Op: api.WriteFile, Path: &path, Data: &data})
_, err = m.clients.Call(ctx, client, api.ClientRequest{Op: api.WriteFile, Path: &path, Data: &data})
return err
}
@@ -775,7 +755,7 @@ func clientShellPath(p string) (string, error) {
func (m *Manager) runClientByre(ctx context.Context, client, sub, args string) (string, error) {
command := clientByre + " " + sub + args
exitCode, output, err := m.runClientBash(ctx, client, command, fileTimeout)
exitCode, output, err := m.runClientBash(ctx, client, command)
if err != nil {
return "", err
}
@@ -785,9 +765,13 @@ func (m *Manager) runClientByre(ctx context.Context, client, sub, args string) (
return output, nil
}
func (m *Manager) runClientBash(ctx context.Context, client, command string, timeout time.Duration) (int, string, error) {
ms := int(timeout.Milliseconds())
cctx, cancel := context.WithTimeout(ctx, timeout+replyGrace)
func untilReply(ctx context.Context) (context.Context, context.CancelFunc) {
return context.WithTimeout(context.WithoutCancel(ctx), mcp.Timeout(ctx)+mcp.Grace)
}
func (m *Manager) runClientBash(ctx context.Context, client, command string) (int, string, error) {
ms := int(mcp.Timeout(ctx).Milliseconds())
cctx, cancel := untilReply(ctx)
defer cancel()
reply, err := m.clients.Call(cctx, client, api.ClientRequest{Op: api.Bash, Command: &command, TimeoutMs: &ms})
if err != nil {
@@ -814,18 +798,8 @@ func (m *Manager) clientHands(ctx context.Context, raw json.RawMessage) mcp.Resu
if args.Arguments == nil {
args.Arguments = &map[string]any{}
}
ms := func(m map[string]any, key string) time.Duration {
v, _ := m[key].(float64)
return time.Duration(v) * time.Millisecond
}
timeout := handsTimeout + ms(*args.Arguments, "wait_ms") + 2*ms(*args.Arguments, "timeout_ms")
steps, _ := (*args.Arguments)["steps"].([]any)
for _, s := range steps {
step, _ := s.(map[string]any)
timeout += ms(step, "wait")
}
cctx, cancel := context.WithTimeout(ctx, timeout)
(*args.Arguments)["timeout_ms"] = mcp.Timeout(ctx).Milliseconds()
cctx, cancel := untilReply(ctx)
defer cancel()
reply, err := m.clients.Call(cctx, args.Client, api.ClientRequest{Op: api.Hands, Tool: &args.Tool, Arguments: args.Arguments})
if err != nil {
internal/agents/files.go
+1 -3
View File
@@ -262,9 +262,7 @@ func (m *Manager) grep(ctx context.Context, l *live, raw json.RawMessage) (strin
return "", err
}
}
cctx, cancel := context.WithTimeout(ctx, fileTimeout)
defer cancel()
return grep.Search(cctx, args.Pattern, l.dir, path, args.options())
return grep.Search(ctx, args.Pattern, l.dir, path, args.options())
}
type listArgs struct {
internal/agents/sign.go
+4 -1
View File
@@ -10,7 +10,10 @@ import (
"g.fc.run/ai/byre/internal/clients"
)
const signTimeout = 2 * time.Minute
const (
signTimeout = 2 * time.Minute
replyGrace = 30 * time.Second
)
func (m *Manager) signer(ctx context.Context, user api.User) func(payload []byte) ([]byte, error) {
return func(payload []byte) ([]byte, error) {
internal/codex/session.go
+1 -1
View File
@@ -270,7 +270,7 @@ func dynamicTools(servers []*mcp.Server) []DynamicToolSpec {
Type: DynamicToolSpecTypeFunction,
Name: t.Name,
Description: t.Description,
InputSchema: t.InputSchema,
InputSchema: t.Schema(),
})
}
specs = append(specs, DynamicToolSpec{
internal/mcp/server.go
+67 -1
View File
@@ -4,6 +4,7 @@ import (
"context"
"encoding/json"
"fmt"
"time"
)
const protocolVersion = "2025-06-18"
@@ -39,7 +40,72 @@ func TextResult(text string, isError bool) Result {
return Result{Content: []Content{{Type: "text", Text: text}}, IsError: isError}
}
const (
MaxTimeout = 10 * time.Minute
Grace = 5 * time.Second
)
type timeoutKey struct{}
func Timeout(ctx context.Context) time.Duration {
return ctx.Value(timeoutKey{}).(time.Duration)
}
func (t Tool) Schema() json.RawMessage {
s := map[string]any{}
err := json.Unmarshal(t.InputSchema, &s)
if err != nil {
panic(fmt.Sprintf("tool %s: input schema: %s", t.Name, err))
}
props, _ := s["properties"].(map[string]any)
if props == nil {
props = map[string]any{}
s["properties"] = props
}
if props["timeout_ms"] == nil {
props["timeout_ms"] = map[string]any{"type": "integer", "description": fmt.Sprintf("Give up on the call after this long (at most %d).", MaxTimeout.Milliseconds())}
}
required, _ := s["required"].([]any)
s["required"] = append(required, "timeout_ms")
out, err := json.Marshal(s)
if err != nil {
panic(fmt.Sprintf("tool %s: input schema: %s", t.Name, err))
}
return out
}
func (t Tool) Call(ctx context.Context, args json.RawMessage) Result {
p := struct {
TimeoutMs *int `json:"timeout_ms"`
}{}
err := json.Unmarshal(args, &p)
if err != nil {
return TextResult(fmt.Sprintf("malformed %s call: %s", t.Name, err), true)
}
if p.TimeoutMs == nil || *p.TimeoutMs <= 0 || *p.TimeoutMs > int(MaxTimeout.Milliseconds()) {
return TextResult(fmt.Sprintf("timeout_ms is required, between 1 and %d", MaxTimeout.Milliseconds()), true)
}
timeout := time.Duration(*p.TimeoutMs) * time.Millisecond
cctx, cancel := context.WithTimeout(context.WithValue(ctx, timeoutKey{}, timeout), timeout)
defer cancel()
done := make(chan Result, 1)
go func() {
done <- t.run(cctx, args)
}()
late := time.NewTimer(timeout + Grace)
defer late.Stop()
select {
case res := <-done:
return res
case <-ctx.Done():
return TextResult(ctx.Err().Error(), true)
case <-late.C:
return TextResult(fmt.Sprintf("%s gave no result within %s", t.Name, timeout), true)
}
}
func (t Tool) run(ctx context.Context, args json.RawMessage) Result {
if t.Rich != nil {
return t.Rich(ctx, args)
}
@@ -127,7 +193,7 @@ func (s *Server) dispatch(ctx context.Context, req request) (any, *rpcError) {
case "tools/list":
tools := []toolDesc{}
for _, t := range s.Tools {
tools = append(tools, toolDesc{Name: t.Name, Description: t.Description, InputSchema: t.InputSchema})
tools = append(tools, toolDesc{Name: t.Name, Description: t.Description, InputSchema: t.Schema()})
}
return map[string]any{"tools": tools}, nil