Skip to content

Commit 70ea8de

Browse files
doozMenStijn Willems
doozMen
authored and
Stijn Willems
committed
fix continuation (#9)
1 parent 5d50718 commit 70ea8de

4 files changed

Lines changed: 187 additions & 24 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: 122 additions & 1 deletion
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 {
@@ -525,6 +532,86 @@ import Logging
525532
completion: .contentProcessed { [weak self] error in
526533
guard let self = self else { return }
527534

535+
if let error = error {
536+
self.logger.error("Send error: \(error)")
537+
538+
// Schedule reconnection check on a separate task
539+
Task { [weak self] in
540+
guard let self = self else { return }
541+
let isStopping = await self.isStopping
542+
if !isStopping && self.reconnectionConfig.enabled {
543+
let isConnected = await self.isConnected
544+
if isConnected && error.isConnectionLost {
545+
self.logger.warning(
546+
"Connection appears broken, will attempt to reconnect..."
547+
)
548+
await self.setIsConnected(false)
549+
try? await Task.sleep(for: .milliseconds(500))
550+
551+
let currentIsStopping = await self.isStopping
552+
if !currentIsStopping {
553+
self.connection.cancel()
554+
try? await self.connect()
555+
Task { @MainActor in
556+
if !sendContinuationResumed {
557+
sendContinuationResumed = true
558+
if let error = error {
559+
self.logger.error("Send error: \(error)")
560+
561+
// Check if we should attempt to reconnect on send failure
562+
let isStopping = await self.isStopping // Await actor-isolated property
563+
if !isStopping && self.reconnectionConfig.enabled {
564+
let isConnected = await self.isConnected
565+
if isConnected {
566+
if error.isConnectionLost {
567+
self.logger.warning(
568+
"Connection appears broken, will attempt to reconnect..."
569+
)
570+
571+
// Schedule connection restart
572+
Task { [weak self] in // Operate on self's executor
573+
guard let self = self else { return }
574+
575+
await self.setIsConnected(false)
576+
577+
try? await Task.sleep(for: .milliseconds(500))
578+
579+
let currentIsStopping = await self.isStopping
580+
if !currentIsStopping {
581+
// Cancel the connection, then attempt to reconnect fully.
582+
self.connection.cancel()
583+
try? await self.connect()
584+
}
585+
}
586+
}
587+
if let error = error {
588+
self.logger.error("Send error: \(error)")
589+
590+
// Schedule reconnection check on a separate task
591+
Task { [weak self] in
592+
guard let self = self else { return }
593+
let isStopping = await self.isStopping
594+
if !isStopping && self.reconnectionConfig.enabled {
595+
let isConnected = await self.isConnected
596+
if isConnected && error.isConnectionLost {
597+
self.logger.warning(
598+
"Connection appears broken, will attempt to reconnect..."
599+
)
600+
await self.setIsConnected(false)
601+
try? await Task.sleep(for: .milliseconds(500))
602+
603+
let currentIsStopping = await self.isStopping
604+
if !currentIsStopping {
605+
self.connection.cancel()
606+
try? await self.connect()
607+
completion: .contentProcessed { [weak self, continuation] error in
608+
guard let self = self else {
609+
continuation.resume(
610+
throwing: MCPError.internalError(
611+
"Transport deallocated during send"))
612+
return
613+
}
614+
528615
if let error = error {
529616
self.logger.error("Send error: \(error)")
530617

@@ -751,6 +838,40 @@ import Logging
751838
continuation.resume(throwing: MCPError.connectionClosed)
752839
} else {
753840
continuation.resume(returning: Data())
841+
Task { @MainActor in
842+
if !receiveContinuationResumed {
843+
receiveContinuationResumed = true
844+
if let error = error {
845+
continuation.resume(throwing: MCPError.transportError(error))
846+
} else if let content = content {
847+
continuation.resume(returning: content)
848+
} else if isComplete {
849+
self.logger.trace("Connection completed by peer")
850+
continuation.resume(throwing: MCPError.connectionClosed)
851+
} else {
852+
// EOF: Resume with empty data instead of throwing an error
853+
continuation.resume(returning: Data())
854+
}
855+
}
856+
if let error = error {
857+
continuation.resume(throwing: MCPError.transportError(error))
858+
} else if let content = content {
859+
continuation.resume(returning: content)
860+
} else if isComplete {
861+
self.logger.trace("Connection completed by peer")
862+
continuation.resume(throwing: MCPError.connectionClosed)
863+
} else {
864+
continuation.resume(returning: Data())
865+
[weak self, continuation] content, _, isComplete, error in
866+
if let error = error {
867+
continuation.resume(throwing: MCPError.transportError(error))
868+
} else if let content = content {
869+
continuation.resume(returning: content)
870+
} else if isComplete {
871+
self?.logger.trace("Connection completed by peer")
872+
continuation.resume(throwing: MCPError.connectionClosed)
873+
} else {
874+
continuation.resume(returning: Data())
754875
}
755876
}
756877
}

‎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)