diff --git a/README.md b/README.md index cee86e4..0250041 100644 --- a/README.md +++ b/README.md @@ -202,7 +202,11 @@ curl -s -X POST -H "Authorization: Bearer $TOKEN" -H 'Content-Type: application/ runtime with a lower operator ceiling rejects the request rather than silently shortening it. - **Redial policy** (openab-pty §9.2): stop on `4001` expired · `4002` replaced · `4004` session ended · `4010` revoked · handshake `401`; redial with backoff (1 s → 30 s) on `1000` / `4006` / - errors until the grant deadline. Nothing on the attached socket is trusted as identity — the + `4011` (liveness) / errors until the grant deadline. A connection that stayed up for 60 s resets + the backoff. The Mac keeps a receive outstanding during tool calls, so it answers the runtime's + pings, and pings the runtime itself every 15 s. On these attaches `exec` refuses a + `timeout_secs` over 100 s, because the runtime fails any call still running after 120 s; use + `exec_start` + `exec_poll` for long work. Nothing on the attached socket is trusted as identity — the grant is the identity; `Tailscale-User-Login` there would be pod-supplied. - `--no-attach` disables the plane (`/attach` → 404). diff --git a/Sources/InstanceMCPCore/AttachManager.swift b/Sources/InstanceMCPCore/AttachManager.swift index 5390913..1100d3c 100644 --- a/Sources/InstanceMCPCore/AttachManager.swift +++ b/Sources/InstanceMCPCore/AttachManager.swift @@ -220,27 +220,45 @@ public actor AttachManager { // MARK: helpers + /// Every grant, whatever its profile, is told how long a call may run. The + /// openab-pty runtime fails a call with no reply after 120 s, and older + /// runtimes or clients also drop the Mac during long calls (#68). + static let longCallNotice = """ + Keep each tool call on this Mac under about 60 seconds: a longer call can time out or \ + disconnect you. For long work use `exec_start`, then check it with `exec_poll` (when those \ + tools are listed); never `sleep` inside `exec`. + """ + + /// Calls on one attach run concurrently (#68), so parallel input calls can + /// interleave on the screen. + static let uiOrderNotice = """ + Tool calls run concurrently: issue `mouse`, `key` and `osascript` steps one at a time and \ + wait for each result before sending the next. + """ + static func desktopInstructions(profile: ToolProfile, base: String?) -> String? { - if profile == .observe { - let head = base.map { $0 + "\n\n" } ?? "" + let head = base.map { $0 + "\n\n" } ?? "" + switch profile { + case .observe: return head + """ You reached this computer through OpenAB Connect under the `observe` profile: a human lets \ you look, not act. Only `sys_info` and `screenshot` are available — you cannot click, type \ or run anything here. If the task needs input on this computer, ask the human to grant \ the `desktop` profile instead. - """ + """ + "\n\n" + longCallNotice + case .desktop: + return head + """ + You reached this Mac through OpenAB Connect: a human lent it to your sandbox session for a \ + limited time and is likely watching the screen. This is the `desktop` profile — there is no \ + `exec` tool here (prefer the shell in your own session for local work); drive the Mac through \ + `screenshot`, `mouse`, `key` and `osascript`, and — when `browser_*` tools are listed — through \ + the browser directly: `browser_navigate` then `browser_snapshot` gives you the page as text, \ + no screenshot needed. If a tool starts failing with "not attached", the grant ended; ask the \ + human to lend the Mac again. + """ + "\n\n" + uiOrderNotice + "\n\n" + longCallNotice + case .owner: + return head + uiOrderNotice + "\n\n" + longCallNotice } - guard profile == .desktop else { return base } - let head = base.map { $0 + "\n\n" } ?? "" - return head + """ - You reached this Mac through OpenAB Connect: a human lent it to your sandbox session for a \ - limited time and is likely watching the screen. This is the `desktop` profile — there is no \ - `exec` tool here (prefer the shell in your own session for local work); drive the Mac through \ - `screenshot`, `mouse`, `key` and `osascript`, and — when `browser_*` tools are listed — through \ - the browser directly: `browser_navigate` then `browser_snapshot` gives you the page as text, \ - no screenshot needed. If a tool starts failing with "not attached", the grant ended; ask the \ - human to lend the Mac again. - """ } /// Build the exact runtime mint request. Keeping the body shape in a pure diff --git a/Sources/InstanceMCPCore/ReverseAttachClient.swift b/Sources/InstanceMCPCore/ReverseAttachClient.swift index 6e54d6e..891f477 100644 --- a/Sources/InstanceMCPCore/ReverseAttachClient.swift +++ b/Sources/InstanceMCPCore/ReverseAttachClient.swift @@ -6,10 +6,19 @@ import Foundation /// /// The runtime is the MCP client here — it sends `initialize` first, then relays /// the CLI's `tools/list` / `tools/call` with its own integer ids. We answer each -/// frame with the same id and otherwise treat the socket exactly like an HTTP -/// request/response pair: one frame in, one frame out, no server-initiated +/// frame with the same id: one frame in, one frame out, no server-initiated /// messages. /// +/// Liveness: the runtime pings every 20 s and drops a Mac that stays silent for +/// more than three intervals. `URLSessionWebSocketTask` only answers a ping while +/// a `receive()` is outstanding, so a `receive()` is **always** outstanding: each +/// frame is answered in its own task, replies go out through one writer task, and +/// a slow tool call no longer silences the socket. Calls also run concurrently +/// (the runtime allows many in flight). When a socket ends, for any reason, the +/// calls it still has running are cancelled, and `exec` kills its process group. +/// As a second safeguard we send our own ping every `keepaliveInterval` while +/// attached; any inbound frame counts as alive on the runtime side. +/// /// Retry ownership is ours (the pod cannot reach us). Close codes decide: /// /// | code | meaning (§9.2) | we | @@ -19,8 +28,12 @@ import Foundation /// | 4004 | session ended | stop | /// | 4006 | runtime replaced | redial with backoff while the grant is valid | /// | 4010 | grant revoked | stop | +/// | 4011 | liveness: we were silent past the runtime's ping budget | redial with backoff while the grant is valid | /// | 1000 / error | dropped | redial with backoff while the grant is valid | /// +/// A connection that stayed up for `stableAfter` resets the backoff, so a drop +/// after a long attach redials in about a second. +/// /// A `401` on the handshake means the verifier is gone (expired, revoked, or the /// pod was replaced): stop and report it, there is nothing to wait for. /// @@ -81,8 +94,11 @@ public actor ReverseAttachClient { public var secretFile: URL? = nil /// How long to wait after a credential refusal before trying again. public var credentialRetry: TimeInterval = 300 - /// A connection that stayed up this long resets the backoff (switchboard). + /// A connection that stayed up this long resets the backoff. public var stableAfter: TimeInterval = 60 + /// How often we ping the peer while attached, below the runtime's 20 s + /// ping interval. 0 turns our own pings off (tests). + public var keepaliveInterval: TimeInterval = 15 public init(runtime: URL, session: String, secret: String, profile: ToolProfile, deadline: Date) { self.runtime = runtime; self.session = session; self.secret = secret @@ -131,6 +147,7 @@ public actor ReverseAttachClient { case 4002: return .stop(.replaced) case 4004: return .stop(.sessionEnded) case 4010: return .stop(.revoked) + case 4011: return .redial // liveness: silent past the ping budget (openab-pty#86) default: return .redial // 1000, 1006, 4006, anything else } } @@ -159,6 +176,8 @@ public actor ReverseAttachClient { /// Frames answered on the current socket; surfaced for status. public private(set) var callsServed: Int = 0 public private(set) var attachedSince: Date? + /// Calls still running, so our own cancellation can reach them. + private var inflight: [UUID: Task] = [:] /// `server` should already be `scoped(to: config.profile)`; this type does /// not re-scope so a caller can pass a pre-built, instruction-tailored server. @@ -180,6 +199,9 @@ public actor ReverseAttachClient { runLoopTask = nil task?.cancel(with: .normalClosure, reason: nil) task = nil + // Directly, not via the receive loop: during a redial backoff there is + // no loop left to do it. + cancelInflight() set(.ended(.cancelled)) } @@ -197,7 +219,7 @@ public actor ReverseAttachClient { let started = Date() let outcome = await dialOnce() if Task.isCancelled { return } - if config.kind == .switchboard, Date().timeIntervalSince(started) >= config.stableAfter { + if Date().timeIntervalSince(started) >= config.stableAfter { backoff = config.initialBackoff } let base: TimeInterval @@ -246,6 +268,10 @@ public actor ReverseAttachClient { } /// One connection lifetime. Returns what to do next. + /// + /// The receive loop never waits on a tool: each frame is answered in its own + /// task and the reply is queued to a single writer, so a `receive()` is + /// always outstanding and URLSession keeps answering the runtime's pings. private func dialOnce() async -> Disposition { let req = attachRequest() let session = URLSession(configuration: .ephemeral) @@ -255,9 +281,36 @@ public actor ReverseAttachClient { task = ws ws.resume() + // One writer per socket: WebSocket sends must not interleave. + let (outbox, outboxIn) = AsyncStream.makeStream() + let sendFailed = LockedFlag() + let writer = Task { [weak self] in + for await reply in outbox { + do { try await ws.send(.string(reply)) } catch { + // The socket is broken; wake the receive loop so it redials. + _ = sendFailed.trySet() + ws.cancel() + return + } + await self?.noteServed() + } + } + var keepalive: Task? + defer { + // Every exit (stop code, retryable close, send failure, our own + // cancel) ends this socket's calls. The runtime has already failed + // them for the caller, so nothing they do can be delivered, and an + // `exec` left running would outlive the grant (it kills its process + // group on cancellation). Long work belongs in `exec_start`, whose + // jobs are not tied to a socket. + cancelInflight() + outboxIn.finish() + writer.cancel() + keepalive?.cancel() + } + // The handshake outcome is only observable by receiving. A rejected // upgrade surfaces as an error whose task has an HTTPURLResponse. - var served = 0 var firstFrame = true while true { let message: URLSessionWebSocketTask.Message @@ -270,8 +323,8 @@ public actor ReverseAttachClient { return Self.disposition(handshakeStatus: http.statusCode, kind: config.kind) } let code = ws.closeCode.rawValue - if code != 0 { - log("attach \(config.session): closed \(code) after \(served) calls") + if code != 0, !sendFailed.isSet { + log("attach \(config.session): closed \(code) after \(callsServed) calls") finishAttached() return Self.disposition(closeCode: code, kind: config.kind) } @@ -282,8 +335,10 @@ public actor ReverseAttachClient { if firstFrame { firstFrame = false attachedSince = Date() + callsServed = 0 set(.attached) log("attach \(config.session): attached as \(config.profile.rawValue) to \(config.runtime.host ?? "?")") + keepalive = startKeepalive(ws) } let text: String switch message { @@ -291,17 +346,47 @@ public actor ReverseAttachClient { case .data(let d): text = String(decoding: d, as: UTF8.self) @unknown default: continue } - if let reply = await answer(text) { - do { try await ws.send(.string(reply)) } catch { - finishAttached() - return .redial + // No per-socket cap: the runtime already bounds calls in flight (64), + // as the HTTP endpoint relies on its clients. Calls pipelined on one + // socket may finish out of order (each reply carries its id), so a + // caller that needs UI steps in order must wait for each reply. + let id = UUID() + inflight[id] = Task { [weak self] in + guard let self else { return } + if let reply = await self.answer(text) { outboxIn.yield(reply) } + await self.inflightDone(id) + } + } + } + + /// Our own ping, so the runtime hears from us even if nothing else is sent. + private func startKeepalive(_ ws: URLSessionWebSocketTask) -> Task? { + let interval = config.keepaliveInterval + guard interval > 0 else { return nil } + let session = config.session + let log = self.log + return Task { + while !Task.isCancelled { + try? await Task.sleep(nanoseconds: UInt64(interval * 1_000_000_000)) + if Task.isCancelled { return } + ws.sendPing { error in + if let error { log("attach \(session): keepalive ping failed: \(error.localizedDescription)") } } - served += 1 - callsServed = served } } } + private func noteServed() { callsServed += 1 } + + private func inflightDone(_ id: UUID) { inflight[id] = nil } + + /// On every socket exit and on our own `cancel()`: the calls of a dead + /// socket cannot be answered, so none of them carries into the next one. + private func cancelInflight() { + for t in inflight.values { t.cancel() } + inflight.removeAll() + } + private func finishAttached() { attachedSince = nil task = nil @@ -317,6 +402,7 @@ public actor ReverseAttachClient { return encode(JSONRPCResponse(id: .null, error: e)) case .success(let r): rpc = r } + if let refusal = refuseLongExec(rpc) { return encode(refusal) } guard let response = await server.handle(rpc) else { return nil } if rpc.method == "tools/call" { let name = rpc.params?["name"]?.stringValue ?? "?" @@ -325,6 +411,28 @@ public actor ReverseAttachClient { return encode(response) } + /// The longest `exec` / `osascript` timeout we accept on an openab-pty + /// attach. The runtime fails any call that has no reply after 120 s + /// (`REQUEST_TIMEOUT`), so a longer one could only end in an error for the + /// caller. + public static let maxAttachedExecTimeout: TimeInterval = 100 + + /// An `exec` or `osascript` over `maxAttachedExecTimeout` on an openab-pty + /// attach is refused as a tool error (the model sees it) pointing at + /// `exec_start`, rather than clamped, so the caller is not surprised by a + /// silent early kill. + func refuseLongExec(_ rpc: JSONRPCRequest) -> JSONRPCResponse? { + guard config.kind == .openabPty, rpc.method == "tools/call", + let tool = rpc.params?["name"]?.stringValue, tool == "exec" || tool == "osascript", + let asked = rpc.params?["arguments"]?["timeout_secs"]?.doubleValue, + asked > Self.maxAttachedExecTimeout else { return nil } + let limit = Int(Self.maxAttachedExecTimeout) + return JSONRPCResponse(id: rpc.id ?? .null, result: ToolResult.error( + "\(tool) timeout_secs \(Int(asked)) is over this attach's limit of \(limit) s: the runtime fails " + + "any call still running after 120 s. For long work use exec_start, then check it with " + + "exec_poll; never sleep inside exec.").json) + } + private func encode(_ r: JSONRPCResponse) -> String { String(decoding: (try? JSONCoding.encoder.encode(r)) ?? Data("{}".utf8), as: UTF8.self) } diff --git a/Sources/InstanceMCPCore/Tools/ExecTool.swift b/Sources/InstanceMCPCore/Tools/ExecTool.swift index ee2d75b..3590391 100644 --- a/Sources/InstanceMCPCore/Tools/ExecTool.swift +++ b/Sources/InstanceMCPCore/Tools/ExecTool.swift @@ -15,7 +15,7 @@ public struct ExecTool: Tool { "properties": [ "command": ["type": "string", "description": "Shell command line, run via `zsh -c`."], "cwd": ["type": "string", "description": "Working directory. Default: user's home."], - "timeout_secs": ["type": "number", "description": "Kill after this many seconds. Default 60, max 600."], + "timeout_secs": ["type": "number", "description": "Kill after this many seconds. Default 60, max 600 (100 on openab-pty attaches; use exec_start + exec_poll for longer work)."], "max_output_bytes": ["type": "integer", "description": "Per-stream cap. Default 65536, max 1048576."], "env": ["type": "object", "additionalProperties": ["type": "string"], "description": "Extra environment variables."], ], @@ -103,22 +103,43 @@ public struct ExecTool: Tool { } let fired = LockedFlag() + let killer = ProcessGroupKiller(pid) let timer = DispatchSource.makeTimerSource(queue: .global()) timer.schedule(deadline: .now() + timeout) - timer.setEventHandler { if fired.trySet() { killpg(pid, SIGKILL) } } + timer.setEventHandler { if fired.trySet() { killer.kill() } } timer.resume() - let status: Int32 = await withCheckedContinuation { cont in - DispatchQueue.global().async { - var st: Int32 = 0 - while waitpid(pid, &st, 0) < 0 && errno == EINTR {} - cont.resume(returning: st) + // Cancelling the calling task (the caller went away, or the grant was + // revoked) kills the process group while its leader is alive. + let status: Int32 = await withTaskCancellationHandler { + await withCheckedContinuation { cont in + DispatchQueue.global().async { + // Wait without reaping first, so a kill racing the exit can + // never hit a recycled pid. + var info = siginfo_t() + while waitid(P_PID, id_t(pid), &info, WEXITED | WNOWAIT) < 0 && errno == EINTR {} + killer.markExited() + var st: Int32 = 0 + while waitpid(pid, &st, 0) < 0 && errno == EINTR {} + cont.resume(returning: st) + } } + } onCancel: { + killer.kill() } timer.cancel() // Pipes reach EOF once every writer in the (now dead) session has exited. - await withCheckedContinuation { (cont: CheckedContinuation) in - group.notify(queue: .global()) { cont.resume() } + // A child backgrounded by the shell (`cmd &`) can hold them open after the + // leader exits; it can no longer be killed safely (the pid may be + // recycled), but a cancelled call stops waiting for it and returns. + let drained = OneShotResume() + await withTaskCancellationHandler { + await withCheckedContinuation { (cont: CheckedContinuation) in + drained.set(cont) + group.notify(queue: .global()) { drained.fire() } + } + } onCancel: { + drained.fire() } let exitCode: Int32 @@ -162,6 +183,31 @@ final class OutputCollector: @unchecked Sendable { } } +/// `killpg(pid, SIGKILL)` until the leader has exited, never after: once it is +/// reaped its pid may belong to someone else. +final class ProcessGroupKiller: @unchecked Sendable { + private let lock = NSLock(); private let pid: pid_t; private var exited = false + init(_ pid: pid_t) { self.pid = pid } + func kill() { lock.lock(); defer { lock.unlock() }; if !exited { killpg(pid, SIGKILL) } } + func markExited() { lock.lock(); exited = true; lock.unlock() } +} + +/// Resumes a continuation exactly once, whether `fire()` comes before or after `set`. +final class OneShotResume: @unchecked Sendable { + private let lock = NSLock(); private var cont: CheckedContinuation?; private var fired = false + func set(_ c: CheckedContinuation) { + lock.lock() + if fired { lock.unlock(); c.resume(); return } + cont = c; lock.unlock() + } + func fire() { + lock.lock() + if fired { lock.unlock(); return } + fired = true; let c = cont; cont = nil; lock.unlock() + c?.resume() + } +} + final class LockedFlag: @unchecked Sendable { private let lock = NSLock(); private var v = false var isSet: Bool { lock.lock(); defer { lock.unlock() }; return v } diff --git a/Sources/InstanceMCPCore/Tools/OsascriptTool.swift b/Sources/InstanceMCPCore/Tools/OsascriptTool.swift index 8a8080e..03f2af0 100644 --- a/Sources/InstanceMCPCore/Tools/OsascriptTool.swift +++ b/Sources/InstanceMCPCore/Tools/OsascriptTool.swift @@ -17,7 +17,7 @@ public struct OsascriptTool: Tool { "properties": [ "script": ["type": "string"], "language": ["type": "string", "enum": ["applescript", "javascript"], "default": "applescript"], - "timeout_secs": ["type": "number", "description": "Default 15, max 300."], + "timeout_secs": ["type": "number", "description": "Default 15, max 300 (100 on openab-pty attaches)."], ], "required": ["script"], ] diff --git a/Tests/InstanceMCPCoreTests/InstanceMCPCoreTests.swift b/Tests/InstanceMCPCoreTests/InstanceMCPCoreTests.swift index 627b009..6d4f2ec 100644 --- a/Tests/InstanceMCPCoreTests/InstanceMCPCoreTests.swift +++ b/Tests/InstanceMCPCoreTests/InstanceMCPCoreTests.swift @@ -292,6 +292,39 @@ extension HTTPRequest { final class ExecToolTests: XCTestCase { let tool = ExecTool() + /// Cancelling the calling task kills the process group (#68 review): a call + /// dropped with its socket must not keep running on the Mac. + func testCancellingTheCallKillsTheProcessGroup() async throws { + let marker = FileManager.default.temporaryDirectory.appendingPathComponent("oab-exec-cancel-\(UUID().uuidString)") + defer { try? FileManager.default.removeItem(at: marker) } + let started = Date() + let call = Task { + try await ExecTool.run(command: "sleep 3; touch '\(marker.path)'", cwd: "/tmp", + env: ProcessInfo.processInfo.environment, timeout: 30, cap: 1024) + } + try await Task.sleep(nanoseconds: 300_000_000) + call.cancel() + let r = try await call.value + XCTAssertLessThan(Date().timeIntervalSince(started), 2.5, "the call returns as soon as it is cancelled") + XCTAssertEqual(r.exitCode, 128 + SIGKILL) + XCTAssertFalse(r.timedOut) + try await Task.sleep(nanoseconds: 3_500_000_000) + XCTAssertFalse(FileManager.default.fileExists(atPath: marker.path), "the shell's children died with it") + } + + func testCancelledCallDoesNotWaitForABackgroundedChildHoldingThePipes() async throws { + let started = Date() + let call = Task { + try await ExecTool.run(command: "sleep 4 & exit 0", cwd: "/tmp", + env: ProcessInfo.processInfo.environment, timeout: 30, cap: 1024) + } + try await Task.sleep(nanoseconds: 500_000_000) + call.cancel() + let r = try await call.value + XCTAssertLessThan(Date().timeIntervalSince(started), 2.5, "the call returns once cancelled, not when the child exits") + XCTAssertEqual(r.exitCode, 0) + } + func testStdoutStderrExit() async throws { let r = try await tool.call(arguments: ["command": "echo out; echo err 1>&2; exit 3"]) XCTAssertTrue(r.isError) diff --git a/Tests/InstanceMCPCoreTests/ReverseAttachLivenessTests.swift b/Tests/InstanceMCPCoreTests/ReverseAttachLivenessTests.swift new file mode 100644 index 0000000..147f9fe --- /dev/null +++ b/Tests/InstanceMCPCoreTests/ReverseAttachLivenessTests.swift @@ -0,0 +1,347 @@ +import Foundation +import Network +import XCTest +@testable import InstanceMCPCore + +/// A tool slower than the fake runtime's ping budget. +struct SlowTool: Tool { + let name = "slow" + let description = "sleeps" + let inputSchema: JSONValue = ["type": "object"] + let seconds: Double + func call(arguments: JSONValue) async throws -> ToolResult { + try await Task.sleep(nanoseconds: UInt64(seconds * 1_000_000_000)) + return .text("slow done") + } +} + +/// Plays the openab-pty runtime's liveness rule (`runtime/src/tools.rs` +/// `serve_attached`), scaled down: ping every `pingInterval`, count an interval +/// with no inbound frame (pong, ping or text) as missed, and close with 1000 +/// after more than `maxMissed` — what evicted the Mac in #68. After +/// `initialize` it sends a slow `tools/call`, then a fast `echo`. +final class LivenessRuntime: @unchecked Sendable { + private(set) var port: UInt16 = 0 + private let listener: NWListener + private let queue = DispatchQueue(label: "liveness-runtime") + let pingInterval: TimeInterval? + let maxMissed: Int + private var conn: NWConnection? + private var timer: DispatchSourceTimer? + private var missed = 0 + private var _evicted = false + private var _pongs = 0 + private var _clientPings = 0 + private var _replyOrder: [Int] = [] + let finished = XCTestExpectation(description: "both calls answered") + + var evicted: Bool { queue.sync { _evicted } } + var pongs: Int { queue.sync { _pongs } } + var clientPings: Int { queue.sync { _clientPings } } + var replyOrder: [Int] { queue.sync { _replyOrder } } + + /// `pingInterval == nil`: never ping (to observe the client's own pings). + init(pingInterval: TimeInterval?, maxMissed: Int = 3) throws { + self.pingInterval = pingInterval + self.maxMissed = maxMissed + let params = NWParameters.tcp + let ws = NWProtocolWebSocket.Options() + ws.autoReplyPing = false // deliver the client's pings to us so we can count them + params.defaultProtocolStack.applicationProtocols.insert(ws, at: 0) + listener = try NWListener(using: params, on: .any) + let sem = DispatchSemaphore(value: 0) + listener.stateUpdateHandler = { [weak self] st in + if case .ready = st { self?.port = self?.listener.port?.rawValue ?? 0; sem.signal() } + if case .failed = st { sem.signal() } + } + listener.newConnectionHandler = { [weak self] c in self?.accept(c) } + listener.start(queue: queue) + sem.wait() + } + + func stop() { + queue.sync { timer?.cancel(); conn?.cancel() } + listener.cancel() + } + + private func accept(_ c: NWConnection) { + conn = c + c.stateUpdateHandler = { [weak self] st in + if case .ready = st { self?.ready(c) } + } + c.start(queue: queue) + } + + private func ready(_ c: NWConnection) { + receiveLoop(c) + if let interval = pingInterval { + let t = DispatchSource.makeTimerSource(queue: queue) + t.schedule(deadline: .now() + interval, repeating: interval) + t.setEventHandler { [weak self] in self?.tick(c) } + t.resume() + timer = t + } + sendText(c, ["jsonrpc": "2.0", "id": 1, "method": "initialize", + "params": ["protocolVersion": "2025-06-18", "capabilities": [:], + "clientInfo": ["name": "openab-pty", "version": "t"]]]) + } + + private func tick(_ c: NWConnection) { + guard !_evicted else { return } + missed += 1 + if missed > maxMissed { + _evicted = true + timer?.cancel() + let meta = NWProtocolWebSocket.Metadata(opcode: .close) + meta.closeCode = .protocolCode(.normalClosure) + let ctx = NWConnection.ContentContext(identifier: "close", metadata: [meta]) + c.send(content: nil, contentContext: ctx, isComplete: true, completion: .contentProcessed { _ in c.cancel() }) + return + } + let meta = NWProtocolWebSocket.Metadata(opcode: .ping) + meta.setPongHandler(queue) { [weak self] error in + guard let self, error == nil else { return } + self._pongs += 1 + self.missed = 0 + } + let ctx = NWConnection.ContentContext(identifier: "ping", metadata: [meta]) + c.send(content: Data(), contentContext: ctx, isComplete: true, completion: .idempotent) + } + + private func sendText(_ c: NWConnection, _ v: JSONValue) { + let meta = NWProtocolWebSocket.Metadata(opcode: .text) + let ctx = NWConnection.ContentContext(identifier: "t", metadata: [meta]) + c.send(content: try! JSONCoding.encoder.encode(v), contentContext: ctx, isComplete: true, completion: .idempotent) + } + + private func receiveLoop(_ c: NWConnection) { + c.receiveMessage { [weak self] data, context, _, error in + guard let self, error == nil else { return } + let meta = context?.protocolMetadata(definition: NWProtocolWebSocket.definition) as? NWProtocolWebSocket.Metadata + switch meta?.opcode { + case .ping?: + self._clientPings += 1 + self.missed = 0 + let pong = NWProtocolWebSocket.Metadata(opcode: .pong) + let ctx = NWConnection.ContentContext(identifier: "pong", metadata: [pong]) + c.send(content: data ?? Data(), contentContext: ctx, isComplete: true, completion: .idempotent) + case .pong?: + self.missed = 0 + case .close?: + return + default: + self.missed = 0 + if let data, let v = try? JSONCoding.decoder.decode(JSONValue.self, from: data) { + self.onReply(c, v) + } + } + self.receiveLoop(c) + } + } + + private func onReply(_ c: NWConnection, _ v: JSONValue) { + guard let id = v["id"]?.intValue else { return } + if id == 1 { + // The slow call first, then a fast one right behind it. + sendText(c, ["jsonrpc": "2.0", "id": 2, "method": "tools/call", "params": ["name": "slow", "arguments": [:]]]) + sendText(c, ["jsonrpc": "2.0", "id": 3, "method": "tools/call", "params": ["name": "echo", "arguments": ["msg": "hi"]]]) + return + } + _replyOrder.append(id) + if _replyOrder.count == 2 { finished.fulfill() } + } +} + +final class ReverseAttachLivenessTests: XCTestCase { + func liveClient(port: UInt16, slow: Double, keepalive: TimeInterval) -> ReverseAttachClient { + var cfg = ReverseAttachClient.Config(runtime: URL(string: "ws://127.0.0.1:\(port)")!, session: "laptop", + secret: "s", profile: .owner, deadline: Date().addingTimeInterval(30)) + cfg.keepaliveInterval = keepalive + let server = MCPServer(name: "t", version: "0", tools: [EchoTool(), SlowTool(seconds: slow)]) + return ReverseAttachClient(config: cfg, server: server) + } + + /// #68: a tool call longer than the runtime's ping budget must not silence + /// the socket. The runtime here evicts after ~0.8 s without an inbound frame + /// and the tool takes 2.5 s. Our own keepalive is off, so only answered + /// pings can keep the attach alive. + func testSlowToolCallSurvivesTheRuntimePingBudget() async throws { + let runtime = try LivenessRuntime(pingInterval: 0.2, maxMissed: 3) + defer { runtime.stop() } + let client = liveClient(port: runtime.port, slow: 2.5, keepalive: 0) + await client.start() + await fulfillment(of: [runtime.finished], timeout: 10) + XCTAssertFalse(runtime.evicted, "the runtime evicted the Mac during a slow call") + XCTAssertGreaterThanOrEqual(runtime.pongs, 5, "pings must be answered while the tool runs") + XCTAssertEqual(runtime.replyOrder, [3, 2], "the fast call is answered while the slow one runs") + let state = await client.state + XCTAssertEqual(state, .attached) + await client.cancel() + } + + /// The second safeguard: while attached we ping the runtime ourselves. + func testClientSendsItsOwnKeepalivePings() async throws { + let runtime = try LivenessRuntime(pingInterval: nil) + defer { runtime.stop() } + let client = liveClient(port: runtime.port, slow: 0.1, keepalive: 0.2) + await client.start() + await fulfillment(of: [runtime.finished], timeout: 10) + try await Task.sleep(nanoseconds: 1_200_000_000) + XCTAssertGreaterThanOrEqual(runtime.clientPings, 3, "keepalive pings reach the runtime") + await client.cancel() + } + + func watchedClient(port: UInt16, tool: WatchedTool) -> ReverseAttachClient { + var cfg = ReverseAttachClient.Config(runtime: URL(string: "ws://127.0.0.1:\(port)")!, session: "laptop", + secret: "s", profile: .owner, deadline: Date().addingTimeInterval(30)) + cfg.keepaliveInterval = 0 + cfg.initialBackoff = 0.2 + return ReverseAttachClient(config: cfg, server: MCPServer(name: "t", version: "0", tools: [tool])) + } + + /// A call still running when its socket closes is cancelled, and nothing of + /// it reaches the next connection: the runtime already failed it. + func testCallsRunningWhenTheSocketClosesAreCancelled() async throws { + let runtime = try DroppingRuntime(closeAfter: 0.3) + defer { runtime.stop() } + let tool = WatchedTool(seconds: 3) + let client = watchedClient(port: runtime.port, tool: tool) + await client.start() + await fulfillment(of: [runtime.secondAttach], timeout: 10) + try await Task.sleep(nanoseconds: 300_000_000) + XCTAssertEqual(tool.counts.started, 1) + XCTAssertEqual(tool.counts.cancelled, 1, "the dropped call was cancelled") + try await Task.sleep(nanoseconds: 3_000_000_000) + XCTAssertEqual(tool.counts.completed, 0, "the dropped call never finished") + XCTAssertEqual(runtime.repliesOnSecondConnection, [1], "only the new initialize is answered") + await client.cancel() + } + + /// Our own cancel (grant revoked) reaches calls in flight. + func testRevokingTheGrantCancelsCallsInFlight() async throws { + let runtime = try DroppingRuntime(closeAfter: nil) + defer { runtime.stop() } + let tool = WatchedTool(seconds: 3) + let client = watchedClient(port: runtime.port, tool: tool) + await client.start() + await fulfillment(of: [runtime.callSent], timeout: 10) + try await Task.sleep(nanoseconds: 300_000_000) + XCTAssertEqual(tool.counts.started, 1) + await client.cancel() + try await Task.sleep(nanoseconds: 300_000_000) + XCTAssertEqual(tool.counts.cancelled, 1) + XCTAssertEqual(tool.counts.completed, 0) + } +} + +/// A slow tool that records whether it finished or was cancelled. +struct WatchedTool: Tool { + final class Counts: @unchecked Sendable { + private let lock = NSLock() + private var s = 0, c = 0, d = 0 + var started: Int { lock.lock(); defer { lock.unlock() }; return s } + var cancelled: Int { lock.lock(); defer { lock.unlock() }; return c } + var completed: Int { lock.lock(); defer { lock.unlock() }; return d } + func start() { lock.lock(); s += 1; lock.unlock() } + func cancel() { lock.lock(); c += 1; lock.unlock() } + func complete() { lock.lock(); d += 1; lock.unlock() } + } + let name = "watched" + let description = "sleeps, cancellation-aware" + let inputSchema: JSONValue = ["type": "object"] + let seconds: Double + let counts = Counts() + func call(arguments: JSONValue) async throws -> ToolResult { + counts.start() + do { try await Task.sleep(nanoseconds: UInt64(seconds * 1_000_000_000)) } catch { + counts.cancel(); throw error + } + counts.complete() + return .text("watched done") + } +} + +/// On the first connection: `initialize`, then a slow `watched` call, then +/// (if `closeAfter` is set) a plain 1000 close while the call runs. On the +/// second connection: `initialize` only, recording every reply id. +final class DroppingRuntime: @unchecked Sendable { + private(set) var port: UInt16 = 0 + private let listener: NWListener + private let queue = DispatchQueue(label: "dropping-runtime") + let closeAfter: TimeInterval? + private var conns: [NWConnection] = [] + private var _second: [Int] = [] + let callSent = XCTestExpectation(description: "slow call sent") + let secondAttach = XCTestExpectation(description: "redialled and initialized") + + var repliesOnSecondConnection: [Int] { queue.sync { _second } } + + init(closeAfter: TimeInterval?) throws { + self.closeAfter = closeAfter + let params = NWParameters.tcp + params.defaultProtocolStack.applicationProtocols.insert(NWProtocolWebSocket.Options(), at: 0) + listener = try NWListener(using: params, on: .any) + let sem = DispatchSemaphore(value: 0) + listener.stateUpdateHandler = { [weak self] st in + if case .ready = st { self?.port = self?.listener.port?.rawValue ?? 0; sem.signal() } + if case .failed = st { sem.signal() } + } + listener.newConnectionHandler = { [weak self] c in self?.accept(c) } + listener.start(queue: queue) + sem.wait() + } + + func stop() { + queue.sync { conns.forEach { $0.cancel() } } + listener.cancel() + } + + private func accept(_ c: NWConnection) { + conns.append(c) + let index = conns.count + c.stateUpdateHandler = { [weak self] st in + guard let self, case .ready = st else { return } + self.receiveLoop(c, index) + self.sendText(c, ["jsonrpc": "2.0", "id": 1, "method": "initialize", + "params": ["protocolVersion": "2025-06-18", "capabilities": [:], + "clientInfo": ["name": "openab-pty", "version": "t"]]]) + } + c.start(queue: queue) + } + + private func sendText(_ c: NWConnection, _ v: JSONValue) { + let meta = NWProtocolWebSocket.Metadata(opcode: .text) + let ctx = NWConnection.ContentContext(identifier: "t", metadata: [meta]) + c.send(content: try! JSONCoding.encoder.encode(v), contentContext: ctx, isComplete: true, completion: .idempotent) + } + + private func close(_ c: NWConnection) { + let meta = NWProtocolWebSocket.Metadata(opcode: .close) + meta.closeCode = .protocolCode(.normalClosure) + let ctx = NWConnection.ContentContext(identifier: "close", metadata: [meta]) + c.send(content: nil, contentContext: ctx, isComplete: true, completion: .contentProcessed { _ in c.cancel() }) + } + + private func receiveLoop(_ c: NWConnection, _ index: Int) { + c.receiveMessage { [weak self] data, context, _, error in + guard let self, error == nil else { return } + let meta = context?.protocolMetadata(definition: NWProtocolWebSocket.definition) as? NWProtocolWebSocket.Metadata + if meta?.opcode == .close { return } + if meta?.opcode == .text, let data, let v = try? JSONCoding.decoder.decode(JSONValue.self, from: data), + let id = v["id"]?.intValue { + if index == 1, id == 1 { + self.sendText(c, ["jsonrpc": "2.0", "id": 2, "method": "tools/call", + "params": ["name": "watched", "arguments": [:]]]) + self.callSent.fulfill() + if let after = self.closeAfter { + self.queue.asyncAfter(deadline: .now() + after) { self.close(c) } + } + } else if index == 2 { + self._second.append(id) + if id == 1 { self.secondAttach.fulfill() } + } + } + self.receiveLoop(c, index) + } + } +} diff --git a/Tests/InstanceMCPCoreTests/ReverseAttachTests.swift b/Tests/InstanceMCPCoreTests/ReverseAttachTests.swift index a374247..4d454a8 100644 --- a/Tests/InstanceMCPCoreTests/ReverseAttachTests.swift +++ b/Tests/InstanceMCPCoreTests/ReverseAttachTests.swift @@ -69,11 +69,68 @@ final class ReverseAttachPolicyTests: XCTestCase { } func testRecoverableCloseCodesRedial() { - for code in [1000, 1001, 1006, 4006, 4009] { + for code in [1000, 1001, 1006, 4006, 4009, 4011] { XCTAssertEqual(ReverseAttachClient.disposition(closeCode: code), .redial, "code \(code)") } } + /// openab-pty#86: the runtime closes a Mac that went silent past its ping + /// budget with 4011. That is a drop, not an ending: the grant is still good. + func testLivenessCloseRedials() { + XCTAssertEqual(ReverseAttachClient.disposition(closeCode: 4011), .redial) + XCTAssertEqual(ReverseAttachClient.disposition(closeCode: 4011, kind: .switchboard), .redial) + } + + func testLongExecIsRefusedOnARuntimeAttachWithAPointerToExecStart() async { + let c = ReverseAttachClient.Config(runtime: URL(string: "ws://h:1")!, session: "s", secret: "x", + profile: .owner, deadline: .distantFuture) + let client = ReverseAttachClient(config: c, server: fullServer()) + func call(_ args: JSONValue) async -> JSONValue { + let text = await client.answer(String(decoding: rpc("tools/call", id: 9, params: ["name": "exec", "arguments": args]), as: UTF8.self)) + return decodeJSON(Data(text!.utf8)) + } + let refused = await call(["command": "sleep 1500", "timeout_secs": 600]) + XCTAssertEqual(refused["id"]?.intValue, 9) + XCTAssertEqual(refused["result"]?["isError"]?.boolValue, true, "a tool error the model can read: \(refused)") + let message = refused["result"]?["content"]?.arrayValue?.first?["text"]?.stringValue ?? "" + XCTAssertTrue(message.contains("exec_start"), message) + XCTAssertTrue(message.contains("100"), message) + + // osascript has the same limit (its own cap is 300 s). + let osa = await client.answer(String(decoding: rpc("tools/call", id: 11, params: ["name": "osascript", "arguments": ["script": "delay 200", "timeout_secs": 200]]), as: UTF8.self)) + let osaResult = decodeJSON(Data(osa!.utf8))["result"] + XCTAssertEqual(osaResult?["isError"]?.boolValue, true, "\(String(describing: osaResult))") + XCTAssertTrue((osaResult?["content"]?.arrayValue?.first?["text"]?.stringValue ?? "").hasPrefix("osascript timeout_secs 200")) + + for ok: JSONValue in [["command": "id"], ["command": "id", "timeout_secs": 100]] { + let r = await call(ok) + XCTAssertEqual(r["result"]?["content"]?.arrayValue?.first?["text"]?.stringValue, "ran exec", "\(ok)") + } + // Other tools keep their own limits. + let start = await client.answer(String(decoding: rpc("tools/call", id: 10, params: ["name": "exec_start", "arguments": ["command": "x", "timeout_secs": 3600]]), as: UTF8.self)) + XCTAssertEqual(decodeJSON(Data(start!.utf8))["result"]?["content"]?.arrayValue?.first?["text"]?.stringValue, "ran exec_start") + } + + func testLongExecIsNotRefusedOnASwitchboardAttach() async { + let c = ReverseAttachClient.Config.switchboard(url: URL(string: "ws://127.0.0.1:1/vm/attach")!, secret: "x", profile: .owner) + let client = ReverseAttachClient(config: c, server: fullServer()) + let text = await client.answer(String(decoding: rpc("tools/call", id: 1, params: ["name": "exec", "arguments": ["command": "id", "timeout_secs": 600]]), as: UTF8.self)) + XCTAssertEqual(decodeJSON(Data(text!.utf8))["result"]?["content"]?.arrayValue?.first?["text"]?.stringValue, "ran exec") + } + + func testEveryGrantProfileIsToldToKeepCallsShort() { + for p in ToolProfile.allCases { + let text = AttachManager.desktopInstructions(profile: p, base: "base") ?? "" + XCTAssertTrue(text.hasPrefix("base\n\n"), p.rawValue) + XCTAssertTrue(text.contains("60 seconds"), p.rawValue) + XCTAssertTrue(text.contains("`exec_start`") && text.contains("`exec_poll`"), p.rawValue) + XCTAssertTrue(text.contains("never `sleep` inside `exec`"), p.rawValue) + XCTAssertEqual(text.contains("one at a time"), p != .observe, p.rawValue) + } + XCTAssertEqual(AttachManager.desktopInstructions(profile: .owner, base: nil), + AttachManager.uiOrderNotice + "\n\n" + AttachManager.longCallNotice) + } + func testHandshakeRejectionStopsButThrottleWaits() { XCTAssertEqual(ReverseAttachClient.disposition(handshakeStatus: 401), .stop(.handshakeRejected(401))) XCTAssertEqual(ReverseAttachClient.disposition(handshakeStatus: 404), .stop(.handshakeRejected(404))) @@ -378,6 +435,26 @@ final class ReverseAttachEndToEndTests: XCTestCase { XCTAssertTrue(states.all.contains { if case .waitingToRedial = $0 { return true }; return false }) } + /// A connection that stayed up for `stableAfter` resets the backoff on an + /// openab-pty attach too, not just on a switchboard. + func testRuntimeAttachResetsBackoffAfterAStableConnection() async throws { + let runtime = try FakeRuntime(secret: "s3", closeWith: .privateCode(4006)) + defer { runtime.stop() } + let states = StateSink() + var cfg = ReverseAttachClient.Config(runtime: URL(string: "ws://127.0.0.1:\(runtime.port)")!, session: "laptop", + secret: "s3", profile: .owner, deadline: Date().addingTimeInterval(3)) + cfg.initialBackoff = 0.2 + cfg.stableAfter = 0 // every connection counts as stable + let client = ReverseAttachClient(config: cfg, server: fullServer(), onStateChange: { states.push($0) }) + await client.start() + try await Task.sleep(nanoseconds: 3_500_000_000) + let waits = states.all.compactMap { s -> Double? in + if case .waitingToRedial(let w) = s { return w }; return nil + } + XCTAssertGreaterThanOrEqual(waits.count, 3, "\(states.all)") + XCTAssertTrue(waits.allSatisfy { $0 <= 0.2 + 1e-9 }, "backoff must not double after a stable attach: \(waits)") + } + func testHandshakeRefusalStops() async throws { // Nothing listens here → connection refused → redial, then deadline. // (A real 401 is exercised by openab-pty's own suite; here we prove the