modelcontextprotocol/swift-sdk
Watch
1
0
Fork
You've already forked swift-sdk
 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Update transports to send and receive data instead of strings (#42)

Browse source 
* Update transports to send and receive data instead of strings

* Rename MockTransport helper methods

* Adopt Task.sleep(for:...) instead of variant taking nanoseconds

* Update comment

* Revert inadvertant change to visibility of send method
This commit is contained in:
Mattt 2025-04-03 04:42:56 -07:00 • committed by GitHub
commit a9bcbf9ad9
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 138 additions and 147 deletions
Download patch file Download diff file Expand all files Collapse all files

66
Sources/MCP/Base/Transports.swift
View file

@ -19,11 +19,11 @@ public protocol Transport: Actor {
/// Disconnects from the transport
func disconnect() async
/// Sends a message string
func send(_ message: String) async throws
/// Sends data
func send(_ data: Data) async throws
/// Receives message strings as an async sequence
func receive() -> AsyncThrowingStream<String, Swift.Error>
/// Receives data in an async sequence
func receive() -> AsyncThrowingStream<Data, Swift.Error>
}
/// Standard input/output transport implementation
@ -33,8 +33,8 @@ public actor StdioTransport: Transport {
public nonisolated let logger: Logger
private var isConnected = false
private let messageStream: AsyncStream<String>
private let messageContinuation: AsyncStream<String>.Continuation
private let messageStream: AsyncStream<Data>
private let messageContinuation: AsyncStream<Data>.Continuation
public init(
input: FileDescriptor = FileDescriptor.standardInput,
@ -50,7 +50,7 @@ public actor StdioTransport: Transport {
factory: { _ in SwiftLogNoOpLogHandler() })
// Create message stream
var continuation: AsyncStream<String>.Continuation!
var continuation: AsyncStream<Data>.Continuation!
self.messageStream = AsyncStream { continuation = $0 }
self.messageContinuation = continuation
}
@ -105,15 +105,13 @@ public actor StdioTransport: Transport {
let messageData = pendingData[..<newlineIndex]
pendingData = pendingData[(newlineIndex + 1)...]
if let message = String(data: messageData, encoding: .utf8),
!message.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty
{
logger.debug("Message received", metadata: ["message": "\(message)"])
messageContinuation.yield(message)
if !messageData.isEmpty {
logger.debug("Message received", metadata: ["size": "\(messageData.count)"])
messageContinuation.yield(Data(messageData))
}
}
} catch let error where Error.isResourceTemporarilyUnavailable(error) {
try? await Task.sleep(nanoseconds: 10_000_000) // 10ms backoff
try? await Task.sleep(for: .milliseconds(10))
continue
} catch {
if !Task.isCancelled {
@ -133,17 +131,16 @@ public actor StdioTransport: Transport {
logger.info("Transport disconnected")
}
public func send(_ message: String) async throws {
public func send(_ message: Data) async throws {
guard isConnected else {
throw Error.transportError(Errno.socketNotConnected)
}
let message = message + "\n"
guard let data = message.data(using: .utf8) else {
throw Error.transportError(Errno.invalidArgument)
}
// Add newline as delimiter
var messageWithNewline = message
messageWithNewline.append(UInt8(ascii: "\n"))
var remaining = data
var remaining = messageWithNewline
while !remaining.isEmpty {
do {
let written = try remaining.withUnsafeBytes { buffer in
@ -153,7 +150,7 @@ public actor StdioTransport: Transport {
remaining = remaining.dropFirst(written)
}
} catch let error where Error.isResourceTemporarilyUnavailable(error) {
try await Task.sleep(nanoseconds: 10_000_000) // 10ms backoff
try await Task.sleep(for: .milliseconds(10))
continue
} catch {
throw Error.transportError(error)
@ -161,7 +158,7 @@ public actor StdioTransport: Transport {
}
}
public func receive() -> AsyncThrowingStream<String, Swift.Error> {
public func receive() -> AsyncThrowingStream<Data, Swift.Error> {
return AsyncThrowingStream { continuation in
Task {
for await message in messageStream {
@ -182,8 +179,8 @@ public actor StdioTransport: Transport {
public nonisolated let logger: Logger
private var isConnected = false
private let messageStream: AsyncThrowingStream<String, Swift.Error>
private let messageContinuation: AsyncThrowingStream<String, Swift.Error>.Continuation
private let messageStream: AsyncThrowingStream<Data, Swift.Error>
private let messageContinuation: AsyncThrowingStream<Data, Swift.Error>.Continuation
// Track connection state for continuations
private var connectionContinuationResumed = false
@ -198,7 +195,7 @@ public actor StdioTransport: Transport {
)
// Create message stream
var continuation: AsyncThrowingStream<String, Swift.Error>.Continuation!
var continuation: AsyncThrowingStream<Data, Swift.Error>.Continuation!
self.messageStream = AsyncThrowingStream { continuation = $0 }
self.messageContinuation = continuation
}
@ -289,14 +286,14 @@ public actor StdioTransport: Transport {
logger.info("Network transport disconnected")
}
public func send(_ message: String) async throws {
public func send(_ message: Data) async throws {
guard isConnected else {
throw MCP.Error.internalError("Transport not connected")
}
guard let data = (message + "\n").data(using: .utf8) else {
throw MCP.Error.internalError("Failed to encode message")
}
// Add newline as delimiter
var messageWithNewline = message
messageWithNewline.append(UInt8(ascii: "\n"))
// Use a local actor-isolated variable to track continuation state
var sendContinuationResumed = false
@ -309,7 +306,7 @@ public actor StdioTransport: Transport {
}
connection.send(
content: data,
content: messageWithNewline,
completion: .contentProcessed { [weak self] error in
guard let self = self else { return }
@ -329,7 +326,7 @@ public actor StdioTransport: Transport {
}
}
public func receive() -> AsyncThrowingStream<String, Swift.Error> {
public func receive() -> AsyncThrowingStream<Data, Swift.Error> {
return AsyncThrowingStream { continuation in
Task {
do {
@ -357,11 +354,10 @@ public actor StdioTransport: Transport {
let messageData = buffer[..<newlineIndex]
buffer = buffer[(newlineIndex + 1)...]
if let message = String(data: messageData, encoding: .utf8),
!message.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty
{
logger.debug("Message received", metadata: ["message": "\(message)"])
messageContinuation.yield(message)
if !messageData.isEmpty {
logger.debug(
"Message received", metadata: ["size": "\(messageData.count)"])
messageContinuation.yield(Data(messageData))
}
}
} catch let error as NWError {