Add support for JSON-RPC batching (#67)

* Update StdioTransport documentation comments to indicate support for JSON-RPC batches

* Add server support for JSON-RPC batching

* Add client support for JSON-RPC batching

* Fix bug that caused error messages to be wrapped unnecessarily

Tests/MCPTests/ClientTests.swift

@@ -226,4 +226,141 @@ struct ClientTests {
         // Disconnect client
         await client.disconnect()
     }
+
+    @Test("Batch request - success")
+    func testBatchRequestSuccess() async throws {
+        let transport = MockTransport()
+        let client = Client(name: "TestClient", version: "1.0")
+
+        try await client.connect(transport: transport)
+        try await Task.sleep(for: .milliseconds(10))  // Allow connection tasks
+
+        let request1 = Ping.request()
+        let request2 = Ping.request()
+        var resultTask1: Task<Ping.Result, Swift.Error>?
+        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)
+        }
+
+        // Check if one batch message was sent
+        let sentMessages = await transport.sentMessages
+        #expect(sentMessages.count == 1)
+
+        guard let batchData = sentMessages.first?.data(using: .utf8) else {
+            #expect(Bool(false), "Failed to get batch data")
+            return
+        }
+
+        // Verify the sent batch contains the two requests
+        let decoder = JSONDecoder()
+        let sentRequests = try decoder.decode([AnyRequest].self, from: batchData)
+        #expect(sentRequests.count == 2)
+        #expect(sentRequests.first?.id == request1.id)
+        #expect(sentRequests.first?.method == Ping.name)
+        #expect(sentRequests.last?.id == request2.id)
+        #expect(sentRequests.last?.method == Ping.name)
+
+        // Prepare batch response
+        let response1 = Response<Ping>(id: request1.id, result: .init())
+        let response2 = Response<Ping>(id: request2.id, result: .init())
+        let anyResponse1 = try AnyResponse(response1)
+        let anyResponse2 = try AnyResponse(response2)
+
+        // Queue the batch response
+        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
+
+        #expect(Bool(true))  // Reaching here means success
+
+        await client.disconnect()
+    }
+
+    @Test("Batch request - mixed success/error")
+    func testBatchRequestMixed() async throws {
+        let transport = MockTransport()
+        let client = Client(name: "TestClient", version: "1.0")
+
+        try await client.connect(transport: transport)
+        try await Task.sleep(for: .milliseconds(10))
+
+        let request1 = Ping.request()  // Success
+        let request2 = Ping.request()  // Error
+
+        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))
+        }
+
+        // Check if one batch message was sent
+        #expect(await transport.sentMessages.count == 1)
+
+        // Prepare batch response (success for 1, error for 2)
+        let response1 = Response<Ping>(id: request1.id, result: .init())
+        let error = MCPError.internalError("Simulated batch error")
+        let response2 = Response<Ping>(id: request2.id, error: error)
+        let anyResponse1 = try AnyResponse(response1)
+        let anyResponse2 = try AnyResponse(response2)
+
+        // Queue the batch response
+        try await transport.queue(batch: [anyResponse1, anyResponse2])
+
+        // Wait for results and verify
+        #expect(resultTasks.count == 2)
+        guard resultTasks.count == 2 else {
+            #expect(Bool(false), "Expected 2 result tasks")
+            return
+        }
+
+        let task1 = resultTasks[0]
+        let task2 = resultTasks[1]
+
+        _ = try await task1.value  // Task 1 should succeed
+
+        do {
+            _ = try await task2.value  // Task 2 should fail
+            #expect(Bool(false), "Task 2 should have thrown an error")
+        } catch let mcpError as MCPError {
+            if case .internalError(let message) = mcpError {
+                #expect(message == "Simulated batch error")
+            } else {
+                #expect(Bool(false), "Expected internalError, got \(mcpError)")
+            }
+        } catch {
+            #expect(Bool(false), "Expected MCPError, got \(error)")
+        }
+
+        await client.disconnect()
+    }
+
+    @Test("Batch request - empty")
+    func testBatchRequestEmpty() async throws {
+        let transport = MockTransport()
+        let client = Client(name: "TestClient", version: "1.0")
+
+        try await client.connect(transport: transport)
+        try await Task.sleep(for: .milliseconds(10))
+
+        // Call withBatch but don't add any requests
+        try await client.withBatch { _ in
+            // No requests added
+        }
+
+        // Check that no messages were sent
+        #expect(await transport.sentMessages.isEmpty)
+
+        await client.disconnect()
+    }
 }

