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
+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
+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)
+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)))
+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)
}
}
+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]
+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."}