diff --git a/CHANGELOG.md b/CHANGELOG.md index 6f8ca46c4..a472a0499 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,7 @@ - Fix host-routed screen observations with Accessibility elements by validating their semantic owner separately from the screen raster target. #715, #710. - Fix application name and bundle resolution being blocked by reaped processes lingering in LaunchServices; require repeated native absence while retaining refusal for uncertain or changing process identities. #709. +- Let hosts embedding `PeekabooMCPServer` supply their own MCP transport with the same completion and cleanup lifecycle as stdio. Thanks @semyoren! #716. ## 4.3.4 - 2026-09-11 diff --git a/Core/PeekabooCore/Sources/PeekabooAgentRuntime/MCP/Server/PeekabooMCPServer.swift b/Core/PeekabooCore/Sources/PeekabooAgentRuntime/MCP/Server/PeekabooMCPServer.swift index 63a94d907..d67533696 100644 --- a/Core/PeekabooCore/Sources/PeekabooAgentRuntime/MCP/Server/PeekabooMCPServer.swift +++ b/Core/PeekabooCore/Sources/PeekabooAgentRuntime/MCP/Server/PeekabooMCPServer.swift @@ -21,6 +21,8 @@ public enum TransportType: CustomStringConvertible, Sendable { /// Peekaboo MCP Server implementation public actor PeekabooMCPServer { + @TaskLocal private static var toolCallServingGeneration: UUID? + private enum StrictCallTool: MCP.Method { typealias Parameters = Value typealias Result = CallTool.Result @@ -56,6 +58,11 @@ public actor PeekabooMCPServer { private let toolRegistry: MCPToolRegistry private let logger: os.Logger private let toolContext: MCPToolContext + private var servingGeneration: UUID? + private var startupTask: Task? + private var shutdownTask: Task? + private var acceptsToolCalls = false + private var activeToolCalls: [UUID: Task] = [:] private let serverName = PeekabooMCPVersion.serverName private let serverVersion = PeekabooMCPVersion.current @@ -153,24 +160,10 @@ public actor PeekabooMCPServer { throw MCP.MCPError.methodNotFound("Server deallocated") } - let request = try ToolCallRequest(params: params) - - guard let tool = await self.toolRegistry.tool(named: request.name) else { - throw MCP.MCPError.invalidParams("Tool '\(request.name)' not found") + guard let generation = Self.toolCallServingGeneration else { + throw MCP.MCPError.internalError("MCP request has no serving lifetime") } - - let arguments = ToolArguments(value: .object(request.arguments)) - do { - try MCPToolArgumentValidator.validateClosedProperties(tool: tool, arguments: arguments) - } catch let error as MCPToolArgumentSchemaError { - throw MCP.MCPError.invalidParams( - "Invalid arguments for tool '\(request.name)': \(error.localizedDescription)") - } - - // Execute tool on main thread - let response = try await self.toolContext.execute(tool: tool, arguments: arguments) - - return Self.callToolResult(from: response, toolName: request.name) + return try await self.handleToolCall(params, generation: generation) } // Resources list handler (empty for now, but prevents inspector errors) @@ -220,6 +213,77 @@ public actor PeekabooMCPServer { } } + private func handleToolCall(_ parameters: Value, generation: UUID) async throws -> CallTool.Result { + guard self.acceptsToolCalls, self.servingGeneration == generation else { + throw MCP.MCPError.internalError("MCP server is shutting down") + } + let id = UUID() + let execution = Task { try await self.executeToolCall(parameters) } + self.activeToolCalls[id] = execution + defer { self.activeToolCalls[id] = nil } + return try await withTaskCancellationHandler { + try await execution.value + } onCancel: { + execution.cancel() + } + } + + private func executeToolCall(_ parameters: Value) async throws -> CallTool.Result { + try Task.checkCancellation() + let request = try ToolCallRequest(params: parameters) + guard let tool = await self.toolRegistry.tool(named: request.name) else { + throw MCP.MCPError.invalidParams("Tool '\(request.name)' not found") + } + try Task.checkCancellation() + let arguments = ToolArguments(value: .object(request.arguments)) + do { + try MCPToolArgumentValidator.validateClosedProperties(tool: tool, arguments: arguments) + } catch let error as MCPToolArgumentSchemaError { + throw MCP.MCPError.invalidParams( + "Invalid arguments for tool '\(request.name)': \(error.localizedDescription)") + } + let response = try await self.toolContext.execute(tool: tool, arguments: arguments) + return Self.callToolResult(from: response, toolName: request.name) + } + + private func startServing(transport: any Transport, generation: UUID) async throws { + let startup = Task { + try await Self.$toolCallServingGeneration.withValue(generation) { + try await self.server.start(transport: transport) + } + } + self.startupTask = startup + try await startup.value + } + + private func stopServing(generation: UUID? = nil) async { + if let generation, self.servingGeneration != generation { + return + } + if let shutdownTask = self.shutdownTask { + await shutdownTask.value + return + } + self.acceptsToolCalls = false + let calls = Array(self.activeToolCalls.values) + for call in calls { + call.cancel() + } + let startup = self.startupTask + startup?.cancel() + let shutdown = Task { + // A cancelled connect can still complete; finish startup before disconnecting its SDK session. + _ = try? await startup?.value + await self.server.stop() + // SDK shutdown does not own incoming tool tasks. Keep their context alive until they drain. + for call in calls { + _ = try? await call.value + } + } + self.shutdownTask = shutdown + await shutdown.value + } + static func callToolResult(from response: ToolResponse, toolName: String? = nil) -> CallTool.Result { let fields = MCPToolResponseMetadataProjector.externalFields(from: response.meta, toolName: toolName) let metadata = fields.isEmpty ? nil : Metadata(additionalFields: fields) @@ -269,13 +333,20 @@ public actor PeekabooMCPServer { } func startForTesting(transport: any Transport) async throws { - try await self.server.start(transport: transport) + let generation = UUID() + self.servingGeneration = generation + self.acceptsToolCalls = true + try await self.startServing(transport: transport, generation: generation) } @discardableResult func stopForTesting() async -> Bool { - await self.server.stop() - return await self.releaseToolContextForTeardown() + await self.stopServing() + let released = await self.releaseToolContextForTeardown() + self.servingGeneration = nil + self.startupTask = nil + self.shutdownTask = nil + return released } private func releaseToolContextForTeardown() async -> Bool { @@ -288,8 +359,8 @@ public actor PeekabooMCPServer { public func serve(transport: TransportType, port: Int = 8080) async throws { self.logger.info("Starting Peekaboo MCP server on \(transport) transport, version: \(self.serverVersion)") - do { - let serverTransport: any Transport = switch transport { + try await self.run { + switch transport { case .stdio: EOFDrainingTransport(wrapping: StdioTransport()) case .http: @@ -298,17 +369,53 @@ public actor PeekabooMCPServer { case .sse: throw MCPError.notImplemented("SSE server transport not yet implemented") } - try await self.server.start(transport: serverTransport) + } + } + + /// Serves over a transport supplied by the host process. + /// + /// For hosts that embed the server instead of spawning the CLI — an application that + /// links `PeekabooCore` and speaks MCP over a connection it owns. The lifecycle is the + /// one `serve(transport:port:)` gives the built-in stdio transport: the server runs + /// until the transport completes, and the tool context is released on the way out, + /// on success and on failure alike. + public func serve(transport: any Transport) async throws { + self.logger.info("Starting Peekaboo MCP server on a host transport, version: \(self.serverVersion)") + try await self.run { transport } + } - // Keep the server running - await self.server.waitUntilCompleted() + private func run(makingTransport: () throws -> any Transport) async throws { + guard self.servingGeneration == nil else { + throw MCPError.executionFailed("MCP server is already serving a transport") + } + let generation = UUID() + self.servingGeneration = generation + self.acceptsToolCalls = true + defer { + self.servingGeneration = nil + self.startupTask = nil + self.shutdownTask = nil + } + do { + let serverTransport = try makingTransport() + try await withTaskCancellationHandler { + try Task.checkCancellation() + try await self.startServing(transport: serverTransport, generation: generation) + try Task.checkCancellation() + await self.server.waitUntilCompleted() + try Task.checkCancellation() + } onCancel: { + Task { await self.stopServing(generation: generation) } + } } catch { + await self.stopServing(generation: generation) let cleanupConfirmed = await self.releaseToolContextForTeardown() if !cleanupConfirmed { self.logger.error("Browser session cleanup remains pending after MCP server failure") } throw error } + await self.stopServing(generation: generation) let cleanupConfirmed = await self.releaseToolContextForTeardown() guard cleanupConfirmed else { throw MCPError.executionFailed( diff --git a/Core/PeekabooCore/Tests/PeekabooTests/MCP/PeekabooMCPHostTransportTests.swift b/Core/PeekabooCore/Tests/PeekabooTests/MCP/PeekabooMCPHostTransportTests.swift new file mode 100644 index 000000000..2adb4db0b --- /dev/null +++ b/Core/PeekabooCore/Tests/PeekabooTests/MCP/PeekabooMCPHostTransportTests.swift @@ -0,0 +1,288 @@ +import Foundation +import Logging +import MCP +import TachikomaMCP +import Testing +@testable import PeekabooAgentRuntime +@testable import PeekabooCore + +@Suite(.serialized) +struct PeekabooMCPHostTransportTests { + @Test(.timeLimit(.minutes(1)), arguments: [false, true]) + @MainActor + func `shutdown cancels and drains an accepted tool before serve returns`(cancelsServerTask: Bool) async throws { + let provider = HeldPermissionsProvider() + defer { provider.release() } + let context = await MCPToolTestHelpers.makeContext(permissionsStatusProvider: provider) + let server = try await PeekabooMCPServer(toolContext: context) + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + // InMemoryTransport drops messages sent before the receiving peer connects. + try await serverTransport.connect() + let client = Client(name: "MCPShutdownTests", version: "1") + defer { Task { await client.disconnect(); await server.stopForTesting() } } + var completedWithDrainedCall = false + let serving = Task { + defer { completedWithDrainedCall = provider.finished } + try await server.serve(transport: serverTransport) + } + _ = try await client.connect(transport: clientTransport) + let call: RequestContext = try await client.callTool(name: "permissions", arguments: [:]) + var started = provider.started.makeAsyncIterator() + #expect(await started.next() != nil) + if cancelsServerTask { + serving.cancel() + } else { + await clientTransport.disconnect() + } + let release = Task { @MainActor in + try? await Task.sleep(for: .milliseconds(100)) + provider.release() + } + + if cancelsServerTask { + await #expect(throws: CancellationError.self) { try await serving.value } + } else { + try await serving.value + } + await release.value + await client.disconnect() + _ = try? await call.value + #expect(provider.sawCancellation) + #expect(provider.finished) + #expect(completedWithDrainedCall) + } + + @Test + @MainActor + func `cancelling serve disconnects a host transport only once`() async throws { + let context = await MCPToolTestHelpers.makeContext() + let server = try await PeekabooMCPServer(toolContext: context) + let (clientTransport, underlying) = await InMemoryTransport.createConnectedPair() + try await underlying.connect() + let transport = HeldDisconnectTransport(wrapping: underlying) + let client = Client(name: "MCPSingleShutdownTests", version: "1") + let serving = Task { try await server.serve(transport: transport) } + _ = try await client.connect(transport: clientTransport) + + serving.cancel() + var started = transport.disconnectStarted.makeAsyncIterator() + #expect(await started.next() != nil) + try await Task.sleep(for: .milliseconds(100)) + #expect(await transport.disconnectCount == 1) + await transport.release() + await #expect(throws: CancellationError.self) { try await serving.value } + #expect(await transport.disconnectCount == 1) + await client.disconnect() + } + + @Test + @MainActor + func `a concurrent serve refuses before connecting a second transport`() async throws { + let context = await MCPToolTestHelpers.makeContext() + let server = try await PeekabooMCPServer(toolContext: context) + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + // InMemoryTransport drops messages sent before the receiving peer connects. + try await serverTransport.connect() + let client = Client(name: "MCPConcurrentServeTests", version: "1") + let serving = Task { try await server.serve(transport: serverTransport) } + _ = try await client.connect(transport: clientTransport) + let second = LifecycleHostTransport(failsToConnect: false) + await #expect(throws: (any Error).self) { try await server.serve(transport: second) } + #expect(await second.connectCount == 0) + await client.disconnect() + try await serving.value + } + + @Test + @MainActor + func `cancellation waits for startup before disconnecting its session`() async throws { + let context = await MCPToolTestHelpers.makeContext() + let server = try await PeekabooMCPServer(toolContext: context) + let transport = LifecycleHostTransport(failsToConnect: false, holdsConnect: true) + let serving = Task { try await server.serve(transport: transport) } + var started = transport.connectStarted.makeAsyncIterator() + #expect(await started.next() != nil) + + serving.cancel() + try await Task.sleep(for: .milliseconds(100)) + #expect(await transport.disconnectCount == 0) + await transport.releaseConnect() + await #expect(throws: CancellationError.self) { try await serving.value } + #expect(await transport.connectSawCancellation) + #expect(await transport.disconnectCount == 1) + } + + @Test + @MainActor + func `host transport startup failure releases its snapshot owner`() async throws { + let context = await MCPToolTestHelpers.makeContext() + let server = try await PeekabooMCPServer(toolContext: context) + let snapshots = await MCPToolUISnapshotStore(owner: server.snapshotOwnerForTesting()) + let snapshot = await snapshots.createSnapshot() + let transport = LifecycleHostTransport() + + await #expect(throws: HostTransportError.connectionFailed) { + try await server.serve(transport: transport) + } + + #expect(await transport.connectCount == 1) + #expect(await transport.disconnectCount == 1) + #expect(await snapshots.getSnapshot(id: snapshot.id) == nil) + #expect(await !snapshots.hasOwnerState()) + } + + @Test + @MainActor + func `host transport completion disconnects the SDK session`() async throws { + let context = await MCPToolTestHelpers.makeContext() + let server = try await PeekabooMCPServer(toolContext: context) + let transport = LifecycleHostTransport(failsToConnect: false) + + try await server.serve(transport: transport) + + #expect(await transport.connectCount == 1) + #expect(await transport.disconnectCount == 1) + } +} + +@MainActor +private final class HeldPermissionsProvider: PermissionsStatusProviding { + let started: AsyncStream + private let continuation: AsyncStream.Continuation + private var waiter: CheckedContinuation? + private(set) var finished = false + private(set) var sawCancellation = false + + init() { + let pair = AsyncStream.makeStream() + self.started = pair.stream + self.continuation = pair.continuation + } + + func permissionsStatus() async throws -> PermissionsStatus { + self.continuation.yield(()) + await withCheckedContinuation { self.waiter = $0 } + self.sawCancellation = Task.isCancelled + self.finished = true + return PermissionsStatus(screenRecording: true, accessibility: true, postEvent: true) + } + + func release() { + self.waiter?.resume() + self.waiter = nil + self.continuation.finish() + } +} + +private enum HostTransportError: Error, Equatable { + case connectionFailed +} + +private actor LifecycleHostTransport: Transport { + nonisolated let logger = Logger(label: "boo.peekaboo.tests.host-transport") + nonisolated let connectStarted: AsyncStream + private let connectContinuation: AsyncStream.Continuation + private let failsToConnect: Bool + private let holdsConnect: Bool + private var connectWaiter: CheckedContinuation? + private(set) var connectSawCancellation = false + private(set) var connectCount = 0 + private(set) var disconnectCount = 0 + + init(failsToConnect: Bool = true, holdsConnect: Bool = false) { + self.failsToConnect = failsToConnect + self.holdsConnect = holdsConnect + let pair = AsyncStream.makeStream() + self.connectStarted = pair.stream + self.connectContinuation = pair.continuation + } + + func connect() async throws { + self.connectCount += 1 + self.connectContinuation.yield(()) + if self.holdsConnect { + await withCheckedContinuation { self.connectWaiter = $0 } + } + self.connectSawCancellation = Task.isCancelled + if self.failsToConnect { + throw HostTransportError.connectionFailed + } + } + + func releaseConnect() { + self.connectWaiter?.resume() + self.connectWaiter = nil + self.connectContinuation.finish() + } + + func disconnect() async { + self.disconnectCount += 1 + } + + func send(_: Data) async throws { + throw HostTransportError.connectionFailed + } + + func receive() -> AsyncThrowingStream { + AsyncThrowingStream { $0.finish() } + } +} + +private actor HeldDisconnectTransport: Transport { + nonisolated let logger = Logger(label: "boo.peekaboo.tests.held-disconnect") + nonisolated let disconnectStarted: AsyncStream + private let continuation: AsyncStream.Continuation + private let underlying: any Transport + private var waiters: [CheckedContinuation] = [] + private var released = false + private(set) var disconnectCount = 0 + + init(wrapping underlying: any Transport) { + self.underlying = underlying + let pair = AsyncStream.makeStream() + self.disconnectStarted = pair.stream + self.continuation = pair.continuation + } + + func connect() async throws { + try await self.underlying.connect() + } + + func send(_ data: Data) async throws { + try await self.underlying.send(data) + } + + func receive() -> AsyncThrowingStream { + AsyncThrowingStream { continuation in + let reader = Task { + do { + for try await data in await self.underlying.receive() { + continuation.yield(data) + } + continuation.finish() + } catch { + continuation.finish(throwing: error) + } + } + continuation.onTermination = { _ in reader.cancel() } + } + } + + func disconnect() async { + self.disconnectCount += 1 + await self.underlying.disconnect() + self.continuation.yield(()) + if !self.released { + await withCheckedContinuation { self.waiters.append($0) } + } + } + + func release() { + self.released = true + for waiter in self.waiters { + waiter.resume() + } + self.waiters.removeAll() + self.continuation.finish() + } +} diff --git a/Core/PeekabooCore/Tests/PeekabooTests/MCP/PeekabooMCPServerTests.swift b/Core/PeekabooCore/Tests/PeekabooTests/MCP/PeekabooMCPServerTests.swift index e2a448f98..7b431b19c 100644 --- a/Core/PeekabooCore/Tests/PeekabooTests/MCP/PeekabooMCPServerTests.swift +++ b/Core/PeekabooCore/Tests/PeekabooTests/MCP/PeekabooMCPServerTests.swift @@ -35,6 +35,37 @@ struct PeekabooMCPServerTests { #expect(!names.contains("swipe")) } + @Test + @MainActor + func `serve runs on a host-supplied transport until it completes`() async throws { + let context = await MCPToolTestHelpers.makeContext() + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + // InMemoryTransport drops messages sent before the receiving peer connects. + try await serverTransport.connect() + let server = try await PeekabooMCPServer(toolContext: context) + let snapshots = await MCPToolUISnapshotStore(owner: server.snapshotOwnerForTesting()) + let snapshot = await snapshots.createSnapshot() + let client = Client(name: "PeekabooHostTransportTests", version: "1.0") + + let serving = Task { try await server.serve(transport: serverTransport) } + do { + _ = try await client.connect(transport: clientTransport) + let (tools, _) = try await client.listTools() + #expect(tools.contains { $0.name == "see" }) + #expect(await snapshots.hasOwnerState()) + } catch { + await client.disconnect() + await server.stopForTesting() + _ = try? await serving.value + throw error + } + + await client.disconnect() + try await serving.value + #expect(await snapshots.getSnapshot(id: snapshot.id) == nil) + #expect(await !snapshots.hasOwnerState()) + } + @Test func `each direct MCP server owns one isolated snapshot namespace`() async throws { let first = try await makeServer() @@ -542,7 +573,8 @@ struct PeekabooMCPServerTests { @Test @MainActor func `wire decoder rejects non-object tool calls and preserves omitted arguments`() async throws { - let context = await MCPToolTestHelpers.makeContext() + let context = await MCPToolTestHelpers.makeContext( + permissionsStatusProvider: WireDecoderPermissionsStatusProvider()) let session = try await MCPWireSession.connect(context: context) do { @@ -762,12 +794,21 @@ struct PeekabooMCPServerTests { } } +@MainActor +private struct WireDecoderPermissionsStatusProvider: PermissionsStatusProviding { + func permissionsStatus() async throws -> PermissionsStatus { + PermissionsStatus(screenRecording: true, accessibility: true, postEvent: true) + } +} + private struct MCPWireSession { let client: Client let server: PeekabooMCPServer static func connect(context: MCPToolContext) async throws -> Self { let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + // InMemoryTransport drops messages sent before the receiving peer connects. + try await serverTransport.connect() let server = try await PeekabooMCPServer(toolContext: context) let client = Client(name: "PeekabooClickWireTests", version: "1.0") try await server.startForTesting(transport: serverTransport) diff --git a/docs/MCP.md b/docs/MCP.md index 2064debca..299b9e8ff 100644 --- a/docs/MCP.md +++ b/docs/MCP.md @@ -57,6 +57,16 @@ Supported transports: - **stdio**: supported and default. - **http / sse**: recognized flags, but server transports are not implemented yet. +Applications linking `PeekabooCore` can pass an MCP Swift SDK `Transport` to +`PeekabooMCPServer.serve(transport:)`. The host owns connection setup and transport policy; `serve` waits for +the transport to complete, then disconnects the SDK session and releases its tool context. Startup failures also +disconnect the session and release that context. Shutdown closes tool-call admission, cancels accepted calls, and +waits for them to finish before releasing their context; cancelling the serving task uses the same cleanup path. +Host-provided transports and tool services must cooperate with task cancellation. Use one server instance per accepted connection; +concurrent `serve` calls on the same instance are rejected before connecting a second transport. A successful +return confirms cleanup; incomplete cleanup throws. This entry point does not add a built-in HTTP or SSE server. +For an accepted `NetworkTransport` connection, disable SDK reconnection so a disconnected peer ends the session. + Peekaboo validates numeric arguments before a tool or mutation lane runs. Fields published as `integer` accept exact whole values (including whole-number JSON doubles and integer strings) but reject fractional, non-finite, and out-of-range values. Fields published as `number` must be finite. Rejections report `mutation_dispatched: false` and