fix continuation (#9)

Co-authored-by: Stijn Willems <stijn@promptping.ai>
fix continuation

fix rebase with upstream

fix rebase with upstream

fix continuation (#9)

Co-authored-by: Stijn Willems <stijn@promptping.ai>

revert: undo fix continuation (#9) — improper platform bump and unnecessary guards

proper fix

fix: ensure CheckedContinuation always resumes in NetworkTransport send/receive paths

- Capture continuation directly in completion closures alongside [weak self]
- Resume with error when actor deallocates during async operation
- Add test for withBatch non-Void early-return path
- Prevents permanent deadlock from unresumed continuations

Author: doozMen

.claude/settings.local.json

@@ -0,0 +1,15 @@
+{
+  "permissions": {
+    "allow": [
+      "Bash(jj log:*)",
+      "Bash(jj diff:*)",
+      "Bash(jj bookmark:*)",
+      "Bash(jj show:*)",
+      "Bash(jj git:*)",
+      "Bash(gh repo:*)",
+      "Bash(gh pr:*)",
+      "Bash(jj status:*)",
+      "Bash(gh api:*)"
+    ]
+  }
+}

.swift-version

@@ -0,0 +1 @@
+6.3-snapshot-2026-03-05

Sources/MCP/Base/Transports/NetworkTransport.swift

@@ -511,9 +511,6 @@ import Logging
             var messageWithNewline = message
             messageWithNewline.append(UInt8(ascii: "\n"))
 
-            // Use a local actor-isolated variable to track continuation state
-            var sendContinuationResumed = false
-
             try await withCheckedThrowingContinuation {
                 [weak self] (continuation: CheckedContinuation<Void, Swift.Error>) in
                 guard let self = self else {
@@ -525,50 +522,43 @@ import Logging
                     content: messageWithNewline,
                     contentContext: .defaultMessage,
                     isComplete: true,
-                    completion: .contentProcessed { [weak self] error in
-                        guard let self = self else { return }
-
-                        Task { @MainActor in
-                            if !sendContinuationResumed {
-                                sendContinuationResumed = true
-                                if let error = error {
-                                    self.logger.error("Send error: \(error)")
-
-                                    // Check if we should attempt to reconnect on send failure
-                                    let isStopping = await self.isStopping  // Await actor-isolated property
-                                    if !isStopping && self.reconnectionConfig.enabled {
-                                        let isConnected = await self.isConnected
-                                        if isConnected {
-                                            if error.isConnectionLost {
-                                                self.logger.warning(
-                                                    "Connection appears broken, will attempt to reconnect..."
-                                                )
-
-                                                // Schedule connection restart
-                                                Task { [weak self] in  // Operate on self's executor
-                                                    guard let self = self else { return }
-
-                                                    await self.setIsConnected(false)
-
-                                                    try? await Task.sleep(for: .milliseconds(500))
-
-                                                    let currentIsStopping = await self.isStopping
-                                                    if !currentIsStopping {
-                                                        // Cancel the connection, then attempt to reconnect fully.
-                                                        self.connection.cancel()
-                                                        try? await self.connect()
-                                                    }
-                                                }
-                                            }
+                    completion: .contentProcessed { [weak self, continuation] error in
+                        guard let self = self else {
+                            continuation.resume(
+                                throwing: MCPError.internalError(
+                                    "Transport deallocated during send"))
+                            return
+                        }
+
+                        if let error = error {
+                            self.logger.error("Send error: \(error)")
+
+                            // Schedule reconnection check on a separate task
+                            Task { [weak self] in
+                                guard let self = self else { return }
+                                let isStopping = await self.isStopping
+                                if !isStopping && self.reconnectionConfig.enabled {
+                                    let isConnected = await self.isConnected
+                                    if isConnected && error.isConnectionLost {
+                                        self.logger.warning(
+                                            "Connection appears broken, will attempt to reconnect..."
+                                        )
+                                        await self.setIsConnected(false)
+                                        try? await Task.sleep(for: .milliseconds(500))
+
+                                        let currentIsStopping = await self.isStopping
+                                        if !currentIsStopping {
+                                            self.connection.cancel()
+                                            try? await self.connect()
                                         }
                                     }
-
-                                    continuation.resume(
-                                        throwing: MCPError.internalError("Send error: \(error)"))
-                                } else {
-                                    continuation.resume()
                                 }
                             }
+
+                            continuation.resume(
+                                throwing: MCPError.internalError("Send error: \(error)"))
+                        } else {
+                            continuation.resume()
                         }
                     })
             }
@@ -747,8 +737,6 @@ import Logging
         /// - Returns: The received data chunk
         /// - Throws: Network errors or transport failures
         private func receiveData() async throws -> Data {
-            var receiveContinuationResumed = false
-
             return try await withCheckedThrowingContinuation {
                 [weak self] (continuation: CheckedContinuation<Data, Swift.Error>) in
                 guard let self = self else {
@@ -758,22 +746,16 @@ import Logging
 
                 let maxLength = bufferConfig.maxReceiveBufferSize ?? Int.max
                 connection.receive(minimumIncompleteLength: 1, maximumLength: maxLength) {
-                    content, _, isComplete, error in
-                    Task { @MainActor in
-                        if !receiveContinuationResumed {
-                            receiveContinuationResumed = true
-                            if let error = error {
-                                continuation.resume(throwing: MCPError.transportError(error))
-                            } else if let content = content {
-                                continuation.resume(returning: content)
-                            } else if isComplete {
-                                self.logger.trace("Connection completed by peer")
-                                continuation.resume(throwing: MCPError.connectionClosed)
-                            } else {
-                                // EOF: Resume with empty data instead of throwing an error
-                                continuation.resume(returning: Data())
-                            }
-                        }
+                    [weak self, continuation] content, _, isComplete, error in
+                    if let error = error {
+                        continuation.resume(throwing: MCPError.transportError(error))
+                    } else if let content = content {
+                        continuation.resume(returning: content)
+                    } else if isComplete {
+                        self?.logger.trace("Connection completed by peer")
+                        continuation.resume(throwing: MCPError.connectionClosed)
+                    } else {
+                        continuation.resume(returning: Data())
                     }
                 }
             }

Sources/MCP/Client/Client.swift

@@ -592,24 +592,27 @@ public actor Client {
     ///                   Use this object to add requests to the batch.
     /// - Throws: `MCPError.internalError` if the client is not connected.
     ///           Can also rethrow errors from the `body` closure or from sending the batch request.
-    public func withBatch(body: @escaping @Sendable (Batch) async throws -> Void) async throws {
+    @discardableResult
+    public func withBatch<T: Sendable>(
+        body: @escaping @Sendable (Batch) async throws -> T
+    ) async throws -> T {
         guard let connection = connection else {
             throw MCPError.internalError("Client connection not initialized")
         }
 
         // Create Batch actor, passing self (Client)
         let batch = Batch(client: self)
 
-        // Populate the batch actor by calling the user's closure.
-        try await body(batch)
+        // Populate the batch actor by calling the user's closure and capture result.
+        let result = try await body(batch)
 
         // Get the collected requests from the batch actor
         let requests = await batch.requests
 
         // Check if there are any requests to send
         guard !requests.isEmpty else {
             await logger?.debug("Batch requested but no requests were added.")
-            return  // Nothing to send
+            return result  // Return result even if no requests
         }
 
         await logger?.debug(
@@ -620,6 +623,7 @@ public actor Client {
         try await connection.send(data)
 
         // Responses will be handled asynchronously by the message loop and handleBatchResponse/handleResponse.
+        return result
     }
 
     // MARK: - Lifecycle

Tests/MCPTests/ClientTests.swift

@@ -345,12 +345,10 @@ struct ClientTests {
 
         let request1 = Ping.request()
         let request2 = Ping.request()
-        nonisolated(unsafe) var resultTask1: Task<Ping.Result, Swift.Error>?
-        nonisolated(unsafe) var resultTask2: Task<Ping.Result, Swift.Error>?
-
-        try await client.withBatch { batch in
-            resultTask1 = try await batch.addRequest(request1)
-            resultTask2 = try await batch.addRequest(request2)
+        let (resultTask1, resultTask2) = try await client.withBatch { batch in
+            let task1 = try await batch.addRequest(request1)
+            let task2 = try await batch.addRequest(request2)
+            return (task1, task2)
         }
 
         // Check if batch message was sent (after initialize and initialized notification)
@@ -381,13 +379,8 @@ struct ClientTests {
         try await transport.queue(batch: [anyResponse1, anyResponse2])
 
         // Wait for results and verify
-        guard let task1 = resultTask1, let task2 = resultTask2 else {
-            #expect(Bool(false), "Result tasks not created")
-            return
-        }
-
-        _ = try await task1.value  // Should succeed
-        _ = try await task2.value  // Should succeed
+        _ = try await resultTask1.value  // Should succeed
+        _ = try await resultTask2.value  // Should succeed
 
         #expect(Bool(true))  // Reaching here means success
 
@@ -426,11 +419,11 @@ struct ClientTests {
         let request1 = Ping.request()  // Success
         let request2 = Ping.request()  // Error
 
-        nonisolated(unsafe) var resultTasks: [Task<Ping.Result, Swift.Error>] = []
-
-        try await client.withBatch { batch in
-            resultTasks.append(try await batch.addRequest(request1))
-            resultTasks.append(try await batch.addRequest(request2))
+        let resultTasks = try await client.withBatch { batch in
+            [
+                try await batch.addRequest(request1),
+                try await batch.addRequest(request2),
+            ]
         }
 
         // Check if batch message was sent (after initialize and initialized notification)
@@ -514,6 +507,49 @@ struct ClientTests {
         await client.disconnect()
     }
 
+    @Test("Batch request - empty with non-Void return")
+    func testBatchRequestEmptyNonVoid() async throws {
+        let transport = MockTransport()
+        let client = Client(name: "TestClient", version: "1.0")
+
+        // Set up a task to handle the initialize response
+        let initTask = Task {
+            try await Task.sleep(for: .milliseconds(10))
+            if let lastMessage = await transport.sentMessages.last,
+                let data = lastMessage.data(using: .utf8),
+                let request = try? JSONDecoder().decode(Request<Initialize>.self, from: data)
+            {
+                let response = Initialize.response(
+                    id: request.id,
+                    result: .init(
+                        protocolVersion: Version.latest,
+                        capabilities: .init(),
+                        serverInfo: .init(name: "TestServer", version: "1.0"),
+                        instructions: nil
+                    )
+                )
+                try await transport.queue(response: response)
+            }
+        }
+
+        try await client.connect(transport: transport)
+        try await Task.sleep(for: .milliseconds(10))
+        initTask.cancel()
+
+        // Call withBatch with non-Void return but don't add any requests
+        let result: Int = try await client.withBatch { _ in
+            42
+        }
+
+        // Verify the closure's return value is passed through
+        #expect(result == 42)
+
+        // Check that only initialize message and initialized notification were sent
+        #expect(await transport.sentMessages.count == 2)
+
+        await client.disconnect()
+    }
+
     @Test("Notify method sends notifications")
     func testClientNotify() async throws {
         let transport = MockTransport()