ai/byre
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 15 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
hands/Sources/byre-hands/Act.swift
+10 -9
View File
@@ -56,14 +56,14 @@ func runStep(_ step: [String: Any], _ id: Int, _ w: AXUIElement, _ table: RefTab
for (k, right, count) in [("click", false, 1), ("double_click", false, 2), ("right_click", true, 1)] {
if step[k] != nil {
let p = try point(step[k], info, table)
await click(p, right: right, count: count)
try await click(p, right: right, count: count)
return "\(k) at screen \(Int(p.x)),\(Int(p.y))"
}
}
if step["drag"] != nil {
let from = try point(step["drag"], info, table)
let to = try point(step["to"], info, table)
await drag(from, to)
try await drag(from, to)
return "dragged"
}
if step["scroll"] != nil {
@@ -77,7 +77,7 @@ func runStep(_ step: [String: Any], _ id: Int, _ w: AXUIElement, _ table: RefTab
return "scrolled"
}
if let s = step.string("type") {
await type(s)
try await type(s)
return "typed \(s.count) characters"
}
if let k = step.string("key") {
@@ -85,7 +85,7 @@ func runStep(_ step: [String: Any], _ id: Int, _ w: AXUIElement, _ table: RefTab
return "key \(k)"
}
if let ms = try step.ms("wait") {
await pause(ms)
try await pause(ms)
return "waited \(ms)ms"
}
throw ToolError("unknown step \(step.keys.sorted())")
@@ -93,11 +93,11 @@ func runStep(_ step: [String: Any], _ id: Int, _ w: AXUIElement, _ table: RefTab
func waitFor(_ args: [String: Any], acted: Bool, _ present: (String) async throws -> Bool) async throws -> String? {
let wait = try args.ms("wait_ms")
let timeout = Double(try args.ms("timeout_ms") ?? 10000) / 1000
let timeout = Double(try args.ms("wait_timeout_ms") ?? 10000) / 1000
if let wait {
await pause(wait)
try await pause(wait)
} else if acted {
await pause(300)
try await pause(300)
}
guard let needle = args.string("wait_for") else {
return nil
@@ -110,16 +110,17 @@ func waitFor(_ args: [String: Any], acted: Bool, _ present: (String) async throw
if Date().timeIntervalSince(start) > timeout {
return String(format: "gave up waiting for %@ after %.1fs", needle.debugDescription, timeout)
}
await pause(200)
try await pause(200)
}
}
func runSteps(_ steps: [[String: Any]], _ id: Int, _ w: AXUIElement, _ table: RefTable) async throws -> [String] {
var log: [String] = []
for (i, step) in steps.enumerated() {
try Task.checkCancellation()
do {
log.append("\(i + 1). " + (try await runStep(step, id, w, table)))
await pause(50)
try await pause(50)
} catch let e as ToolError {
log.append("\(i + 1). " + e.message + "; stopped here")
break
hands/Sources/byre-hands/Browser.swift
+4 -4
View File
@@ -65,12 +65,12 @@ func tabs() throws -> String {
func waitLoaded(_ t: TabTarget, timeout: Double) async throws -> String {
let start = Date()
await pause(300)
try await pause(300)
while Date().timeIntervalSince(start) < timeout {
if (try? js("document.readyState", t)) == "complete" {
return String(format: "loaded after %.1fs", Date().timeIntervalSince(start))
}
await pause(200)
try await pause(200)
}
return String(format: "still loading after %.1fs", timeout)
}
@@ -97,7 +97,7 @@ func browserTool(_ args: [String: Any]) async throws -> [Content] {
} else {
_ = try osa("tell application \"Orion\" to set URL of \(target.spec) to \(quote(url))")
}
log.append(try await waitLoaded(target, timeout: Double(try args.ms("timeout_ms") ?? 15000) / 1000))
log.append(try await waitLoaded(target, timeout: Double(try args.ms("wait_timeout_ms") ?? 15000) / 1000))
log.append(try js("document.title + ' | ' + location.href", target))
case "js":
log.append(try js(try args.require("code", args.string), target))
@@ -125,7 +125,7 @@ func browserTool(_ args: [String: Any]) async throws -> [Content] {
if let tab = target.tab {
_ = try osa("tell application \"Orion\" to set current tab of window \(win) to tab \(tab) of window \(win)")
log.append("selected tab \(tab) to capture it")
await pause(300)
try await pause(300)
}
let id = try orionWindow(win)
let shot = try await capture(window: id, display: nil, crop: args.dict("crop"), maxEdge: args.double("max_edge") ?? 1568)
hands/Sources/byre-hands/Chrome.swift
+7 -6
View File
@@ -153,6 +153,7 @@ final class Chrome: @unchecked Sendable {
}
func send(_ method: String, _ params: [String: Any] = [:], session: String? = nil, timeout: Int = 30) async throws -> [String: Any] {
try Task.checkCancellation()
let id = lock.withLock {
next += 1
return next
@@ -180,7 +181,7 @@ final class Chrome: @unchecked Sendable {
}
}
Task {
await pause(timeout * 1000)
try? await pause(timeout * 1000)
resolve(id, .failure(ToolError("\(method) got no answer from Chrome in \(timeout)s")))
}
}
@@ -413,7 +414,7 @@ func openChrome() async throws -> (Chrome, String) {
if let chrome = await chromeFromProfile() {
return (chrome, "started Chrome with byre's profile\n")
}
await pause(200)
try await pause(200)
}
throw ToolError("Chrome's debugger for byre's profile didn't answer within 15s; quit any already-running byre-profile Chrome once and retry")
}
@@ -528,7 +529,7 @@ func traceStop(_ c: Chrome, _ t: TabState, _ target: String) async throws -> Str
if Date().timeIntervalSince(start) > 60 {
throw ToolError("Chrome didn't finish the trace within 60s")
}
await pause(100)
try await pause(100)
stream = c.lock.withLock { t.traceStream }
}
var data = Data()
@@ -686,14 +687,14 @@ func chromeTool(_ args: [String: Any]) async throws -> [Content] {
if let err = r.string("errorText") {
throw ToolError("navigating to \(url) failed: \(err)")
}
let timeout = Double(try args.ms("timeout_ms") ?? 15000) / 1000
let timeout = Double(try args.ms("wait_timeout_ms") ?? 15000) / 1000
let start = Date()
await pause(300)
try await pause(300)
var loaded = false
while !loaded, Date().timeIntervalSince(start) < timeout {
loaded = (try? await evaluate(c, t, "document.readyState")) == "complete"
if !loaded {
await pause(200)
try await pause(200)
}
}
log.append(String(format: loaded ? "loaded after %.1fs" : "still loading after %.1fs", Date().timeIntervalSince(start)))
hands/Sources/byre-hands/Input.swift
+20 -13
View File
@@ -18,8 +18,8 @@ let modifierFlags: [String: CGEventFlags] = [
"fn": .maskSecondaryFn,
]
func pause(_ ms: Int) async {
try? await Task.sleep(nanoseconds: UInt64(ms) * 1_000_000)
func pause(_ ms: Int) async throws {
try await Task.sleep(nanoseconds: UInt64(ms) * 1_000_000)
}
let shifted: [Character: String] = [
@@ -56,7 +56,12 @@ func lockInput() async throws -> Int32 {
close(fd)
throw ToolError("another agent has been driving this Mac's apps for 30s, so nothing ran; try again")
}
await pause(100)
do {
try await pause(100)
} catch {
close(fd)
throw error
}
}
return fd
}
@@ -82,7 +87,7 @@ func focusWindow(_ id: Int, _ w: AXUIElement) async throws -> WindowInfo {
}
NSRunningApplication(processIdentifier: info.pid)?.activate()
AXUIElementPerformAction(w, kAXRaiseAction as CFString)
await pause(150)
try await pause(150)
let now = try window(id)
if try !isFront(now) {
throw ToolError("could not bring window \(id) to the front (\(NSWorkspace.shared.frontmostApplication?.localizedName ?? "nothing") is frontmost), so no input was sent")
@@ -90,32 +95,34 @@ func focusWindow(_ id: Int, _ w: AXUIElement) async throws -> WindowInfo {
return now
}
func click(_ p: CGPoint, right: Bool, count: Int) async {
func click(_ p: CGPoint, right: Bool, count: Int) async throws {
let (down, up): (CGEventType, CGEventType) = right ? (.rightMouseDown, .rightMouseUp) : (.leftMouseDown, .leftMouseUp)
let button: CGMouseButton = right ? .right : .left
post(CGEvent(mouseEventSource: nil, mouseType: .mouseMoved, mouseCursorPosition: p, mouseButton: button))
await pause(30)
try await pause(30)
for i in 1...max(1, count) {
for t in [down, up] {
let e = CGEvent(mouseEventSource: nil, mouseType: t, mouseCursorPosition: p, mouseButton: button)
e?.setIntegerValueField(.mouseEventClickState, value: Int64(i))
post(e)
}
await pause(30)
try await pause(30)
}
}
func drag(_ from: CGPoint, _ to: CGPoint) async {
func drag(_ from: CGPoint, _ to: CGPoint) async throws {
post(CGEvent(mouseEventSource: nil, mouseType: .mouseMoved, mouseCursorPosition: from, mouseButton: .left))
post(CGEvent(mouseEventSource: nil, mouseType: .leftMouseDown, mouseCursorPosition: from, mouseButton: .left))
await pause(50)
defer {
post(CGEvent(mouseEventSource: nil, mouseType: .leftMouseUp, mouseCursorPosition: to, mouseButton: .left))
}
try await pause(50)
for i in 1...20 {
let f = CGFloat(i) / 20
let p = CGPoint(x: from.x + (to.x - from.x) * f, y: from.y + (to.y - from.y) * f)
post(CGEvent(mouseEventSource: nil, mouseType: .leftMouseDragged, mouseCursorPosition: p, mouseButton: .left))
await pause(10)
try await pause(10)
}
post(CGEvent(mouseEventSource: nil, mouseType: .leftMouseUp, mouseCursorPosition: to, mouseButton: .left))
}
func scroll(_ p: CGPoint, dx: Int32, dy: Int32) {
@@ -123,7 +130,7 @@ func scroll(_ p: CGPoint, dx: Int32, dy: Int32) {
post(CGEvent(scrollWheelEvent2Source: nil, units: .line, wheelCount: 2, wheel1: dy, wheel2: dx, wheel3: 0))
}
func type(_ s: String) async {
func type(_ s: String) async throws {
for ch in s {
for down in [true, false] {
if let (code, flags) = keystroke(ch) {
@@ -137,7 +144,7 @@ func type(_ s: String) async {
e?.keyboardSetUnicodeString(stringLength: units.count, unicodeString: units)
post(e)
}
await pause(15)
try await pause(15)
}
}
hands/Sources/byre-hands/Server.swift
+47 -2
View File
@@ -30,6 +30,31 @@ struct Tool {
}
let stdoutLock = NSLock()
let maxTimeoutMs = 600_000
final class Once: @unchecked Sendable {
private let lock = NSLock()
private var done = false
func run(_ body: () -> Void) {
let first = lock.withLock {
defer { done = true }
return !done
}
if first {
body()
}
}
}
func withTimeout(_ schema: String) -> Any {
var s = try! JSONSerialization.jsonObject(with: Data(schema.utf8)) as! [String: Any]
var props = s["properties"] as? [String: Any] ?? [:]
props["timeout_ms"] = ["type": "integer", "description": "Give up on the call after this long (at most \(maxTimeoutMs))."]
s["properties"] = props
s["required"] = (s["required"] as? [String] ?? []) + ["timeout_ms"]
return s
}
let mediaDir: String? = {
let args = CommandLine.arguments
@@ -84,7 +109,7 @@ struct Server {
resp["result"] = [:] as [String: Any]
case "tools/list":
resp["result"] = ["tools": tools.map { t in
["name": t.name, "description": t.description, "inputSchema": try! JSONSerialization.jsonObject(with: Data(t.schema.utf8))]
["name": t.name, "description": t.description, "inputSchema": withTimeout(t.schema)]
}]
case "tools/call":
resp["result"] = await call(params)
@@ -104,8 +129,28 @@ struct Server {
guard let tool = tools.first(where: { $0.name == name }) else {
return ["content": [Content.text("unknown tool: \(name)").json], "isError": true]
}
guard let ms = args.int("timeout_ms"), ms > 0, ms <= maxTimeoutMs else {
return ["content": [Content.text("timeout_ms is required, between 1 and \(maxTimeoutMs)").json], "isError": true]
}
do {
let content = try await withDisplayAwake { try await tool.run(args) }
let content = try await withCheckedThrowingContinuation { (c: CheckedContinuation<[Content], Error>) in
let once = Once()
let work = Task {
do {
let content = try await withDisplayAwake { try await tool.run(args) }
once.run { c.resume(returning: content) }
} catch {
once.run { c.resume(throwing: error) }
}
}
Task {
try? await Task.sleep(for: .milliseconds(ms))
once.run {
work.cancel()
c.resume(throwing: ToolError("\(name) gave no result within \(ms)ms"))
}
}
}
return ["content": content.map { $0.json }]
} catch let e as ToolError {
return ["content": [Content.text(e.message).json], "isError": true]
hands/Sources/byre-hands/Tools.swift
+2 -2
View File
@@ -7,7 +7,7 @@ let treeProps = """
let thenProps = """
"wait_ms": {"type": "integer", "description": "Wait this long after the steps, before wait_for, tree and capture (default 300 when there were steps)."},
"timeout_ms": {"type": "integer", "description": "How long wait_for waits before giving up (default 10000). Giving up is reported, not an error."},
"wait_timeout_ms": {"type": "integer", "description": "How long wait_for waits before giving up (default 10000). Giving up is reported, not an error."},
"capture": {"type": "boolean", "description": "Finish with a capture of the window."},
"crop": {"type": "object", "description": "Capture only this rectangle, in window points.", "properties": {"x": {"type": "number"}, "y": {"type": "number"}, "width": {"type": "number"}, "height": {"type": "number"}}},
"max_edge": {"type": "number", "description": "Longest edge of the captured image in pixels, 16 to 8192 (default 1568, which is as much as you can see; captures are never scaled up)."}
@@ -121,7 +121,7 @@ let tools: [Tool] = [
"params": {"type": "object", "description": "For cdp: the method's parameters."},
"wait_for": {"type": "string", "description": "Then wait until this text (case-insensitive) is in the tab's title or page text, polling every 200ms."},
"wait_ms": {"type": "integer", "description": "Wait this long after the op, before wait_for and capture."},
"timeout_ms": {"type": "integer", "description": "How long navigate waits for the load (default 15000) and wait_for waits (default 10000). Giving up is reported, not an error."},
"wait_timeout_ms": {"type": "integer", "description": "How long navigate waits for the load (default 15000) and wait_for waits (default 10000). Giving up is reported, not an error."},
"capture": {"type": "boolean", "description": "Finish with a screenshot of the tab's page."},
"max_edge": {"type": "number", "description": "Longest edge of the screenshot in pixels, 16 to 8192 (default 1568, which is as much as you can see)."},
"save": {"type": "boolean", "description": "With capture, also save the screenshot as a PNG in the task's media folder; the reply gives its path, which a Markdown image in your final message can show."}
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