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
7 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 {

20
Sources/MCP/Client/Client.swift
View file

@ -181,15 +181,10 @@ public actor Client {
do {
let stream = await connection.receive()
for try await string in stream {
for try await data in stream {
if Task.isCancelled { break } // Check inside loop too
// Decode and handle incoming message
guard let data = string.data(using: .utf8) else {
throw Error.parseError("Invalid UTF-8 data")
}
// Attempt to decode string data as AnyResponse or AnyMessage
// Attempt to decode data as AnyResponse or AnyMessage
let decoder = JSONDecoder()
if let response = try? decoder.decode(AnyResponse.self, from: data),
let request = pendingRequests[response.id]
@ -207,7 +202,7 @@ public actor Client {
}
}
} catch let error where Error.isResourceTemporarilyUnavailable(error) {
try? await Task.sleep(nanoseconds: 10_000_000) // 10ms
try? await Task.sleep(for: .milliseconds(10))
continue
} catch {
await logger?.error(
@ -256,12 +251,9 @@ public actor Client {
}
let requestData = try JSONEncoder().encode(request)
guard let requestString = String(data: requestData, encoding: .utf8) else {
throw Error.internalError("Failed to encode request")
}
// Store the pending request first
return try await withCheckedThrowingContinuation { continuation in
// Store the pending request first
Task {
self.addPendingRequest(
id: request.id,
@ -269,9 +261,9 @@ public actor Client {
type: M.Result.self
)
// Send the request
// Send the request data
do {
try await connection.send(requestString)
try await connection.send(requestData)
} catch {
continuation.resume(throwing: error)
self.removePendingRequest(id: request.id)

23
Sources/MCP/Server/Server.swift
View file

@ -172,15 +172,11 @@ public actor Server {
task = Task {
do {
let stream = await transport.receive()
for try await string in stream {
for try await data in stream {
if Task.isCancelled { break } // Check cancellation inside loop
var requestID: ID?
do {
guard let data = string.data(using: .utf8) else {
throw Error.parseError("Invalid UTF-8 data")
}
// Attempt to decode string data as AnyRequest or AnyMessage
let decoder = JSONDecoder()
if let request = try? decoder.decode(AnyRequest.self, from: data) {
@ -203,7 +199,7 @@ public actor Server {
}
} catch let error where Error.isResourceTemporarilyUnavailable(error) {
// Resource temporarily unavailable, retry after a short delay
try? await Task.sleep(nanoseconds: 10_000_000) // 10ms
try? await Task.sleep(for: .milliseconds(10))
continue
} catch {
await logger?.error(
@ -266,19 +262,17 @@ public actor Server {
// MARK: - Sending
/// Send a response to a client
/// Send a response to a request
public func send<M: Method>(_ response: Response<M>) async throws {
guard let connection = connection else {
throw Error.internalError("Server connection not initialized")
}
let encoder = JSONEncoder()
encoder.outputFormatting = [.sortedKeys, .withoutEscapingSlashes]
let responseData = try encoder.encode(response)
if let responseStr = String(data: responseData, encoding: .utf8) {
try await connection.send(responseStr)
}
try await connection.send(responseData)
}
/// Send a notification to connected clients
@ -291,10 +285,7 @@ public actor Server {
encoder.outputFormatting = [.sortedKeys, .withoutEscapingSlashes]
let notificationData = try encoder.encode(notification)
if let notificationStr = String(data: notificationData, encoding: .utf8) {
try await connection.send(notificationStr)
}
try await connection.send(notificationData)
}
// MARK: -
@ -407,7 +398,7 @@ public actor Server {
// Send initialized notification after a short delay
Task {
try? await Task.sleep(nanoseconds: 100_000_000) // 100ms
try? await Task.sleep(for: .milliseconds(10))
try? await self.notify(InitializedNotification.message())
}

24
Tests/MCPTests/ClientTests.swift
View file

@ -27,7 +27,7 @@ struct ClientTests {
try await client.connect(transport: transport)
// Small delay to ensure message loop is started
try await Task.sleep(nanoseconds: 10_000_000) // 10ms
try await Task.sleep(for: .milliseconds(10))
// Create a task for initialize that we'll cancel
let initTask = Task {
@ -35,19 +35,19 @@ struct ClientTests {
}
// Give it a moment to send the request
try await Task.sleep(nanoseconds: 10_000_000) // 10ms
try await Task.sleep(for: .milliseconds(10))
#expect(await transport.sentMessages.count == 1)
#expect(await transport.sentMessages[0].contains(Initialize.name))
#expect(await transport.sentMessages[0].contains(client.name))
#expect(await transport.sentMessages[0].contains(client.version))
#expect(await transport.sentMessages.first?.contains(Initialize.name) == true)
#expect(await transport.sentMessages.first?.contains(client.name) == true)
#expect(await transport.sentMessages.first?.contains(client.version) == true)
// Cancel the initialize task
initTask.cancel()
// Disconnect client to clean up message loop and give time for continuation cleanup
await client.disconnect()
try await Task.sleep(nanoseconds: 50_000_000) // 50ms
try await Task.sleep(for: .milliseconds(50))
}
@Test(
@ -60,7 +60,7 @@ struct ClientTests {
try await client.connect(transport: transport)
// Small delay to ensure message loop is started
try await Task.sleep(nanoseconds: 10_000_000) // 10ms
try await Task.sleep(for: .milliseconds(10))
// Create a task for the ping that we'll cancel
let pingTask = Task {
@ -68,17 +68,17 @@ struct ClientTests {
}
// Give it a moment to send the request
try await Task.sleep(nanoseconds: 10_000_000) // 10ms
try await Task.sleep(for: .milliseconds(10))
#expect(await transport.sentMessages.count == 1)
#expect(await transport.sentMessages[0].contains(Ping.name))
#expect(await transport.sentMessages.first?.contains(Ping.name) == true)
// Cancel the ping task
pingTask.cancel()
// Disconnect client to clean up message loop and give time for continuation cleanup
await client.disconnect()
try await Task.sleep(nanoseconds: 50_000_000) // 50ms
try await Task.sleep(for: .milliseconds(50))
}
@Test("Connection failure handling")
@ -168,7 +168,7 @@ struct ClientTests {
// Wait a bit for any setup to complete
try await Task.sleep(for: .milliseconds(10))
// Send the listPrompts request and immediately provide an error response
let promptsTask = Task {
do {
@ -187,7 +187,7 @@ struct ClientTests {
id: decodedRequest.id,
error: Error.methodNotFound("Test: Prompts capability not available")
)
try await transport.queueResponse(errorResponse)
try await transport.queue(response: errorResponse)
// Try the request now that we have a response queued
do {

131
Tests/MCPTests/Helpers/MockTransport.swift
View file

@ -6,10 +6,28 @@ import Logging
/// Mock transport for testing
actor MockTransport: Transport {
var logger: Logger
let encoder = JSONEncoder()
let decoder = JSONDecoder()
var isConnected = false
private(set) var sentMessages: [String] = []
private var messagesToReceive: [String] = []
private var messageStreamContinuation: AsyncThrowingStream<String, Swift.Error>.Continuation?
private(set) var sentData: [Data] = []
var sentMessages: [String] {
return sentData.compactMap { data in
guard let string = String(data: data, encoding: .utf8) else {
logger.error("Failed to decode sent data as UTF-8")
return nil
}
return string
}
}
private var dataToReceive: [Data] = []
private(set) var receivedMessages: [String] = []
private var dataStreamContinuation: AsyncThrowingStream<Data, Swift.Error>.Continuation?
var shouldFailConnect = false
var shouldFailSend = false
@ -17,85 +35,39 @@ actor MockTransport: Transport {
self.logger = logger
}
func connect() async throws {
public func connect() async throws {
if shouldFailConnect {
throw Error.transportError(POSIXError(.ECONNREFUSED))
}
isConnected = true
}
func disconnect() async {
public func disconnect() async {
isConnected = false
messageStreamContinuation?.finish()
messageStreamContinuation = nil
dataStreamContinuation?.finish()
dataStreamContinuation = nil
}
func send<T: Encodable & Sendable>(_ message: T) async throws {
public func send(_ message: Data) async throws {
if shouldFailSend {
throw Error.transportError(POSIXError(.EIO))
}
let data = try JSONEncoder().encode(message)
let str = String(data: data, encoding: .utf8)!
sentMessages.append(str)
sentData.append(message)
}
func receive() -> AsyncThrowingStream<String, Swift.Error> {
return AsyncThrowingStream<String, Swift.Error> { continuation in
messageStreamContinuation = continuation
// Send any queued messages
for message in messagesToReceive {
public func receive() -> AsyncThrowingStream<Data, Swift.Error> {
return AsyncThrowingStream<Data, Swift.Error> { continuation in
dataStreamContinuation = continuation
for message in dataToReceive {
continuation.yield(message)
if let string = String(data: message, encoding: .utf8) {
receivedMessages.append(string)
}
}
messagesToReceive.removeAll()
dataToReceive.removeAll()
}
}
func queueRequest<M: MCP.Method>(_ request: Request<M>) throws {
let data = try JSONEncoder().encode(request)
let str = String(data: data, encoding: .utf8)!
if let continuation = messageStreamContinuation {
continuation.yield(str)
} else {
sentMessages.append(str)
}
}
func queueResponse<M: MCP.Method>(_ response: Response<M>) throws {
let data = try JSONEncoder().encode(response)
let str = String(data: data, encoding: .utf8)!
if let continuation = messageStreamContinuation {
continuation.yield(str)
} else {
messagesToReceive.append(str)
}
}
func queueNotification<N: MCP.Notification>(_ notification: Message<N>) throws {
let data = try JSONEncoder().encode(notification)
let str = String(data: data, encoding: .utf8)!
if let continuation = messageStreamContinuation {
continuation.yield(str)
} else {
messagesToReceive.append(str)
}
}
func getLastSentMessage<T: Decodable>() -> T? {
print("SENT:", sentMessages)
guard let lastMessage = sentMessages.last else { return nil }
do {
let data = lastMessage.data(using: .utf8)!
return try JSONDecoder().decode(T.self, from: data)
} catch {
return nil
}
}
func clearMessages() {
sentMessages.removeAll()
messagesToReceive.removeAll()
}
func setFailConnect(_ shouldFail: Bool) {
shouldFailConnect = shouldFail
}
@ -103,4 +75,37 @@ actor MockTransport: Transport {
func setFailSend(_ shouldFail: Bool) {
shouldFailSend = shouldFail
}
func queue<M: MCP.Method>(request: Request<M>) throws {
let data = try encoder.encode(request)
if let continuation = dataStreamContinuation {
continuation.yield(data)
} else {
dataToReceive.append(data)
}
}
func queue<M: MCP.Method>(response: Response<M>) throws {
let data = try encoder.encode(response)
dataToReceive.append(data)
}
func queue<N: MCP.Notification>(notification: Message<N>) throws {
let data = try encoder.encode(notification)
dataToReceive.append(data)
}
func decodeLastSentMessage<T: Decodable>() -> T? {
guard let lastMessage = sentData.last else { return nil }
do {
return try decoder.decode(T.self, from: lastMessage)
} catch {
return nil
}
}
func clearMessages() {
sentData.removeAll()
dataToReceive.removeAll()
}
}

31
Tests/MCPTests/ServerTests.swift
View file

@ -22,8 +22,8 @@ struct ServerTests {
let transport = MockTransport()
// Queue an initialize request
try await transport.queueRequest(
Initialize.request(
try await transport.queue(
request: Initialize.request(
.init(
protocolVersion: Version.latest,
capabilities: .init(),
@ -42,7 +42,11 @@ struct ServerTests {
try await Task.sleep(nanoseconds: 10_000_000) // 10ms
#expect(await transport.sentMessages.count == 1)
#expect(await transport.sentMessages[0].contains(Initialize.name))
let messages = await transport.sentMessages
if let response = messages.first {
#expect(response.contains("serverInfo"))
}
// Clean up
await server.stop()
@ -68,13 +72,13 @@ struct ServerTests {
#expect(clientInfo.version == "1.0")
await state.setHookCalled()
}
// Wait for server to initialize
try await Task.sleep(nanoseconds: 10_000_000) // 10ms
try await Task.sleep(for: .milliseconds(10))
// Queue an initialize request
try await transport.queueRequest(
Initialize.request(
try await transport.queue(
request: Initialize.request(
.init(
protocolVersion: Version.latest,
capabilities: .init(),
@ -83,7 +87,7 @@ struct ServerTests {
))
// Wait for message processing and hook execution
try await Task.sleep(nanoseconds: 200_000_000) // 200ms
try await Task.sleep(for: .milliseconds(500))
#expect(await state.wasHookCalled() == true)
#expect(await transport.sentMessages.count >= 1)
@ -94,6 +98,7 @@ struct ServerTests {
}
await server.stop()
await transport.disconnect()
}
@Test("Initialize hook - rejection")
@ -107,13 +112,13 @@ struct ServerTests {
throw Error.invalidRequest("Client not allowed")
}
}
// Wait for server to initialize
try await Task.sleep(nanoseconds: 10_000_000) // 10ms
// Queue an initialize request from blocked client
try await transport.queueRequest(
Initialize.request(
try await transport.queue(
request: Initialize.request(
.init(
protocolVersion: Version.latest,
capabilities: .init(),
@ -124,13 +129,15 @@ struct ServerTests {
// Wait for message processing
try await Task.sleep(nanoseconds: 200_000_000) // 200ms
#expect(await transport.sentMessages.count >= 2)
#expect(await transport.sentMessages.count >= 1)
let messages = await transport.sentMessages
if let response = messages.first {
#expect(response.contains("error"))
#expect(response.contains("Client not allowed"))
}
await server.stop()
await transport.disconnect()
}
}

8
Tests/MCPTests/TransportTests.swift
View file

@ -29,7 +29,7 @@ struct StdioTransportTests {
// Test sending a simple message
let message = #"{"key":"value"}"#
try await transport.send(message)
try await transport.send(message.data(using: .utf8)!)
// Read and verify the output
var buffer = [UInt8](repeating: 0, count: 1024)
@ -57,12 +57,12 @@ struct StdioTransportTests {
try writer.close()
// Start receiving messages
let stream: AsyncThrowingStream<String, Swift.Error> = await transport.receive()
let stream: AsyncThrowingStream<Data, Swift.Error> = await transport.receive()
var iterator = stream.makeAsyncIterator()
// Get first message
let received = try await iterator.next()
#expect(received == #"{"key":"value"}"#)
#expect(received == #"{"key":"value"}"#.data(using: .utf8)!)
await transport.disconnect()
}
@ -79,7 +79,7 @@ struct StdioTransportTests {
try writer.writeAll(invalidJSON.data(using: .utf8)!)
try writer.close()
let stream: AsyncThrowingStream<String, Swift.Error> = await transport.receive()
let stream: AsyncThrowingStream<Data, Swift.Error> = await transport.receive()
var iterator = stream.makeAsyncIterator()
_ = try await iterator.next()