Tests/MCPTests/Helpers/MockTransport.swift

@@ -80,8 +80,7 @@ actor MockTransport: Transport {
         shouldFailSend = shouldFail
     }
 
-    func queue<M: Method>(request: Request<M>) throws {
-        let data = try encoder.encode(request)
+    func queue(data: Data) {
         if let continuation = dataStreamContinuation {
             continuation.yield(data)
         } else {
@@ -89,14 +88,24 @@ actor MockTransport: Transport {
         }
     }
 
+    func queue<M: Method>(request: Request<M>) throws {
+        queue(data: try encoder.encode(request))
+    }
+
     func queue<M: Method>(response: Response<M>) throws {
-        let data = try encoder.encode(response)
-        dataToReceive.append(data)
+        queue(data: try encoder.encode(response))
     }
 
     func queue<N: Notification>(notification: Message<N>) throws {
-        let data = try encoder.encode(notification)
-        dataToReceive.append(data)
+        queue(data: try encoder.encode(notification))
+    }
+
+    func queue(batch requests: [AnyRequest]) throws {
+        queue(data: try encoder.encode(requests))
+    }
+
+    func queue(batch responses: [AnyResponse]) throws {
+        queue(data: try encoder.encode(responses))
     }
 
     func decodeLastSentMessage<T: Decodable>() -> T? {

Tests/MCPTests/ResponseTests.swift

@@ -78,7 +78,7 @@ struct ResponseTests {
             #expect(decodedError.code == -32700)
             #expect(
                 decodedError.localizedDescription
-                    == "Parse error: Invalid JSON: Parse error: Invalid JSON: Invalid syntax")
+                    == "Parse error: Invalid JSON: Invalid syntax")
         } else {
             #expect(Bool(false), "Expected error result")
         }

Tests/MCPTests/ServerTests.swift

@@ -144,4 +144,62 @@ struct ServerTests {
         await server.stop()
         await transport.disconnect()
     }
+
+    @Test("JSON-RPC batch processing")
+    func testJSONRPCBatchProcessing() async throws {
+        let transport = MockTransport()
+        let server = Server(name: "TestServer", version: "1.0")
+
+        // Start the server
+        try await server.start(transport: transport)
+
+        // Initialize the server first
+        try await transport.queue(
+            request: Initialize.request(
+                .init(
+                    protocolVersion: Version.latest,
+                    capabilities: .init(),
+                    clientInfo: .init(name: "TestClient", version: "1.0")
+                )
+            )
+        )
+
+        // Wait for server to initialize and respond
+        try await Task.sleep(for: .milliseconds(100))
+
+        // Clear sent messages
+        await transport.clearMessages()
+
+        // Create a batch with multiple requests
+        let batchJSON = """
+            [
+                {"jsonrpc":"2.0","id":1,"method":"ping","params":{}},
+                {"jsonrpc":"2.0","id":2,"method":"ping","params":{}}
+            ]
+            """
+        let batch = try JSONDecoder().decode([AnyRequest].self, from: batchJSON.data(using: .utf8)!)
+
+        // Send the batch request
+        try await transport.queue(batch: batch)
+
+        // Wait for batch processing
+        try await Task.sleep(for: .milliseconds(100))
+
+        // Verify response
+        let sentMessages = await transport.sentMessages
+        #expect(sentMessages.count == 1)
+
+        if let batchResponse = sentMessages.first {
+            // Should be an array
+            #expect(batchResponse.hasPrefix("["))
+            #expect(batchResponse.hasSuffix("]"))
+
+            // Should contain both request IDs
+            #expect(batchResponse.contains("\"id\":1"))
+            #expect(batchResponse.contains("\"id\":2"))
+        }
+
+        await server.stop()
+        await transport.disconnect()
+    }
 }