Skip to content

Commit 843a4ab

Browse files
doozMenStijn Willems
doozMen
authored and
Stijn Willems
committed
fix continuation (#9)
1 parent 973b46b commit 843a4ab

4 files changed

Lines changed: 116 additions & 85 deletions

File tree

‎.gitignore‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,4 +6,6 @@ DerivedData/
66
.swiftpm/configuration/registries.json
77
.swiftpm/xcode/package.xcworkspace/contents.xcworkspacedata
88
.netrc
9-
.index-build/
9+
.index-build/
10+
.claude/settings.local.json
11+
.swift-version

‎Sources/MCP/Base/Transports/NetworkTransport.swift‎

Lines changed: 51 additions & 62 deletions
Original file line numberDiff line numberDiff line change
@@ -459,7 +459,14 @@ import Logging
459459
content: Heartbeat().data,
460460
contentContext: .defaultMessage,
461461
isComplete: true,
462-
completion: .contentProcessed { [weak self] error in
462+
completion: .contentProcessed { [weak self, continuation] error in
463+
guard let self = self else {
464+
continuation.resume(
465+
throwing: MCPError.internalError(
466+
"Transport deallocated during heartbeat"))
467+
return
468+
}
469+
463470
if let error = error {
464471
continuation.resume(throwing: error)
465472
} else {
@@ -511,9 +518,6 @@ import Logging
511518
var messageWithNewline = message
512519
messageWithNewline.append(UInt8(ascii: "\n"))
513520

514-
// Use a local actor-isolated variable to track continuation state
515-
var sendContinuationResumed = false
516-
517521
try await withCheckedThrowingContinuation {
518522
[weak self] (continuation: CheckedContinuation<Void, Swift.Error>) in
519523
guard let self = self else {
@@ -525,50 +529,43 @@ import Logging
525529
content: messageWithNewline,
526530
contentContext: .defaultMessage,
527531
isComplete: true,
528-
completion: .contentProcessed { [weak self] error in
529-
guard let self = self else { return }
530-
531-
Task { @MainActor in
532-
if !sendContinuationResumed {
533-
sendContinuationResumed = true
534-
if let error = error {
535-
self.logger.error("Send error: \(error)")
536-
537-
// Check if we should attempt to reconnect on send failure
538-
let isStopping = await self.isStopping // Await actor-isolated property
539-
if !isStopping && self.reconnectionConfig.enabled {
540-
let isConnected = await self.isConnected
541-
if isConnected {
542-
if error.isConnectionLost {
543-
self.logger.warning(
544-
"Connection appears broken, will attempt to reconnect..."
545-
)
546-
547-
// Schedule connection restart
548-
Task { [weak self] in // Operate on self's executor
549-
guard let self = self else { return }
550-
551-
await self.setIsConnected(false)
552-
553-
try? await Task.sleep(for: .milliseconds(500))
554-
555-
let currentIsStopping = await self.isStopping
556-
if !currentIsStopping {
557-
// Cancel the connection, then attempt to reconnect fully.
558-
self.connection.cancel()
559-
try? await self.connect()
560-
}
561-
}
562-
}
532+
completion: .contentProcessed { [weak self, continuation] error in
533+
guard let self = self else {
534+
continuation.resume(
535+
throwing: MCPError.internalError(
536+
"Transport deallocated during send"))
537+
return
538+
}
539+
540+
if let error = error {
541+
self.logger.error("Send error: \(error)")
542+
543+
// Schedule reconnection check on a separate task
544+
Task { [weak self] in
545+
guard let self = self else { return }
546+
let isStopping = await self.isStopping
547+
if !isStopping && self.reconnectionConfig.enabled {
548+
let isConnected = await self.isConnected
549+
if isConnected && error.isConnectionLost {
550+
self.logger.warning(
551+
"Connection appears broken, will attempt to reconnect..."
552+
)
553+
await self.setIsConnected(false)
554+
try? await Task.sleep(for: .milliseconds(500))
555+
556+
let currentIsStopping = await self.isStopping
557+
if !currentIsStopping {
558+
self.connection.cancel()
559+
try? await self.connect()
563560
}
564561
}
565-
566-
continuation.resume(
567-
throwing: MCPError.internalError("Send error: \(error)"))
568-
} else {
569-
continuation.resume()
570562
}
571563
}
564+
565+
continuation.resume(
566+
throwing: MCPError.internalError("Send error: \(error)"))
567+
} else {
568+
continuation.resume()
572569
}
573570
})
574571
}
@@ -747,8 +744,6 @@ import Logging
747744
/// - Returns: The received data chunk
748745
/// - Throws: Network errors or transport failures
749746
private func receiveData() async throws -> Data {
750-
var receiveContinuationResumed = false
751-
752747
return try await withCheckedThrowingContinuation {
753748
[weak self] (continuation: CheckedContinuation<Data, Swift.Error>) in
754749
guard let self = self else {
@@ -758,22 +753,16 @@ import Logging
758753

759754
let maxLength = bufferConfig.maxReceiveBufferSize ?? Int.max
760755
connection.receive(minimumIncompleteLength: 1, maximumLength: maxLength) {
761-
content, _, isComplete, error in
762-
Task { @MainActor in
763-
if !receiveContinuationResumed {
764-
receiveContinuationResumed = true
765-
if let error = error {
766-
continuation.resume(throwing: MCPError.transportError(error))
767-
} else if let content = content {
768-
continuation.resume(returning: content)
769-
} else if isComplete {
770-
self.logger.trace("Connection completed by peer")
771-
continuation.resume(throwing: MCPError.connectionClosed)
772-
} else {
773-
// EOF: Resume with empty data instead of throwing an error
774-
continuation.resume(returning: Data())
775-
}
776-
}
756+
[weak self, continuation] content, _, isComplete, error in
757+
if let error = error {
758+
continuation.resume(throwing: MCPError.transportError(error))
759+
} else if let content = content {
760+
continuation.resume(returning: content)
761+
} else if isComplete {
762+
self?.logger.trace("Connection completed by peer")
763+
continuation.resume(throwing: MCPError.connectionClosed)
764+
} else {
765+
continuation.resume(returning: Data())
777766
}
778767
}
779768
}

‎Sources/MCP/Client/Client.swift‎

Lines changed: 8 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -592,24 +592,27 @@ public actor Client {
592592
/// Use this object to add requests to the batch.
593593
/// - Throws: `MCPError.internalError` if the client is not connected.
594594
/// Can also rethrow errors from the `body` closure or from sending the batch request.
595-
public func withBatch(body: @escaping @Sendable (Batch) async throws -> Void) async throws {
595+
@discardableResult
596+
public func withBatch<T: Sendable>(
597+
body: @escaping @Sendable (Batch) async throws -> T
598+
) async throws -> T {
596599
guard let connection = connection else {
597600
throw MCPError.internalError("Client connection not initialized")
598601
}
599602

600603
// Create Batch actor, passing self (Client)
601604
let batch = Batch(client: self)
602605

603-
// Populate the batch actor by calling the user's closure.
604-
try await body(batch)
606+
// Populate the batch actor by calling the user's closure and capture result.
607+
let result = try await body(batch)
605608

606609
// Get the collected requests from the batch actor
607610
let requests = await batch.requests
608611

609612
// Check if there are any requests to send
610613
guard !requests.isEmpty else {
611614
await logger?.debug("Batch requested but no requests were added.")
612-
return // Nothing to send
615+
return result // Return result even if no requests
613616
}
614617

615618
await logger?.debug(
@@ -620,6 +623,7 @@ public actor Client {
620623
try await connection.send(data)
621624

622625
// Responses will be handled asynchronously by the message loop and handleBatchResponse/handleResponse.
626+
return result
623627
}
624628

625629
// MARK: - Lifecycle

‎Tests/MCPTests/ClientTests.swift‎

Lines changed: 54 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -345,12 +345,10 @@ struct ClientTests {
345345

346346
let request1 = Ping.request()
347347
let request2 = Ping.request()
348-
nonisolated(unsafe) var resultTask1: Task<Ping.Result, Swift.Error>?
349-
nonisolated(unsafe) var resultTask2: Task<Ping.Result, Swift.Error>?
350-
351-
try await client.withBatch { batch in
352-
resultTask1 = try await batch.addRequest(request1)
353-
resultTask2 = try await batch.addRequest(request2)
348+
let (resultTask1, resultTask2) = try await client.withBatch { batch in
349+
let task1 = try await batch.addRequest(request1)
350+
let task2 = try await batch.addRequest(request2)
351+
return (task1, task2)
354352
}
355353

356354
// Check if batch message was sent (after initialize and initialized notification)
@@ -381,13 +379,8 @@ struct ClientTests {
381379
try await transport.queue(batch: [anyResponse1, anyResponse2])
382380

383381
// Wait for results and verify
384-
guard let task1 = resultTask1, let task2 = resultTask2 else {
385-
#expect(Bool(false), "Result tasks not created")
386-
return
387-
}
388-
389-
_ = try await task1.value // Should succeed
390-
_ = try await task2.value // Should succeed
382+
_ = try await resultTask1.value // Should succeed
383+
_ = try await resultTask2.value // Should succeed
391384

392385
#expect(Bool(true)) // Reaching here means success
393386

@@ -426,11 +419,11 @@ struct ClientTests {
426419
let request1 = Ping.request() // Success
427420
let request2 = Ping.request() // Error
428421

429-
nonisolated(unsafe) var resultTasks: [Task<Ping.Result, Swift.Error>] = []
430-
431-
try await client.withBatch { batch in
432-
resultTasks.append(try await batch.addRequest(request1))
433-
resultTasks.append(try await batch.addRequest(request2))
422+
let resultTasks = try await client.withBatch { batch in
423+
[
424+
try await batch.addRequest(request1),
425+
try await batch.addRequest(request2),
426+
]
434427
}
435428

436429
// Check if batch message was sent (after initialize and initialized notification)
@@ -514,6 +507,49 @@ struct ClientTests {
514507
await client.disconnect()
515508
}
516509

510+
@Test("Batch request - empty with non-Void return")
511+
func testBatchRequestEmptyNonVoid() async throws {
512+
let transport = MockTransport()
513+
let client = Client(name: "TestClient", version: "1.0")
514+
515+
// Set up a task to handle the initialize response
516+
let initTask = Task {
517+
try await Task.sleep(for: .milliseconds(10))
518+
if let lastMessage = await transport.sentMessages.last,
519+
let data = lastMessage.data(using: .utf8),
520+
let request = try? JSONDecoder().decode(Request<Initialize>.self, from: data)
521+
{
522+
let response = Initialize.response(
523+
id: request.id,
524+
result: .init(
525+
protocolVersion: Version.latest,
526+
capabilities: .init(),
527+
serverInfo: .init(name: "TestServer", version: "1.0"),
528+
instructions: nil
529+
)
530+
)
531+
try await transport.queue(response: response)
532+
}
533+
}
534+
535+
try await client.connect(transport: transport)
536+
try await Task.sleep(for: .milliseconds(10))
537+
initTask.cancel()
538+
539+
// Call withBatch with non-Void return but don't add any requests
540+
let result: Int = try await client.withBatch { _ in
541+
42
542+
}
543+
544+
// Verify the closure's return value is passed through
545+
#expect(result == 42)
546+
547+
// Check that only initialize message and initialized notification were sent
548+
#expect(await transport.sentMessages.count == 2)
549+
550+
await client.disconnect()
551+
}
552+
517553
@Test("Notify method sends notifications")
518554
func testClientNotify() async throws {
519555
let transport = MockTransport()

0 commit comments

Comments
 (0)