From 0baa2a875d5c96da4a0b582432f57cb307078bf5 Mon Sep 17 00:00:00 2001 From: bbdu3 Date: Mon, 9 Nov 2026 12:54:03 +0800 Subject: [PATCH 1/2] Complete batch request tasks when withBatch fails Signed-off-by: bbdu3 --- Sources/MCP/Client/Client.swift | 93 +++++++++++++--- Tests/MCPTests/ClientTests.swift | 182 +++++++++++++++++++++++++++++++ 2 files changed, 259 insertions(+), 16 deletions(-) diff --git a/Sources/MCP/Client/Client.swift b/Sources/MCP/Client/Client.swift index bf07d092..bc32122e 100644 --- a/Sources/MCP/Client/Client.swift +++ b/Sources/MCP/Client/Client.swift @@ -2,6 +2,7 @@ import Logging import struct Foundation.Data import struct Foundation.Date +import struct Foundation.UUID import class Foundation.JSONDecoder import class Foundation.JSONEncoder @@ -467,12 +468,64 @@ public actor Client { // MARK: - Batching + private struct BatchState { + var registeredIDs: Set = [] + // Non-nil once the body/send has finished. Retain failures until all + // asynchronously created continuations have had a chance to register. + var remainingIDs: Set? + var error: Swift.Error? + } + + private var batchStates: [UUID: BatchState] = [:] + + func addPendingBatchRequest( + id: ID, + continuation: CheckedContinuation, + type: T.Type, + batchID: UUID + ) { + var state = batchStates[batchID] ?? BatchState() + state.registeredIDs.insert(id) + state.remainingIDs?.remove(id) + if state.remainingIDs?.isEmpty == true { + batchStates.removeValue(forKey: batchID) + } else { + batchStates[batchID] = state + } + + if let error = state.error { + continuation.resume(throwing: error) + } else { + addPendingRequest(id: id, continuation: continuation, type: type) + } + } + + func finishBatch(id: UUID, requestIDs: Set, error: Swift.Error? = nil) { + var state = batchStates[id] ?? BatchState() + state.error = error + state.remainingIDs = requestIDs.subtracting(state.registeredIDs) + if state.remainingIDs?.isEmpty == true { + batchStates.removeValue(forKey: id) + } else { + batchStates[id] = state + } + + if let error { + for requestID in requestIDs { + // A response, cancellation, or disconnect may have already + // removed the request and resumed its continuation. + removePendingRequest(id: requestID)?.resume(throwing: error) + } + } + } + /// A batch of requests. /// /// Objects of this type are passed as an argument to the closure /// of the ``Client/withBatch(body:)`` method. public actor Batch { unowned let client: Client + let id = UUID() var requests: [AnyRequest] = [] init(client: Client) { @@ -494,10 +547,11 @@ public actor Client { // We are already inside a Task, but need another Task // to bridge to the client actor's context. Task { - await client.addPendingRequest( + await client.addPendingBatchRequest( id: request.id, continuation: continuation, - type: M.Result.self + type: M.Result.self, + batchID: id ) } } @@ -587,24 +641,31 @@ public actor Client { // 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) + do { + // Populate the batch actor by calling the user's closure. + try await body(batch) - // Get the collected requests from the batch actor - let requests = await batch.requests + // 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 - } + // 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 + } - await logger?.debug( - "Sending batch request", metadata: ["count": "\(requests.count)"]) + await logger?.debug( + "Sending batch request", metadata: ["count": "\(requests.count)"]) - // Encode the array of AnyMethod requests into a single JSON payload - let data = try encoder.encode(requests) - try await connection.send(data) + // Encode the array of AnyMethod requests into a single JSON payload + let data = try encoder.encode(requests) + try await connection.send(data) + finishBatch(id: batch.id, requestIDs: Set(requests.map(\.id))) + } catch { + let requests = await batch.requests + finishBatch(id: batch.id, requestIDs: Set(requests.map(\.id)), error: error) + throw error + } // Responses will be handled asynchronously by the message loop and handleBatchResponse/handleResponse. } diff --git a/Tests/MCPTests/ClientTests.swift b/Tests/MCPTests/ClientTests.swift index 8a2882a4..72acdd48 100644 --- a/Tests/MCPTests/ClientTests.swift +++ b/Tests/MCPTests/ClientTests.swift @@ -1,3 +1,4 @@ +import Logging import Foundation import Testing @@ -474,6 +475,135 @@ struct ClientTests { await client.disconnect() } + @Test("Batch failure completes all result tasks", arguments: [false, true]) + func testBatchFailureCompletesTasks(bodyThrows: Bool) async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + let server = Server(name: "TestServer", version: "1.0") + try await server.start(transport: serverTransport) + let transport = BatchFailureTransport(base: clientTransport) + let client = Client(name: "TestClient", version: "1.0") + try await client.connect(transport: transport) + let results = BatchFailureResults() + let expectedError: BatchFailureError = bodyThrows ? .body : .send + + do { + try await client.withBatch { batch in + await results.observe(try await batch.addRequest(Ping.request())) + await results.observe(try await batch.addRequest(Ping.request())) + if bodyThrows { throw BatchFailureError.body } + } + Issue.record("Expected the batch to fail") + } catch { + #expect(error as? BatchFailureError == expectedError) + } + + // Observe failure before disconnecting, since disconnect itself resumes + // pending requests and would hide missing batch cleanup. + let errors = await results.waitForErrors(count: 2) + #expect(errors.count == 2) + #expect(errors.allSatisfy { $0 == expectedError }) + #expect(await transport.batchSendCount == (bodyThrows ? 0 : 1)) + + // A failed batch must not close the connection or break later requests. + try await client.ping() + await client.disconnect() + await server.stop() + } + + @Test("Batch failure handles either registration order", arguments: [false, true]) + func testBatchFailureRegistrationOrder(registerFirst: Bool) async throws { + let client = Client(name: "TestClient", version: "1.0") + let batchID = UUID() + let requestID = ID.random + let lateID = ID.random + let otherBatchID = UUID() + let otherID = ID.random + let results = BatchFailureResults() + let otherResults = BatchFailureResults() + let (registered, registration) = AsyncStream.makeStream() + + func register(id: ID, batch: UUID) -> Task { + Task { + try await withCheckedThrowingContinuation { continuation in + Task { + await client.addPendingBatchRequest( + id: id, continuation: continuation, + type: Ping.Result.self, batchID: batch + ) + registration.yield(()) + } + } + } + } + + var iterator = registered.makeAsyncIterator() + await otherResults.observe(register(id: otherID, batch: otherBatchID)) + _ = await iterator.next() + if registerFirst { + await results.observe(register(id: requestID, batch: batchID)) + _ = await iterator.next() + } + + await client.finishBatch( + id: batchID, requestIDs: [requestID, lateID], error: BatchFailureError.send + ) + if !registerFirst { + await results.observe(register(id: requestID, batch: batchID)) + _ = await iterator.next() + } + let errors = await results.waitForErrors(count: 1) + #expect(errors == [.send]) + // The failure must remain available for every late registration, not + // just the first continuation that arrives after finishBatch. + await results.observe(register(id: lateID, batch: batchID)) + _ = await iterator.next() + let lateErrors = await results.waitForErrors(count: 2) + #expect(lateErrors == [.send, .send]) + #expect(await otherResults.errors.isEmpty) + + // Finishing another batch must deliver its own error, not the first + // batch's failure. Disconnect also checks that no continuation is resumed twice. + await client.finishBatch(id: otherBatchID, requestIDs: [otherID], error: BatchFailureError.body) + let otherErrors = await otherResults.waitForErrors(count: 1) + #expect(otherErrors == [.body]) + registration.finish() + await client.disconnect() + } + + @Test("Batch failure preserves a response that already completed") + func testBatchFailureAfterResponse() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + let server = Server(name: "TestServer", version: "1.0") + try await server.start(transport: serverTransport) + let client = Client(name: "TestClient", version: "1.0") + try await client.connect(transport: clientTransport) + let batchID = UUID() + let requestID = ID.random + let (registered, registration) = AsyncStream.makeStream() + let result = Task { + try await withCheckedThrowingContinuation { continuation in + Task { + await client.addPendingBatchRequest( + id: requestID, continuation: continuation, + type: Ping.Result.self, batchID: batchID + ) + registration.yield(()) + registration.finish() + } + } + } + for await _ in registered {} + try await serverTransport.send(JSONEncoder().encode(Ping.response(id: requestID, result: .init()))) + _ = try await result.value + + // The response has removed the pending request. Failure cleanup must + // not resume the same continuation again. + await client.finishBatch(id: batchID, requestIDs: [requestID], error: BatchFailureError.send) + try await client.ping() + await client.disconnect() + await server.stop() + } + @Test("Batch request - empty") func testBatchRequestEmpty() async throws { let transport = MockTransport() @@ -773,3 +903,55 @@ struct ClientTests { await client.disconnect() } } + +private enum BatchFailureError: Swift.Error, Equatable { + case send + case body +} + +private actor BatchFailureResults { + private(set) var errors: [BatchFailureError?] = [] + + func observe(_ task: Task) { + Task { + switch await task.result { + case .success: + errors.append(nil) + case .failure(let error): + errors.append(error as? BatchFailureError) + } + } + } + + func waitForErrors(count: Int) async -> [BatchFailureError?] { + let deadline = ContinuousClock.now + .seconds(3) + while errors.count < count && ContinuousClock.now < deadline { + try? await Task.sleep(for: .milliseconds(10)) + } + return errors + } +} + +private actor BatchFailureTransport: Transport { + nonisolated let logger = Logger(label: "mcp.test.batch-failure") + let base: InMemoryTransport + private var stream = AsyncThrowingStream { $0.finish() } + private(set) var batchSendCount = 0 + + init(base: InMemoryTransport) { self.base = base } + + func connect() async throws { + try await base.connect() + stream = await base.receive() + } + func disconnect() async { await base.disconnect() } + func receive() -> AsyncThrowingStream { stream } + + func send(_ data: Data) async throws { + if (try? JSONDecoder().decode([AnyRequest].self, from: data)) != nil { + batchSendCount += 1 + throw BatchFailureError.send + } + try await base.send(data) + } +} From e8b76789477ed41c8b08c120741e2e7194b75921 Mon Sep 17 00:00:00 2001 From: bbdu3 Date: Tue, 10 Nov 2026 09:51:45 +0800 Subject: [PATCH 2/2] fix document the failure behavior Signed-off-by: bbdu3 --- README.md | 6 +++ Sources/MCP/Client/Client.swift | 35 +++++++++++---- Tests/MCPTests/ClientTests.swift | 74 +++++++++++++++++++++++++++++--- 3 files changed, 99 insertions(+), 16 deletions(-) diff --git a/README.md b/README.md index 5bf1aaaa..084c8421 100644 --- a/README.md +++ b/README.md @@ -598,6 +598,12 @@ do { #### Request Batching +If `withBatch` throws, any tasks returned by `batch.addRequest(_:)` that have not already +completed fail with the same error. Catch the `withBatch` error to continue awaiting those tasks. + +`addRequest` throws `MCPError.invalidRequest` when an ID is repeated within the same batch. +Callers must still keep request IDs unique across the session, including outside that batch. + Improve performance by sending multiple requests in a single batch: ```swift diff --git a/Sources/MCP/Client/Client.swift b/Sources/MCP/Client/Client.swift index bc32122e..b51d1f32 100644 --- a/Sources/MCP/Client/Client.swift +++ b/Sources/MCP/Client/Client.swift @@ -478,28 +478,29 @@ public actor Client { private var batchStates: [UUID: BatchState] = [:] + // Internal so tests can deterministically exercise registration before and after failure. func addPendingBatchRequest( id: ID, continuation: CheckedContinuation, type: T.Type, batchID: UUID ) { - var state = batchStates[batchID] ?? BatchState() - state.registeredIDs.insert(id) - state.remainingIDs?.remove(id) - if state.remainingIDs?.isEmpty == true { + // Mutate in place to avoid copying the ID sets on every registration. + batchStates[batchID, default: BatchState()].registeredIDs.insert(id) + batchStates[batchID]?.remainingIDs?.remove(id) + let error = batchStates[batchID]?.error + if batchStates[batchID]?.remainingIDs?.isEmpty == true { batchStates.removeValue(forKey: batchID) - } else { - batchStates[batchID] = state } - if let error = state.error { + if let error { continuation.resume(throwing: error) } else { addPendingRequest(id: id, continuation: continuation, type: type) } } + // Internal so tests can finish a batch before its continuations register. func finishBatch(id: UUID, requestIDs: Set, error: Swift.Error? = nil) { var state = batchStates[id] ?? BatchState() state.error = error @@ -527,6 +528,7 @@ public actor Client { unowned let client: Client let id = UUID() var requests: [AnyRequest] = [] + private var requestIDs: Set = [] init(client: Client) { self.client = client @@ -534,14 +536,26 @@ public actor Client { /// Adds a request to the batch and prepares its expected response task. /// The actual sending happens when the `withBatch` scope completes. + /// If `withBatch` throws before this task completes, the task fails with the same error. + /// + /// - Parameter request: A request whose ID has not already been added to this batch. /// - Returns: A `Task` that will eventually produce the result or throw an error. + /// - Throws: `MCPError.invalidRequest` if the ID is already present in this batch, + /// or an error if the request cannot be encoded. + /// - Important: Request IDs must be unique across the session. This method only + /// checks for duplicates within this batch. public func addRequest(_ request: Request) async throws -> Task< M.Result, Swift.Error > { - requests.append(try AnyRequest(request)) + guard !requestIDs.contains(request.id) else { + throw MCPError.invalidRequest("Duplicate request ID in batch: \(request.id)") + } + let anyRequest = try AnyRequest(request) + requestIDs.insert(request.id) + requests.append(anyRequest) // Return a Task that registers the pending request and awaits its result. - // The continuation is resumed when the response arrives. + // The continuation is resumed when a response arrives or the batch fails. return Task { try await withCheckedThrowingContinuation { continuation in // We are already inside a Task, but need another Task @@ -571,6 +585,9 @@ public actor Client { /// returns a `Task` handle representing the asynchronous operation /// for that specific request's result. /// + /// If this method throws, any tasks returned by `batch.addRequest(_:)` that have not + /// already completed fail with the same error. Tasks that already completed keep their results. + /// /// It's recommended to collect these `Task` handles into an array /// within the `body` closure`. After the `withBatch` method returns /// (meaning the batch request has been sent), you can then process diff --git a/Tests/MCPTests/ClientTests.swift b/Tests/MCPTests/ClientTests.swift index 72acdd48..ff7a1ad6 100644 --- a/Tests/MCPTests/ClientTests.swift +++ b/Tests/MCPTests/ClientTests.swift @@ -483,7 +483,7 @@ struct ClientTests { let transport = BatchFailureTransport(base: clientTransport) let client = Client(name: "TestClient", version: "1.0") try await client.connect(transport: transport) - let results = BatchFailureResults() + let results = BatchFailureResults() let expectedError: BatchFailureError = bodyThrows ? .body : .send do { @@ -510,6 +510,66 @@ struct ClientTests { await server.stop() } + @Test("Batch rejects duplicate request IDs", arguments: [ID.string("duplicate"), ID.number(42)]) + func testBatchRejectsDuplicateRequestIDs(id: ID) async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + let server = Server(name: "TestServer", version: "1.0") + try await server.start(transport: serverTransport) + let transport = BatchFailureTransport(base: clientTransport) + let client = Client(name: "TestClient", version: "1.0") + try await client.connect(transport: transport) + let results = BatchFailureResults() + let expectedError = MCPError.invalidRequest("Duplicate request ID in batch: \(id)") + + await #expect(throws: expectedError) { + try await client.withBatch { batch in + await results.observe(try await batch.addRequest(Ping.request(id: id))) + _ = try await batch.addRequest(Ping.request(id: id)) + } + } + + // Reject the duplicate before sending, and fail the original task with + // the error that escaped the body rather than leaving it suspended. + let errors = await results.waitForErrors(count: 1) + #expect(errors == [expectedError]) + #expect(await transport.batchSendCount == 0) + try await client.ping() + await client.disconnect() + await server.stop() + } + + @Test("Catching a duplicate ID error preserves the batch") + func testBatchCaughtDuplicateID() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + let server = Server(name: "TestServer", version: "1.0") + try await server.start(transport: serverTransport) + let transport = BatchFailureTransport(base: clientTransport) + let client = Client(name: "TestClient", version: "1.0") + try await client.connect(transport: transport) + let results = BatchFailureResults() + + do { + try await client.withBatch { batch in + let id = ID.random + await results.observe(try await batch.addRequest(Ping.request(id: id))) + await #expect(throws: MCPError.invalidRequest("Duplicate request ID in batch: \(id)")) { + _ = try await batch.addRequest(Ping.request(id: id)) + } + await results.observe(try await batch.addRequest(Ping.request())) + #expect(await batch.requests.count == 2) + } + Issue.record("Expected the batch send to fail") + } catch { + #expect(error as? BatchFailureError == .send) + } + + let errors = await results.waitForErrors(count: 2) + #expect(errors == [.send, .send]) + #expect(await transport.batchSendCount == 1) + await client.disconnect() + await server.stop() + } + @Test("Batch failure handles either registration order", arguments: [false, true]) func testBatchFailureRegistrationOrder(registerFirst: Bool) async throws { let client = Client(name: "TestClient", version: "1.0") @@ -518,8 +578,8 @@ struct ClientTests { let lateID = ID.random let otherBatchID = UUID() let otherID = ID.random - let results = BatchFailureResults() - let otherResults = BatchFailureResults() + let results = BatchFailureResults() + let otherResults = BatchFailureResults() let (registered, registration) = AsyncStream.makeStream() func register(id: ID, batch: UUID) -> Task { @@ -909,8 +969,8 @@ private enum BatchFailureError: Swift.Error, Equatable { case body } -private actor BatchFailureResults { - private(set) var errors: [BatchFailureError?] = [] +private actor BatchFailureResults { + private(set) var errors: [Failure?] = [] func observe(_ task: Task) { Task { @@ -918,12 +978,12 @@ private actor BatchFailureResults { case .success: errors.append(nil) case .failure(let error): - errors.append(error as? BatchFailureError) + errors.append(error as? Failure) } } } - func waitForErrors(count: Int) async -> [BatchFailureError?] { + func waitForErrors(count: Int) async -> [Failure?] { let deadline = ContinuousClock.now + .seconds(3) while errors.count < count && ContinuousClock.now < deadline { try? await Task.sleep(for: .milliseconds(10))