From d8ae61be3422867ecf74215e7d0563ed0cb1e226 Mon Sep 17 00:00:00 2001 From: Semyon Kotliarenko Date: Sat, 12 Sep 2026 20:51:31 +0200 Subject: [PATCH 1/2] feat(mcp): serve over a host-supplied transport MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add `PeekabooMCPServer.serve(transport: any Transport)` 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 built-in stdio path and the new entry point share one lifecycle: the server runs until the transport completes and the tool context is released on the way out, on success and on failure alike. `serve(transport: TransportType, port:)` keeps its behaviour, including teardown on the not-implemented transports. --- CHANGELOG.md | 2 ++ .../MCP/Server/PeekabooMCPServer.swift | 22 +++++++++++++++++-- .../MCP/PeekabooMCPServerTests.swift | 17 ++++++++++++++ 3 files changed, 39 insertions(+), 2 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index c9ea76cfb..5346614f5 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,8 @@ ## Unreleased +- Let hosts that embed `PeekabooMCPServer` serve over a transport they supply (`serve(transport: any Transport)`), with the same run-until-complete lifecycle and tool-context teardown as the built-in stdio path. #716. + ## 4.3.4 - 2026-09-11 **Highlights:** Restore provider-compatible MCP tools and frontmost daemon captures. diff --git a/Core/PeekabooCore/Sources/PeekabooAgentRuntime/MCP/Server/PeekabooMCPServer.swift b/Core/PeekabooCore/Sources/PeekabooAgentRuntime/MCP/Server/PeekabooMCPServer.swift index 63a94d907..d2050fab1 100644 --- a/Core/PeekabooCore/Sources/PeekabooAgentRuntime/MCP/Server/PeekabooMCPServer.swift +++ b/Core/PeekabooCore/Sources/PeekabooAgentRuntime/MCP/Server/PeekabooMCPServer.swift @@ -288,8 +288,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,6 +298,24 @@ public actor PeekabooMCPServer { case .sse: throw MCPError.notImplemented("SSE server transport not yet implemented") } + } + } + + /// 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 } + } + + private func run(makingTransport: () throws -> any Transport) async throws { + do { + let serverTransport = try makingTransport() try await self.server.start(transport: serverTransport) // Keep the server running diff --git a/Core/PeekabooCore/Tests/PeekabooTests/MCP/PeekabooMCPServerTests.swift b/Core/PeekabooCore/Tests/PeekabooTests/MCP/PeekabooMCPServerTests.swift index e2a448f98..a006cf29c 100644 --- a/Core/PeekabooCore/Tests/PeekabooTests/MCP/PeekabooMCPServerTests.swift +++ b/Core/PeekabooCore/Tests/PeekabooTests/MCP/PeekabooMCPServerTests.swift @@ -35,6 +35,23 @@ 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() + let server = try await PeekabooMCPServer(toolContext: context) + let client = Client(name: "PeekabooHostTransportTests", version: "1.0") + + let serving = Task { try await server.serve(transport: serverTransport) } + _ = try await client.connect(transport: clientTransport) + let (tools, _) = try await client.listTools() + #expect(tools.contains { $0.name == "see" }) + + await client.disconnect() + try await serving.value + } + @Test func `each direct MCP server owns one isolated snapshot namespace`() async throws { let first = try await makeServer() From ad4b686b69c32d03efcca4e3fdbb13305cdcfc4c Mon Sep 17 00:00:00 2001 From: Peter Steinberger Date: Sat, 12 Sep 2026 16:15:49 -0700 Subject: [PATCH 2/2] fix(mcp): disconnect completed and failed server transports SDK completion only awaits the receive loop and retains its transport. Stop the SDK session before releasing tool context on both completion and startup failure, with disconnect regressions and live TCP/stdio EOF proof. --- .../MCP/Server/PeekabooMCPServer.swift | 2 ++ ...ft => PeekabooMCPHostTransportTests.swift} | 36 +++++++++++++++---- docs/MCP.md | 3 +- 3 files changed, 34 insertions(+), 7 deletions(-) rename Core/PeekabooCore/Tests/PeekabooTests/MCP/{PeekabooMCPHostTransportFailureTests.swift => PeekabooMCPHostTransportTests.swift} (55%) diff --git a/Core/PeekabooCore/Sources/PeekabooAgentRuntime/MCP/Server/PeekabooMCPServer.swift b/Core/PeekabooCore/Sources/PeekabooAgentRuntime/MCP/Server/PeekabooMCPServer.swift index d2050fab1..88662813f 100644 --- a/Core/PeekabooCore/Sources/PeekabooAgentRuntime/MCP/Server/PeekabooMCPServer.swift +++ b/Core/PeekabooCore/Sources/PeekabooAgentRuntime/MCP/Server/PeekabooMCPServer.swift @@ -321,12 +321,14 @@ public actor PeekabooMCPServer { // Keep the server running await self.server.waitUntilCompleted() } catch { + await self.server.stop() let cleanupConfirmed = await self.releaseToolContextForTeardown() if !cleanupConfirmed { self.logger.error("Browser session cleanup remains pending after MCP server failure") } throw error } + await self.server.stop() let cleanupConfirmed = await self.releaseToolContextForTeardown() guard cleanupConfirmed else { throw MCPError.executionFailed( diff --git a/Core/PeekabooCore/Tests/PeekabooTests/MCP/PeekabooMCPHostTransportFailureTests.swift b/Core/PeekabooCore/Tests/PeekabooTests/MCP/PeekabooMCPHostTransportTests.swift similarity index 55% rename from Core/PeekabooCore/Tests/PeekabooTests/MCP/PeekabooMCPHostTransportFailureTests.swift rename to Core/PeekabooCore/Tests/PeekabooTests/MCP/PeekabooMCPHostTransportTests.swift index e93885b31..276e7e66a 100644 --- a/Core/PeekabooCore/Tests/PeekabooTests/MCP/PeekabooMCPHostTransportFailureTests.swift +++ b/Core/PeekabooCore/Tests/PeekabooTests/MCP/PeekabooMCPHostTransportTests.swift @@ -6,7 +6,7 @@ import Testing @testable import PeekabooCore @Suite(.serialized) -struct PeekabooMCPHostTransportFailureTests { +struct PeekabooMCPHostTransportTests { @Test @MainActor func `host transport startup failure releases its snapshot owner`() async throws { @@ -14,38 +14,62 @@ struct PeekabooMCPHostTransportFailureTests { let server = try await PeekabooMCPServer(toolContext: context) let snapshots = await MCPToolUISnapshotStore(owner: server.snapshotOwnerForTesting()) let snapshot = await snapshots.createSnapshot() - let transport = FailingHostTransport() + 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) + } } private enum HostTransportError: Error, Equatable { case connectionFailed } -private actor FailingHostTransport: Transport { +private actor LifecycleHostTransport: Transport { nonisolated let logger = Logger(label: "boo.peekaboo.tests.host-transport") + private let failsToConnect: Bool private(set) var connectCount = 0 + private(set) var disconnectCount = 0 + + init(failsToConnect: Bool = true) { + self.failsToConnect = failsToConnect + } func connect() async throws { self.connectCount += 1 - throw HostTransportError.connectionFailed + if self.failsToConnect { + throw HostTransportError.connectionFailed + } } - func disconnect() async {} + func disconnect() async { + self.disconnectCount += 1 + } func send(_: Data) async throws { throw HostTransportError.connectionFailed } func receive() -> AsyncThrowingStream { - AsyncThrowingStream { $0.finish(throwing: HostTransportError.connectionFailed) } + AsyncThrowingStream { $0.finish() } } } diff --git a/docs/MCP.md b/docs/MCP.md index 54c998b08..c7489a3ef 100644 --- a/docs/MCP.md +++ b/docs/MCP.md @@ -59,7 +59,8 @@ Supported transports: 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 releases its tool context. Startup failures also release that context. A successful +the transport to complete, then disconnects the SDK session and releases its tool context. Startup failures also +disconnect the session and release that context. 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.