Skip to content

Commit b761d3e

Browse files
authored
Properly propagate errors in NetworkTransport receive stream (#4)
Changes AsyncStream to AsyncThrowingStream to ensure network errors are correctly propagated to callers. Updates receive() to handle errors from the underlying message stream.
1 parent 1f5f8d6 commit b761d3e

File tree

1 file changed

+12
-7
lines changed

1 file changed

+12
-7
lines changed

Sources/MCP/Base/Transports.swift

Lines changed: 12 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -180,8 +180,8 @@ public actor StdioTransport: Transport {
180180
public nonisolated let logger: Logger
181181

182182
private var isConnected = false
183-
private let messageStream: AsyncStream<String>
184-
private let messageContinuation: AsyncStream<String>.Continuation
183+
private let messageStream: AsyncThrowingStream<String, Swift.Error>
184+
private let messageContinuation: AsyncThrowingStream<String, Swift.Error>.Continuation
185185

186186
// Track connection state for continuations
187187
private var connectionContinuationResumed = false
@@ -196,8 +196,8 @@ public actor StdioTransport: Transport {
196196
)
197197

198198
// Create message stream
199-
var continuation: AsyncStream<String>.Continuation!
200-
self.messageStream = AsyncStream { continuation = $0 }
199+
var continuation: AsyncThrowingStream<String, Swift.Error>.Continuation!
200+
self.messageStream = AsyncThrowingStream { continuation = $0 }
201201
self.messageContinuation = continuation
202202
}
203203

@@ -330,10 +330,14 @@ public actor StdioTransport: Transport {
330330
public func receive() -> AsyncThrowingStream<String, Swift.Error> {
331331
return AsyncThrowingStream { continuation in
332332
Task {
333-
for await message in messageStream {
334-
continuation.yield(message)
333+
do {
334+
for try await message in messageStream {
335+
continuation.yield(message)
336+
}
337+
continuation.finish()
338+
} catch {
339+
continuation.finish(throwing: error)
335340
}
336-
continuation.finish()
337341
}
338342
}
339343
}
@@ -361,6 +365,7 @@ public actor StdioTransport: Transport {
361365
} catch {
362366
if !Task.isCancelled {
363367
logger.error("Receive error: \(error)")
368+
messageContinuation.finish(throwing: error)
364369
}
365370
break
366371
}

0 commit comments

Comments
 (0)