Skip to content
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
13 changes: 13 additions & 0 deletions Sources/MCP/Base/Transports/StdioTransport.swift
Original file line number Diff line number Diff line change
Expand Up @@ -56,6 +56,7 @@ import struct Foundation.Data
private var isConnected = false
private let messageStream: AsyncThrowingStream<Data, Swift.Error>
private let messageContinuation: AsyncThrowingStream<Data, Swift.Error>.Continuation
private var lastSend: Task<Void, Never>?

/// Creates a new stdio transport with the specified file descriptors
///
Expand Down Expand Up @@ -195,6 +196,18 @@ import struct Foundation.Data
/// - Parameter message: The message data to send (without a trailing newline)
/// - Throws: Error if the message cannot be sent
public func send(_ message: Data) async throws {
let previousSend = lastSend
let currentSend = Task {
await previousSend?.value
try await write(message)
}
lastSend = Task {
try? await currentSend.value
}
try await currentSend.value
}

private func write(_ message: Data) async throws {
guard isConnected else {
throw MCPError.transportError(Errno(rawValue: ENOTCONN))
}
Expand Down
44 changes: 44 additions & 0 deletions Tests/MCPTests/StdioTransportTests.swift
Original file line number Diff line number Diff line change
Expand Up @@ -43,6 +43,50 @@ struct StdioTransportTests {
await transport.disconnect()
}

@Test("Concurrent sends preserve message framing under backpressure")
func testConcurrentSendsPreserveMessageFramingUnderBackpressure() async throws {
let (reader, output) = try FileDescriptor.pipe()
let (input, _) = try FileDescriptor.pipe()
let transport = StdioTransport(input: input, output: output, logger: nil)
try await transport.connect()

let largeMessage = Data(repeating: UInt8(ascii: "a"), count: 512 * 1024)
let smallMessage = Data(#"{"id":2}"#.utf8)
let expected =
largeMessage + Data([UInt8(ascii: "\n")])
+ smallMessage + Data([UInt8(ascii: "\n")])

let firstSend = Task {
try await transport.send(largeMessage)
}

// Leave the pipe undrained until the first send has filled its buffer
// and suspended in the EAGAIN retry path.
try await Task.sleep(for: .milliseconds(50))

let secondSend = Task {
try await transport.send(smallMessage)
}

let received = try await Task.detached {
var received = Data()
var buffer = [UInt8](repeating: 0, count: 4096)
while received.count < expected.count {
let count = try buffer.withUnsafeMutableBufferPointer { pointer in
try reader.read(into: UnsafeMutableRawBufferPointer(pointer))
}
received.append(contentsOf: buffer[..<count])
}
return received
}.value

try await firstSend.value
try await secondSend.value
#expect(received == expected)

await transport.disconnect()
}

@Test("Receive Message")
func testStdioTransportReceiveMessage() async throws {
let (input, writer) = try FileDescriptor.pipe()
Expand Down