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 bf07d092..b51d1f32 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,13 +468,67 @@ 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] = [:] + + // Internal so tests can deterministically exercise registration before and after failure. + func addPendingBatchRequest( + id: ID, + continuation: CheckedContinuation, + type: T.Type, + batchID: UUID + ) { + // 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) + } + + 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 + 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] = [] + private var requestIDs: Set = [] init(client: Client) { self.client = client @@ -481,23 +536,36 @@ 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 // 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 ) } } @@ -517,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 @@ -587,24 +658,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..ff7a1ad6 100644 --- a/Tests/MCPTests/ClientTests.swift +++ b/Tests/MCPTests/ClientTests.swift @@ -1,3 +1,4 @@ +import Logging import Foundation import Testing @@ -474,6 +475,195 @@ 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 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") + 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 +963,55 @@ struct ClientTests { await client.disconnect() } } + +private enum BatchFailureError: Swift.Error, Equatable { + case send + case body +} + +private actor BatchFailureResults { + private(set) var errors: [Failure?] = [] + + func observe(_ task: Task) { + Task { + switch await task.result { + case .success: + errors.append(nil) + case .failure(let error): + errors.append(error as? Failure) + } + } + } + + 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)) + } + 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) + } +}