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/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 {