Skip to content

Commit 494b5b5

Browse files
authored
Add configuration for maximum receive buffer size to NetworkTransport (#146)
1 parent 36865fb commit 494b5b5

File tree

1 file changed

+37
-4
lines changed

1 file changed

+37
-4
lines changed

Sources/MCP/Base/Transports/NetworkTransport.swift

Lines changed: 37 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -48,6 +48,12 @@ import Logging
4848
/// // Initialize the transport with the connection
4949
/// let transport = NetworkTransport(connection: connection)
5050
///
51+
/// // For large messages (e.g., images), configure unlimited buffer size
52+
/// let largeBufferTransport = NetworkTransport(
53+
/// connection: connection,
54+
/// bufferConfig: .unlimited
55+
/// )
56+
///
5157
/// // Use the transport with an MCP client
5258
/// let client = Client(name: "MyApp", version: "1.0.0")
5359
/// try await client.connect(transport: transport)
@@ -207,6 +213,26 @@ import Logging
207213
}
208214
}
209215

216+
/// Configuration for buffer behavior.
217+
public struct BufferConfiguration: Hashable, Sendable {
218+
/// Maximum buffer size for receiving data chunks.
219+
/// Set to nil for unlimited (uses system default).
220+
public let maxReceiveBufferSize: Int?
221+
222+
/// Creates a new buffer configuration.
223+
///
224+
/// - Parameter maxReceiveBufferSize: Maximum buffer size in bytes (default: 10MB, nil for unlimited)
225+
public init(maxReceiveBufferSize: Int? = 10 * 1024 * 1024) {
226+
self.maxReceiveBufferSize = maxReceiveBufferSize
227+
}
228+
229+
/// Default buffer configuration with 10MB limit.
230+
public static let `default` = BufferConfiguration()
231+
232+
/// Configuration with no buffer size limit.
233+
public static let unlimited = BufferConfiguration(maxReceiveBufferSize: nil)
234+
}
235+
210236
// State tracking
211237
private var isConnected = false
212238
private var isStopping = false
@@ -228,6 +254,7 @@ import Logging
228254
// Configuration
229255
private let heartbeatConfig: HeartbeatConfiguration
230256
private let reconnectionConfig: ReconnectionConfiguration
257+
private let bufferConfig: BufferConfiguration
231258

232259
/// Creates a new NetworkTransport with the specified NWConnection
233260
///
@@ -236,25 +263,29 @@ import Logging
236263
/// - logger: Optional logger instance for transport events
237264
/// - reconnectionConfig: Configuration for reconnection behavior (default: .default)
238265
/// - heartbeatConfig: Configuration for heartbeat behavior (default: .default)
266+
/// - bufferConfig: Configuration for buffer behavior (default: .default)
239267
public init(
240268
connection: NWConnection,
241269
logger: Logger? = nil,
242270
heartbeatConfig: HeartbeatConfiguration = .default,
243-
reconnectionConfig: ReconnectionConfiguration = .default
271+
reconnectionConfig: ReconnectionConfiguration = .default,
272+
bufferConfig: BufferConfiguration = .default
244273
) {
245274
self.init(
246275
connection,
247276
logger: logger,
248277
heartbeatConfig: heartbeatConfig,
249-
reconnectionConfig: reconnectionConfig
278+
reconnectionConfig: reconnectionConfig,
279+
bufferConfig: bufferConfig
250280
)
251281
}
252282

253283
init(
254284
_ connection: NetworkConnectionProtocol,
255285
logger: Logger? = nil,
256286
heartbeatConfig: HeartbeatConfiguration = .default,
257-
reconnectionConfig: ReconnectionConfiguration = .default
287+
reconnectionConfig: ReconnectionConfiguration = .default,
288+
bufferConfig: BufferConfiguration = .default
258289
) {
259290
self.connection = connection
260291
self.logger =
@@ -265,6 +296,7 @@ import Logging
265296
)
266297
self.reconnectionConfig = reconnectionConfig
267298
self.heartbeatConfig = heartbeatConfig
299+
self.bufferConfig = bufferConfig
268300

269301
// Create message stream
270302
var continuation: AsyncThrowingStream<Data, Swift.Error>.Continuation!
@@ -773,7 +805,8 @@ import Logging
773805
return
774806
}
775807

776-
connection.receive(minimumIncompleteLength: 1, maximumLength: 65536) {
808+
let maxLength = bufferConfig.maxReceiveBufferSize ?? Int.max
809+
connection.receive(minimumIncompleteLength: 1, maximumLength: maxLength) {
777810
content, _, isComplete, error in
778811
Task { @MainActor in
779812
if !receiveContinuationResumed {

0 commit comments

Comments
 (0)