byre/hands/Sources/byre-hands/Server.swift

import AppKit

struct ToolError: Error {
    let message: String

    init(_ message: String) {
        self.message = message
    }
}

enum Content {
    case text(String)
    case png(Data)

    var json: [String: Any] {
        switch self {
        case .text(let s):
            return ["type": "text", "text": s]
        case .png(let d):
            return ["type": "image", "data": d.base64EncodedString(), "mimeType": "image/png"]
        }
    }
}

struct Tool {
    let name: String
    let description: String
    let schema: String
    let run: ([String: Any]) async throws -> [Content]
}

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
    guard let i = args.firstIndex(of: "--media-dir"), i + 1 < args.count else {
        return nil
    }
    return args[i + 1]
}()

@main
struct Server {
    static func main() {
        _ = NSApplication.shared
        NSApp.setActivationPolicy(.prohibited)
        AXUIElementSetMessagingTimeout(AXUIElementCreateSystemWide(), 3)
        if CommandLine.arguments.dropFirst().first == "check" {
            runCheck()
            RunLoop.main.run()
        }
        Thread.detachNewThread {
            let inFlight = DispatchGroup()
            while let line = readLine() {
                inFlight.enter()
                Task.detached {
                    await handle(line)
                    inFlight.leave()
                }
            }
            inFlight.wait()
            exit(0)
        }
        RunLoop.main.run()
    }

    static func handle(_ line: String) async {
        guard let data = line.data(using: .utf8),
              let req = try? JSONSerialization.jsonObject(with: data) as? [String: Any],
              let id = req["id"]
        else {
            return
        }
        var resp: [String: Any] = ["jsonrpc": "2.0", "id": id]
        let params = req["params"] as? [String: Any] ?? [:]
        switch req["method"] as? String ?? "" {
        case "initialize":
            resp["result"] = [
                "protocolVersion": params["protocolVersion"] as? String ?? "2025-06-18",
                "capabilities": ["tools": [:]],
                "serverInfo": ["name": "hands", "version": "0.0.0"],
            ]
        case "ping":
            resp["result"] = [:] as [String: Any]
        case "tools/list":
            resp["result"] = ["tools": tools.map { t in
                ["name": t.name, "description": t.description, "inputSchema": withTimeout(t.schema)]
            }]
        case "tools/call":
            resp["result"] = await call(params)
        default:
            resp["error"] = ["code": -32601, "message": "method not found: \(req["method"] ?? "")"]
        }
        var out = try! JSONSerialization.data(withJSONObject: resp, options: .withoutEscapingSlashes)
        out.append(Data("\n".utf8))
        stdoutLock.withLock {
            FileHandle.standardOutput.write(out)
        }
    }

    static func call(_ params: [String: Any]) async -> [String: Any] {
        let name = params["name"] as? String ?? ""
        let args = params["arguments"] as? [String: Any] ?? [:]
        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 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]
        } catch {
            return ["content": [Content.text("\(error)").json], "isError": true]
        }
    }
}