diff --git a/hands/Sources/byre-hands/Act.swift b/hands/Sources/byre-hands/Act.swift index 10ff96a..7468d11 100644 --- a/hands/Sources/byre-hands/Act.swift +++ b/hands/Sources/byre-hands/Act.swift @@ -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 diff --git a/hands/Sources/byre-hands/Browser.swift b/hands/Sources/byre-hands/Browser.swift index c18a279..9b4a657 100644 --- a/hands/Sources/byre-hands/Browser.swift +++ b/hands/Sources/byre-hands/Browser.swift @@ -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) diff --git a/hands/Sources/byre-hands/Chrome.swift b/hands/Sources/byre-hands/Chrome.swift index 80c58c4..bb1082f 100644 --- a/hands/Sources/byre-hands/Chrome.swift +++ b/hands/Sources/byre-hands/Chrome.swift @@ -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))) diff --git a/hands/Sources/byre-hands/Input.swift b/hands/Sources/byre-hands/Input.swift index 9c6ee0f..a86168a 100644 --- a/hands/Sources/byre-hands/Input.swift +++ b/hands/Sources/byre-hands/Input.swift @@ -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) } } diff --git a/hands/Sources/byre-hands/Server.swift b/hands/Sources/byre-hands/Server.swift index 0f088dd..b225400 100644 --- a/hands/Sources/byre-hands/Server.swift +++ b/hands/Sources/byre-hands/Server.swift @@ -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] diff --git a/hands/Sources/byre-hands/Tools.swift b/hands/Sources/byre-hands/Tools.swift index acd09b8..edf114e 100644 --- a/hands/Sources/byre-hands/Tools.swift +++ b/hands/Sources/byre-hands/Tools.swift @@ -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."} diff --git a/internal/agents/bash.go b/internal/agents/bash.go index fadceed..ccf372a 100644 --- a/internal/agents/bash.go +++ b/internal/agents/bash.go @@ -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): diff --git a/internal/agents/clients.go b/internal/agents/clients.go index beeea23..ba610b8 100644 --- a/internal/agents/clients.go +++ b/internal/agents/clients.go @@ -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 { diff --git a/internal/agents/files.go b/internal/agents/files.go index cc6ece0..879a34d 100644 --- a/internal/agents/files.go +++ b/internal/agents/files.go @@ -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 { diff --git a/internal/agents/sign.go b/internal/agents/sign.go index b46a8ab..639315f 100644 --- a/internal/agents/sign.go +++ b/internal/agents/sign.go @@ -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) { diff --git a/internal/codex/session.go b/internal/codex/session.go index e33a90d..6e1c6c5 100644 --- a/internal/codex/session.go +++ b/internal/codex/session.go @@ -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{ diff --git a/internal/mcp/server.go b/internal/mcp/server.go index 7a4ced0..eebe72f 100644 --- a/internal/mcp/server.go +++ b/internal/mcp/server.go @@ -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