diff --git a/Sources/MCP/Server/Server.swift b/Sources/MCP/Server/Server.swift index 060b51f8..383f541b 100644 --- a/Sources/MCP/Server/Server.swift +++ b/Sources/MCP/Server/Server.swift @@ -4,6 +4,7 @@ import struct Foundation.Data import struct Foundation.Date import class Foundation.JSONDecoder import class Foundation.JSONEncoder +import struct Foundation.UUID /// Model Context Protocol server public actor Server { @@ -165,8 +166,27 @@ public actor Server { private var methodHandlers: [String: RequestHandlerBox] = [:] /// Notification handlers private var notificationHandlers: [String: [NotificationHandlerBox]] = [:] - /// Pending request tasks (for cancellation support) - private var pendingRequestTasks: [ID: Task, Error>] = [:] + /// Built-in cancellation handler, invoked before logging or application handlers. + private var cancellationHandler: NotificationHandlerBox? + private struct ActiveRequest { + let token: UUID + let task: Task, Error> + } + + /// Active request handler tasks keyed by request ID. + private var pendingRequestTasks: [ID: ActiveRequest] = [:] + + private struct DispatchedRequest { + let token: UUID + let task: Task + } + + /// Request dispatch tasks that have not registered a handler task yet. + private var dispatchedRequests: [ID: DispatchedRequest] = [:] + + var trackedRequestTaskCount: Int { + dispatchedRequests.count + pendingRequestTasks.count + } /// Pending requests sent to the client, awaiting responses private var pendingRequests: [ID: AnyPendingRequest] = [:] @@ -208,7 +228,9 @@ public actor Server { ) async throws { self.connection = transport registerDefaultHandlers(initializeHook: initializeHook) + if cancellationHandler == nil { registerCancellationHandler() + } try await transport.connect() await logger?.debug( @@ -231,10 +253,31 @@ public actor Server { } else if let response = try? decoder.decode(AnyResponse.self, from: data) { await handleResponse(response) } else if let request = try? decoder.decode(AnyRequest.self, from: data) { - // Handle request in a separate task to avoid blocking the receive loop - Task { - _ = try? await self.handleRequest(request, sendResponse: true) + if dispatchedRequests[request.id] != nil || pendingRequestTasks[request.id] != nil { + try await send( + AnyMethod.response( + id: request.id, + error: .invalidRequest("Duplicate outstanding request ID") + ) + ) + continue } + + // Store the dispatch task before the actor can process a following + // cancellation notification. + let token = UUID() + let dispatchTask = Task { + _ = try? await self.handleRequest( + request, + sendResponse: true, + trackingToken: token + ) + self.finishDispatch(id: request.id, token: token) + } + dispatchedRequests[request.id] = DispatchedRequest( + token: token, + task: dispatchTask + ) } else if let message = try? decoder.decode(AnyMessage.self, from: data) { try await handleMessage(message) } else { @@ -279,6 +322,18 @@ public actor Server { task?.cancel() task = nil + let dispatchTasks = dispatchedRequests.values.map(\.task) + dispatchedRequests.removeAll() + for dispatchTask in dispatchTasks { + dispatchTask.cancel() + } + + let handlerTasks = pendingRequestTasks.values.map(\.task) + pendingRequestTasks.removeAll() + for handlerTask in handlerTasks { + handlerTask.cancel() + } + // Clear pending requests with errors let pendingRequestsToCancel = self.pendingRequests self.pendingRequests = [:] @@ -735,9 +790,28 @@ public actor Server { /// - request: The request to handle /// - sendResponse: Whether to send the response immediately (true) or return it (false) /// - Returns: The response when sendResponse is false - private func handleRequest(_ request: Request, sendResponse: Bool = true) - async throws -> Response? - { + private func handleRequest( + _ request: Request, + sendResponse: Bool = true, + trackingToken: UUID = UUID() + ) async throws -> Response? { + try Task.checkCancellation() + + let hasDifferentDispatch = + dispatchedRequests[request.id].map { $0.token != trackingToken } ?? false + if pendingRequestTasks[request.id] != nil || hasDifferentDispatch { + let response = AnyMethod.response( + id: request.id, + error: .invalidRequest("Duplicate outstanding request ID") + ) + if sendResponse { + try Task.checkCancellation() + try await send(response) + return nil + } + return response + } + // Check if this is a pre-processed error request (empty method) if request.method.isEmpty && !sendResponse { // This is a placeholder for an invalid request that couldn't be parsed in batch mode @@ -753,6 +827,7 @@ public actor Server { "method": "\(request.method)", "id": "\(request.id)", ]) + try Task.checkCancellation() if configuration.strict { // The client SHOULD NOT send requests other than pings @@ -771,10 +846,12 @@ public actor Server { let response = AnyMethod.response(id: request.id, error: error) if sendResponse { + try Task.checkCancellation() try await send(response) return nil } + try Task.checkCancellation() return response } @@ -784,6 +861,7 @@ public actor Server { let httpContext = await (connection as? any HTTPContextProviding)? .httpRequestContext(for: request.id) let handlerContext = HandlerContext(id: request.id, httpContext: httpContext) + try Task.checkCancellation() // Create a task to handle the request with cancellation support. // Set currentHandlerContext as a task local so handlers see it. @@ -812,16 +890,21 @@ public actor Server { } } - // Store the handler task for potential cancellation - pendingRequestTasks[request.id] = handlerTask + // Store the handler task for potential cancellation. The outer dispatch task remains tracked + // until the request finishes so stop/cancellation can prevent a late send after restart. + pendingRequestTasks[request.id] = ActiveRequest(token: trackingToken, task: handlerTask) - // Ensure cleanup happens regardless of success or failure + // Ensure cleanup happens regardless of success or failure, without removing a newer + // request that reused the same ID after cancellation or restart. defer { + if pendingRequestTasks[request.id]?.token == trackingToken { pendingRequestTasks.removeValue(forKey: request.id) } + } do { let response = try await handlerTask.value + try Task.checkCancellation() if sendResponse { try await send(response) @@ -838,15 +921,29 @@ public actor Server { let response = AnyMethod.response(id: request.id, error: mcpError) if sendResponse { + try Task.checkCancellation() try await send(response) return nil } + try Task.checkCancellation() return response } } private func handleMessage(_ message: Message) async throws { + // Cancellation must be applied before logging or application handlers can suspend. + if message.method == CancelledNotification.name, let cancellationHandler { + do { + try await cancellationHandler(message) + } catch { + await logger?.error( + "Error handling cancellation notification", + metadata: ["error": "\(error)"] + ) + } + } + await logger?.trace( "Processing notification", metadata: ["method": "\(message.method)"]) @@ -969,26 +1066,40 @@ public actor Server { self.isInitialized = true } - /// Cancel and remove a pending request task - private func removePendingRequest(id: ID) -> Task, Error>? { - pendingRequestTasks.removeValue(forKey: id) + private enum CancellationDisposition { + case active + case dispatched + case unknown + } + + /// Cancel an active handler or a request that is still registering its handler. + private func cancelRequest(id: ID) -> CancellationDisposition { + var disposition = CancellationDisposition.unknown + if let active = pendingRequestTasks[id] { + active.task.cancel() + disposition = .active + } + if let dispatched = dispatchedRequests[id] { + dispatched.task.cancel() + if case .unknown = disposition { + disposition = .dispatched + } + } + return disposition + } + + private func finishDispatch(id: ID, token: UUID) { + guard dispatchedRequests[id]?.token == token else { return } + dispatchedRequests.removeValue(forKey: id) } private func registerCancellationHandler() { - onNotification(CancelledNotification.self) { [weak self] message in + cancellationHandler = TypedNotificationHandler { [weak self] message in guard let self = self else { return } let requestId = message.params.requestId let reason = message.params.reason - await self.logger?.debug( - "Received cancellation notification", - metadata: [ - "requestId": requestId.map { "\($0)" } ?? "none", - "reason": reason.map { "\($0)" } ?? "none", - ] - ) - guard let requestId = requestId else { await self.logger?.warning( "Received cancellation notification with no requestId (violates spec MUST)", @@ -997,16 +1108,30 @@ public actor Server { return } - // Cancel the pending request task if it exists and remove from tracking - if let task = await self.removePendingRequest(id: requestId) { - task.cancel() + // Apply cancellation before any logging suspension so registration cannot overtake it. + let disposition = await self.cancelRequest(id: requestId) + await self.logger?.debug( + "Received cancellation notification", + metadata: [ + "requestId": "\(requestId)", + "reason": reason.map { "\($0)" } ?? "none", + ] + ) + + switch disposition { + case .active: await self.logger?.debug( "Cancelled request", metadata: ["requestId": "\(requestId)"] ) - } else { - // Request may have already completed or is unknown - // Per MCP spec, we should ignore this gracefully + case .dispatched: + await self.logger?.debug( + "Cancelled request before handler registration", + metadata: ["requestId": "\(requestId)"] + ) + case .unknown: + // Request may have already completed or is unknown. + // Per MCP spec, ignore this gracefully. await self.logger?.trace( "Cancellation notification for unknown or completed request", metadata: ["requestId": "\(requestId)"] diff --git a/Tests/MCPTests/CancellationTests.swift b/Tests/MCPTests/CancellationTests.swift index 9d4412ef..73d3828b 100644 --- a/Tests/MCPTests/CancellationTests.swift +++ b/Tests/MCPTests/CancellationTests.swift @@ -3,6 +3,49 @@ import Testing @testable import MCP +private actor EarlyCancellationState { + private var handlerStarted = false + private var cancellationProcessed = false + private var applicationCancellationHandlerStarted = false + + func markHandlerStarted() { + handlerStarted = true + } + + func markCancellationProcessed() { + cancellationProcessed = true + } + + func markApplicationCancellationHandlerStarted() { + applicationCancellationHandlerStarted = true + } + + func didStartHandler() -> Bool { + handlerStarted + } + + func didProcessCancellation() -> Bool { + cancellationProcessed + } + + func didStartApplicationCancellationHandler() -> Bool { + applicationCancellationHandlerStarted + } +} + +private func eventually( + timeout: Duration = .seconds(1), + condition: @escaping @Sendable () async -> Bool +) async -> Bool { + let clock = ContinuousClock() + let deadline = clock.now.advanced(by: timeout) + while clock.now < deadline { + if await condition() { return true } + try? await Task.sleep(for: .milliseconds(5)) + } + return await condition() +} + @Suite("Cancellation Tests") struct CancellationTests { @Test("Client sends CancelledNotification") @@ -133,7 +176,9 @@ struct CancellationTests { } await server.withMethodHandler(CallTool.self) { params in - return .init(content: [.text(text: "Result for \(params.name)", annotations: nil, _meta: nil)], isError: false) + return .init( + content: [.text(text: "Result for \(params.name)", annotations: nil, _meta: nil)], + isError: false) } // Start server and connect client @@ -142,7 +187,8 @@ struct CancellationTests { try await Task.sleep(for: .milliseconds(50)) // Use the callTool overload that returns RequestContext (non-async version) - let context: RequestContext = try await client.callTool(name: "testTool", arguments: ["test": "value"]) + let context: RequestContext = try await client.callTool( + name: "testTool", arguments: ["test": "value"]) // Verify we got a context #expect(context.requestID != ID(stringLiteral: "")) @@ -174,7 +220,9 @@ struct CancellationTests { await server.withMethodHandler(CallTool.self) { params in try await Task.sleep(for: .seconds(5)) - return .init(content: [.text(text: "Should not reach here", annotations: nil, _meta: nil)], isError: false) + return .init( + content: [.text(text: "Should not reach here", annotations: nil, _meta: nil)], + isError: false) } // Start server and connect client @@ -183,7 +231,8 @@ struct CancellationTests { try await Task.sleep(for: .milliseconds(50)) // Use the callTool overload that returns RequestContext (non-async version) - let context: RequestContext = try await client.callTool(name: "slowTool", arguments: [:]) + let context: RequestContext = try await client.callTool( + name: "slowTool", arguments: [:]) // Cancel after a short delay try await Task.sleep(for: .milliseconds(50)) @@ -246,4 +295,104 @@ struct CancellationTests { await client.disconnect() await server.stop() } + + @Test("A handler that swallows cancellation cannot send a late response") + func testSwallowedCancellationDoesNotSendResponse() async throws { + let transport = MockTransport() + let state = EarlyCancellationState() + let server = Server(name: "TestServer", version: "1.0") + try await server.start(transport: transport) + await server.withMethodHandler(Ping.self) { _ in + await state.markHandlerStarted() + do { + try await Task.sleep(for: .seconds(5)) + } catch is CancellationError { + // Deliberately swallow cancellation to verify the outer dispatch still blocks send. + } + return Empty() + } + await server.onNotification(CancelledNotification.self) { _ in + await state.markCancellationProcessed() + } + + let request = Ping.request() + try await transport.queue(request: request) + #expect(await eventually { await state.didStartHandler() }) + try await transport.queue( + notification: CancelledNotification.message( + .init(requestId: request.id, reason: "test swallowed cancellation") + ) + ) + + #expect(await eventually { await state.didProcessCancellation() }) + #expect(await eventually { await server.trackedRequestTaskCount == 0 }) + #expect(await transport.sentMessages.isEmpty) + await server.stop() + } + + @Test("Duplicate outstanding request IDs are rejected without replacing tracking") + func testDuplicateOutstandingRequestIDIsRejected() async throws { + let transport = MockTransport() + let server = Server(name: "TestServer", version: "1.0") + await transport.blockHTTPContext() + try await server.start(transport: transport) + + let request = Ping.request() + try await transport.queue(request: request) + #expect(await eventually { await transport.hasRequestedHTTPContext() }) + try await transport.queue(request: request) + + #expect(await eventually { await transport.sentMessages.count == 1 }) + let duplicateResponse = try #require(await transport.sentMessages.first) + #expect(duplicateResponse.contains("Duplicate outstanding request ID")) + #expect(await server.trackedRequestTaskCount == 1) + + let duplicateBatch = try JSONEncoder().encode([request]) + await transport.queue(data: duplicateBatch) + #expect(await eventually { await transport.sentMessages.count == 2 }) + let duplicateBatchResponse = try #require(await transport.sentMessages.last) + #expect(duplicateBatchResponse.contains("Duplicate outstanding request ID")) + #expect(await server.trackedRequestTaskCount == 1) + + await transport.releaseHTTPContext() + #expect(await eventually { await server.trackedRequestTaskCount == 0 }) + await server.stop() + } + + @Test("Cancellation received before handler registration prevents handler start") + func testEarlyCancellationPreventsHandlerStart() async throws { + let transport = MockTransport() + let state = EarlyCancellationState() + let server = Server(name: "TestServer", version: "1.0") + await transport.blockHTTPContext() + await server.onNotification(CancelledNotification.self) { _ in + await state.markApplicationCancellationHandlerStarted() + try await Task.sleep(for: .milliseconds(200)) + await state.markCancellationProcessed() + } + try await server.start(transport: transport) + await server.withMethodHandler(Ping.self) { _ in + await state.markHandlerStarted() + return Empty() + } + + let request = Ping.request() + try await transport.queue(request: request) + #expect(await eventually { await transport.hasRequestedHTTPContext() }) + + try await transport.queue( + notification: CancelledNotification.message( + .init(requestId: request.id, reason: "cancel before registration") + ) + ) + #expect(await eventually { await state.didStartApplicationCancellationHandler() }) + await transport.releaseHTTPContext() + + #expect(await eventually { await server.trackedRequestTaskCount == 0 }) + #expect(await state.didStartHandler() == false) + #expect(await transport.sentMessages.isEmpty) + #expect(await eventually { await state.didProcessCancellation() }) + + await server.stop() + } } diff --git a/Tests/MCPTests/Helpers/MockTransport.swift b/Tests/MCPTests/Helpers/MockTransport.swift index 5979e9ad..726c8e44 100644 --- a/Tests/MCPTests/Helpers/MockTransport.swift +++ b/Tests/MCPTests/Helpers/MockTransport.swift @@ -8,7 +8,7 @@ import struct Foundation.POSIXError @testable import MCP /// Mock transport for testing -actor MockTransport: Transport { +actor MockTransport: Transport, HTTPContextProviding { var logger: Logger let encoder = JSONEncoder() @@ -31,6 +31,9 @@ actor MockTransport: Transport { private(set) var receivedMessages: [String] = [] private var dataStreamContinuation: AsyncThrowingStream.Continuation? + private var shouldBlockHTTPContext = false + private var requestedHTTPContextIDs: Set = [] + private var httpContextContinuations: [ID: CheckedContinuation] = [:] var shouldFailConnect = false var shouldFailSend = false @@ -80,6 +83,33 @@ actor MockTransport: Transport { shouldFailSend = shouldFail } + func httpRequestContext(for id: ID) async -> HTTPRequest? { + requestedHTTPContextIDs.insert(id) + if shouldBlockHTTPContext { + await withCheckedContinuation { continuation in + httpContextContinuations[id] = continuation + } + } + return nil + } + + func blockHTTPContext() { + shouldBlockHTTPContext = true + } + + func hasRequestedHTTPContext() -> Bool { + !requestedHTTPContextIDs.isEmpty + } + + func releaseHTTPContext() { + shouldBlockHTTPContext = false + let continuations = Array(httpContextContinuations.values) + httpContextContinuations.removeAll() + for continuation in continuations { + continuation.resume() + } + } + func queue(data: Data) { if let continuation = dataStreamContinuation { continuation.yield(data)