diff --git a/Sources/MCP/Base/Transports/HTTPServer/HTTPServerTypes.swift b/Sources/MCP/Base/Transports/HTTPServer/HTTPServerTypes.swift index bd56191c..7e14fa70 100644 --- a/Sources/MCP/Base/Transports/HTTPServer/HTTPServerTypes.swift +++ b/Sources/MCP/Base/Transports/HTTPServer/HTTPServerTypes.swift @@ -225,6 +225,18 @@ public protocol HTTPContextProviding: Sendable { func httpRequestContext(for id: ID) async -> HTTPRequest? } +/// A transport-internal mapping between private routed ids and client wire ids. +package protocol RoutedRequestIDProviding: Sendable { + func originalRequestID(for routedID: ID) async -> ID? + + /// Returns `nil` when the wire id has no active exchange or is ambiguous. + func routedRequestID(for originalID: ID) async -> ID? +} + +package protocol CancelledHTTPExchangeCompleting: Sendable { + func completeCancelledExchange(for routedID: ID, reason: String?) async +} + // MARK: - JSON-RPC Message Classification /// Classifies a raw JSON-RPC message for routing purposes. diff --git a/Sources/MCP/Base/Transports/HTTPServer/StatelessHTTPServerTransport.swift b/Sources/MCP/Base/Transports/HTTPServer/StatelessHTTPServerTransport.swift index abe9576c..d31a0db6 100644 --- a/Sources/MCP/Base/Transports/HTTPServer/StatelessHTTPServerTransport.swift +++ b/Sources/MCP/Base/Transports/HTTPServer/StatelessHTTPServerTransport.swift @@ -29,7 +29,10 @@ import Logging /// - Session management is handled externally or not needed /// /// For full streaming and session support, use ``StatefulHTTPServerTransport`` instead. -public actor StatelessHTTPServerTransport: Transport, HTTPContextProviding { +/// A cancelled request ends with an empty SSE response so no JSON-RPC reply is sent. +public actor StatelessHTTPServerTransport: + Transport, HTTPContextProviding, RoutedRequestIDProviding, CancelledHTTPExchangeCompleting +{ public nonisolated let logger: Logger // MARK: - Dependencies @@ -48,15 +51,28 @@ public actor StatelessHTTPServerTransport: Transport, HTTPContextProviding { // MARK: - Response waiters - /// Maps request ID → continuation waiting for the server's response. - /// When the server calls `send()` with a response, the matching continuation is resumed. - private var responseWaiters: [String: CheckedContinuation] = [:] + private enum ResponseOutcome { + case data(Data) + case cancelled + } + + private struct ResponseWaiter { + let originalID: ID + let continuation: CheckedContinuation + } + + /// Maps a transport-private exchange ID to the matching HTTP response waiter. + private var responseWaiters: [String: ResponseWaiter] = [:] /// Maps request ID → originating HTTP request, surfaced to handlers via /// ``Server/currentHTTPContext``. Entries live only while a JSON-RPC request /// is in flight. private var httpRequestContexts: [String: HTTPRequest] = [:] + /// Preserves the existing raw-ID context lookup for callers outside `Server`. + /// Multiple entries are possible because JSON-RPC IDs are scoped to each client. + private var exchangeIDsByRequestID: [ID: [String]] = [:] + // MARK: - Init /// Creates a new stateless HTTP server transport. @@ -117,14 +133,20 @@ public actor StatelessHTTPServerTransport: Transport, HTTPContextProviding { switch kind { case .response(let id): - guard let continuation = responseWaiters.removeValue(forKey: id) else { + guard let waiter = responseWaiters.removeValue(forKey: id) else { logger.debug( "No waiter for response, may have timed out", metadata: ["requestID": "\(id)"] ) return } - continuation.resume(returning: data) + do { + let response = try restoringResponseID(in: data, to: waiter.originalID) + waiter.continuation.resume(returning: .data(response)) + } catch { + waiter.continuation.resume(throwing: error) + throw error + } case .notification(let method): logger.debug( @@ -153,6 +175,9 @@ public actor StatelessHTTPServerTransport: Transport, HTTPContextProviding { /// - **GET**: 405 Method Not Allowed /// - **DELETE**: 405 Method Not Allowed /// - Others: 405 Method Not Allowed + /// + /// If the adapter cancels this task because its HTTP exchange ended, it must not write the + /// returned response. The adapter owns disconnect and timeout behavior. public func handleRequest(_ request: HTTPRequest) async -> HTTPResponse { if terminated { return .error( @@ -190,6 +215,12 @@ public actor StatelessHTTPServerTransport: Transport, HTTPContextProviding { .parseError("Invalid JSON-RPC message") ) } + if case .notification = messageKind, + let json = try? JSONSerialization.jsonObject(with: body) as? [String: Any], + json.keys.contains("id") + { + return .error(statusCode: 400, .parseError("Invalid JSON-RPC request id")) + } // Build validation context let context = HTTPValidationContext( @@ -206,7 +237,12 @@ public actor StatelessHTTPServerTransport: Transport, HTTPContextProviding { // Handle by message type switch messageKind { - case .notification, .response: + case .notification: + // Yield to server and return 202 Accepted + incomingContinuation.yield(body) + return .accepted() + + case .response: // Yield to server and return 202 Accepted incomingContinuation.yield(body) return .accepted() @@ -221,32 +257,245 @@ public actor StatelessHTTPServerTransport: Transport, HTTPContextProviding { requestID: String, request: HTTPRequest ) async -> HTTPResponse { - httpRequestContexts[requestID] = request + let exchangeID = makeExchangeID(excluding: requestID) + let routedBody: Data + let originalID: ID + do { + (routedBody, originalID) = try routingRequest(body, as: exchangeID) + } catch { + return .error( + statusCode: 400, + .parseError("Invalid JSON-RPC request id") + ) + } + + registerHTTPContext(request, exchangeID: exchangeID, requestID: originalID) // Yield the incoming message to the server - incomingContinuation.yield(body) + incomingContinuation.yield(routedBody) // Wait for the server to process and send a response - let responseData: Data + let outcome: ResponseOutcome do { - responseData = try await withCheckedThrowingContinuation { continuation in - responseWaiters[requestID] = continuation + outcome = try await withTaskCancellationHandler { + try await withCheckedThrowingContinuation { continuation in + responseWaiters[exchangeID] = ResponseWaiter( + originalID: originalID, + continuation: continuation + ) + } + } onCancel: { + Task { await self.cancelWaitingExchange(exchangeID) } } + } catch MCPError.connectionClosed { + removeHTTPContext(exchangeID: exchangeID, requestID: originalID) + return .error(statusCode: 503, .connectionClosed) } catch { - httpRequestContexts.removeValue(forKey: requestID) + removeHTTPContext(exchangeID: exchangeID, requestID: originalID) return .error( statusCode: 500, .internalError("Error processing request: \(error.localizedDescription)") ) } - httpRequestContexts.removeValue(forKey: requestID) - return .data(responseData, headers: [HTTPHeaderName.contentType: ContentType.json]) + removeHTTPContext(exchangeID: exchangeID, requestID: originalID) + switch outcome { + case .data(let responseData): + return .data(responseData, headers: [HTTPHeaderName.contentType: ContentType.json]) + case .cancelled: + if request.header(HTTPHeaderName.accept)?.lowercased().contains(ContentType.sse) == true { + return .stream( + AsyncThrowingStream { $0.finish() }, + headers: [HTTPHeaderName.contentType: ContentType.sse] + ) + } + let cancellation = AnyMethod.response( + id: originalID, + error: .serverError(code: -32800, message: "Request cancelled") + ) + guard let data = try? JSONEncoder().encode(cancellation) else { + return .error(statusCode: 500, .internalError("Failed to encode cancellation")) + } + return .data(data, headers: [HTTPHeaderName.contentType: ContentType.json]) + } + } + + private func makeExchangeID(excluding requestID: String) -> String { + var exchangeID: String + repeat { + exchangeID = UUID().uuidString + } while exchangeID == requestID + || responseWaiters[exchangeID] != nil + || httpRequestContexts[exchangeID] != nil + || exchangeIDsByRequestID[.string(exchangeID)] != nil + return exchangeID + } + + private func routingRequest(_ body: Data, as exchangeID: String) throws -> (Data, ID) { + try replacingTopLevelID(in: body, with: .string(exchangeID)) + } + + private func restoringResponseID(in data: Data, to originalID: ID) throws -> Data { + try replacingTopLevelID(in: data, with: originalID).0 + } + + /// Replace only the top-level ID. Re-encoding the whole message can change numeric payloads. + private func replacingTopLevelID(in data: Data, with replacement: ID) throws -> (Data, ID) { + guard (try JSONSerialization.jsonObject(with: data)) is [String: Any] else { + throw MCPError.parseError("Invalid JSON-RPC object") + } + + let bytes = Array(data) + var index = 0 + func skipWhitespace() { + while index < bytes.count && [9, 10, 13, 32].contains(bytes[index]) { index += 1 } + } + func stringEnd() throws -> Int { + guard index < bytes.count, bytes[index] == 34 else { + throw MCPError.parseError("Invalid JSON-RPC member") + } + index += 1 + var escaped = false + while index < bytes.count { + let byte = bytes[index] + index += 1 + if escaped { escaped = false } + else if byte == 92 { escaped = true } + else if byte == 34 { return index } + } + throw MCPError.parseError("Invalid JSON string") + } + + skipWhitespace() + index += 1 // Opening brace, checked by JSONSerialization above. + var idRange: Range? + while index < bytes.count { + skipWhitespace() + if bytes[index] == 125 { break } + let keyStart = index + let keyEnd = try stringEnd() + let key = try JSONDecoder().decode(String.self, from: Data(bytes[keyStart.. valueStart && [9, 10, 13, 32].contains(bytes[valueEnd - 1]) { + valueEnd -= 1 + } + if key == "id" { + guard idRange == nil else { + throw MCPError.parseError("Duplicate JSON-RPC request id") + } + idRange = valueStart.. HTTPRequest? { - httpRequestContexts[id.description] + if case .string(let exchangeID) = id, + let request = httpRequestContexts[exchangeID] + { + return request + } + guard let exchangeIDs = exchangeIDsByRequestID[id], + exchangeIDs.count == 1, + let exchangeID = exchangeIDs.first + else { + return nil + } + return httpRequestContexts[exchangeID] + } + + package func originalRequestID(for routedID: ID) -> ID? { + guard case .string(let exchangeID) = routedID else { return nil } + return responseWaiters[exchangeID]?.originalID + } + + package func routedRequestID(for originalID: ID) -> ID? { + guard let exchangeIDs = exchangeIDsByRequestID[originalID], + exchangeIDs.count == 1, + let exchangeID = exchangeIDs.first + else { + return nil + } + return .string(exchangeID) } // MARK: - Termination @@ -258,12 +507,19 @@ public actor StatelessHTTPServerTransport: Transport, HTTPContextProviding { logger.debug("Stateless HTTP server transport terminated") // Cancel all waiting continuations - for (id, continuation) in responseWaiters { - continuation.resume(throwing: MCPError.connectionClosed) - logger.debug("Cancelled waiter for request", metadata: ["requestID": "\(id)"]) + for (exchangeID, waiter) in responseWaiters { + waiter.continuation.resume(throwing: MCPError.connectionClosed) + logger.debug( + "Cancelled waiter for request", + metadata: [ + "exchangeID": "\(exchangeID)", + "requestID": "\(waiter.originalID)", + ] + ) } responseWaiters.removeAll() httpRequestContexts.removeAll() + exchangeIDsByRequestID.removeAll() // Close incoming stream incomingContinuation.finish() diff --git a/Sources/MCP/Server/Server.swift b/Sources/MCP/Server/Server.swift index 55a8e567..0de2c106 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 { @@ -172,8 +173,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] = [:] @@ -215,7 +235,9 @@ public actor Server { ) async throws { self.connection = transport registerDefaultHandlers(initializeHook: initializeHook) - registerCancellationHandler() + if cancellationHandler == nil { + registerCancellationHandler() + } try await transport.connect() await logger?.debug( @@ -237,10 +259,33 @@ 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 { @@ -285,6 +330,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 = [:] @@ -725,9 +782,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 @@ -743,6 +819,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 @@ -761,10 +838,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 } @@ -773,7 +852,11 @@ public actor Server { // that don't carry HTTP context (stdio, in-memory) don't conform. let httpContext = await (connection as? any HTTPContextProviding)? .httpRequestContext(for: request.id) - let handlerContext = HandlerContext(id: request.id, httpContext: httpContext) + let handlerRequestID = + await (connection as? any RoutedRequestIDProviding)? + .originalRequestID(for: request.id) ?? request.id + let handlerContext = HandlerContext(id: handlerRequestID, 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. @@ -802,16 +885,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 { - pendingRequestTasks.removeValue(forKey: request.id) + 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) @@ -828,15 +916,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)"]) @@ -956,25 +1058,41 @@ 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) { 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)", @@ -983,16 +1101,38 @@ 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() + // Route the wire ID before cancellation. An ambiguous ID must not pick one client. + guard let routedID = await self.routedRequestID(for: requestId) else { return } + let disposition = await self.cancelRequest(id: routedID) + switch disposition { + case .active, .dispatched: + await (self.connection as? any CancelledHTTPExchangeCompleting)? + .completeCancelledExchange(for: routedID, reason: reason) + case .unknown: + break + } + 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)"] @@ -1001,6 +1141,13 @@ public actor Server { } } + private func routedRequestID(for originalID: ID) async -> ID? { + guard let provider = connection as? any RoutedRequestIDProviding else { + return originalID + } + return await provider.routedRequestID(for: originalID) + } + /// Cancel a request by sending a CancelledNotification to the client. /// /// This is used when the server needs to cancel an in-progress request it made to the client diff --git a/Tests/MCPTests/CancellationTests.swift b/Tests/MCPTests/CancellationTests.swift index 9d4412ef..5c5a45b6 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") @@ -129,11 +172,15 @@ struct CancellationTests { // Register a tool handler await server.withMethodHandler(ListTools.self) { _ in - .init(tools: [Tool(name: "testTool", description: "A test tool", inputSchema: .object([:]))]) + .init(tools: [ + Tool(name: "testTool", description: "A test tool", inputSchema: .object([:])) + ]) } 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,14 +189,16 @@ 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: "")) // Get the result let result = try await context.value - #expect(result.content == [.text(text: "Result for testTool", annotations: nil, _meta: nil)]) + #expect( + result.content == [.text(text: "Result for testTool", annotations: nil, _meta: nil)]) #expect(result.isError == false) await client.disconnect() @@ -169,12 +218,16 @@ struct CancellationTests { // Register a tool handler that takes time await server.withMethodHandler(ListTools.self) { _ in - .init(tools: [Tool(name: "slowTool", description: "A slow tool", inputSchema: .object([:]))]) + .init(tools: [ + Tool(name: "slowTool", description: "A slow tool", inputSchema: .object([:])) + ]) } 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 +236,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 +300,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/HTTPServerTransportTests.swift b/Tests/MCPTests/HTTPServerTransportTests.swift index 8e1c89b4..0b77a803 100644 --- a/Tests/MCPTests/HTTPServerTransportTests.swift +++ b/Tests/MCPTests/HTTPServerTransportTests.swift @@ -29,6 +29,15 @@ private func makeNotificationBody(method: String = "notifications/initialized") return try! JSONSerialization.data(withJSONObject: json) } +private func makeCancelledNotificationBody(requestID: Any) -> Data { + let json: [String: Any] = [ + "jsonrpc": "2.0", + "method": "notifications/cancelled", + "params": ["requestId": requestID], + ] + return try! JSONSerialization.data(withJSONObject: json) +} + private func makeRequestBody(id: String = "2", method: String = "tools/list") -> Data { let json: [String: Any] = [ "jsonrpc": "2.0", @@ -89,12 +98,24 @@ private func makeStatelessPOSTRequest(body: Data) -> HTTPRequest { method: "POST", headers: [ "Content-Type": "application/json", - "Accept": "application/json", + "Accept": "application/json, text/event-stream", ], body: body ) } +/// `requestId` is `Any` so tests can exercise both string and integer JSON-RPC ids. +private func makeCancelledNotificationBody(requestId: Any, reason: String? = nil) -> Data { + var params: [String: Any] = ["requestId": requestId] + if let reason { params["reason"] = reason } + let json: [String: Any] = [ + "jsonrpc": "2.0", + "method": "notifications/cancelled", + "params": params, + ] + return try! JSONSerialization.data(withJSONObject: json) +} + private func makeStatefulTransport( sessionIDGenerator: any SessionIDGenerator = UUIDSessionIDGenerator() ) -> StatefulHTTPServerTransport { @@ -139,6 +160,43 @@ private actor ChunkCollector { func getChunks() -> [Data] { chunks } } +/// Guards a one-shot resume so two racing tasks can't double-resume a continuation. +private actor OnceGuard { + private var fired = false + func tryFire() -> Bool { + guard !fired else { return false } + fired = true + return true + } +} + +/// Races `operation` against `timeout` without blocking on `operation` if it never +/// completes. Returns `operation`'s result if it finishes first, or `nil` if the +/// timeout elapses first — turning an indefinite hang into a bounded, observable +/// test outcome. If `operation` really does hang, its task is simply abandoned +/// (not cancelled: a suspended `CheckedContinuation`-based wait doesn't observe +/// cancellation), which is fine for a short-lived test process. +private func raceAgainstTimeout( + _ timeout: Duration, + operation: @escaping @Sendable () async -> T +) async -> T? { + let guardActor = OnceGuard() + return await withCheckedContinuation { (continuation: CheckedContinuation) in + Task { + let value = await operation() + if await guardActor.tryFire() { + continuation.resume(returning: value) + } + } + Task { + try? await Task.sleep(for: timeout) + if await guardActor.tryFire() { + continuation.resume(returning: nil) + } + } + } +} + private func drainSSEStream( _ response: HTTPResponse, maxChunks: Int = 10, @@ -158,6 +216,25 @@ private func drainSSEStream( return await collector.getChunks() } +private func nextValue( + from stream: AsyncStream, + timeout: Duration = .seconds(1) +) async -> T? { + await withTaskGroup(of: T?.self) { group in + group.addTask { + var iterator = stream.makeAsyncIterator() + return await iterator.next() + } + group.addTask { + try? await Task.sleep(for: timeout) + return nil + } + let value = await group.next() ?? nil + group.cancelAll() + return value + } +} + /// Initializes a stateful transport session and returns the session ID. /// Spawns a background task to consume the receive stream and send the init response. private func initializeSession( @@ -178,9 +255,13 @@ private func initializeSession( let id = json["id"] { let idString: String - if let s = id as? String { idString = s } - else if let n = id as? Int { idString = String(n) } - else { continue } + if let s = id as? String { + idString = s + } else if let n = id as? Int { + idString = String(n) + } else { + continue + } let responseJSON: [String: Any] = [ "jsonrpc": "2.0", @@ -852,6 +933,42 @@ struct StatefulHTTPServerTransportTests { @Suite("StatelessHTTPServerTransport Tests") struct StatelessHTTPServerTransportTests { + @Test("Routing changes only the top-level ID, preserving numeric payloads") + func testRoutingPreservesPayloadBytes() async throws { + let transport = makeStatelessTransport() + try await transport.connect() + let stream = await transport.receive() + var iterator = stream.makeAsyncIterator() + let request = Data(#"{"jsonrpc":"2.0","id":7,"method":"tools/call","params":{"amounts":[0.1,19.99,0.3],"id":"nested"}}"#.utf8) + let pending = Task { + await transport.handleRequest(makeStatelessPOSTRequest(body: request)) + } + let routed = try #require(try await iterator.next()) + let routedText = String(decoding: routed, as: UTF8.self) + #expect(routedText.contains(#""amounts":[0.1,19.99,0.3],"id":"nested""#)) + let json = try #require(JSONSerialization.jsonObject(with: routed) as? [String: Any]) + let routedID = try #require(json["id"] as? String) + #expect(routedID != "7") + + let response = Data(#"{"jsonrpc":"2.0","id":"\#(routedID)","result":{"amounts":[0.1,19.99,0.3],"id":"nested"}}"#.utf8) + try await transport.send(response) + let completed = await pending.value + #expect(completed.statusCode == 200) + #expect(String(decoding: completed.bodyData ?? Data(), as: UTF8.self) + == #"{"jsonrpc":"2.0","id":7,"result":{"amounts":[0.1,19.99,0.3],"id":"nested"}}"#) + await transport.disconnect() + } + + @Test("Boolean and fractional request IDs are rejected") + func testInvalidRequestIDs() async { + let transport = makeStatelessTransport() + for id in ["true", "0.1", "1.0", "null"] { + let body = Data(#"{"jsonrpc":"2.0","id":\#(id),"method":"tools/list"}"#.utf8) + let response = await transport.handleRequest(makeStatelessPOSTRequest(body: body)) + #expect(response.statusCode == 400) + } + } + // MARK: - Lifecycle @Test("Connect succeeds") @@ -895,7 +1012,17 @@ struct StatelessHTTPServerTransportTests { try await transport.connect() let requestBody = makeRequestBody(id: "42", method: "tools/list") - let responseBody = makeResponseBody(id: "42") + + let responseRouter = Task { + let stream = await transport.receive() + var iterator = stream.makeAsyncIterator() + let body = try #require(try await iterator.next()) + let json = try #require( + try JSONSerialization.jsonObject(with: body) as? [String: Any] + ) + let exchangeID = try #require(json["id"] as? String) + try await transport.send(makeResponseBody(id: exchangeID)) + } // handleRequest blocks waiting for response let handleTask = Task { @@ -904,16 +1031,183 @@ struct StatelessHTTPServerTransportTests { ) } - // Give handleRequest time to register the waiter - try await Task.sleep(for: .milliseconds(50)) - - // Consume the request from receive and send the response - try await transport.send(responseBody) - + try await responseRouter.value let httpResponse = await handleTask.value #expect(httpResponse.statusCode == 200) - #expect(httpResponse.bodyData == responseBody) #expect(httpResponse.headers[HTTPHeaderName.contentType] == ContentType.json) + let responseData = try #require(httpResponse.bodyData) + let responseJSON = try #require( + try JSONSerialization.jsonObject(with: responseData) as? [String: Any] + ) + #expect(responseJSON["id"] as? String == "42") + let result = try #require(responseJSON["result"] as? [String: Any]) + let tools = try #require(result["tools"] as? [Any]) + #expect(tools.isEmpty) + + await transport.disconnect() + } + + @Test("Concurrent requests sharing a wire ID route to their own HTTP exchange") + func testConcurrentRequestsWithSameWireID() async throws { + let transport = makeStatelessTransport() + try await transport.connect() + + func request(marker: String) -> HTTPRequest { + let body = try! JSONSerialization.data(withJSONObject: [ + "jsonrpc": "2.0", + "id": 7, + "method": "tools/list", + "params": ["marker": marker], + ]) + return HTTPRequest( + method: "POST", + headers: [ + "Content-Type": "application/json", + "Accept": "application/json", + "Authorization": "Bearer \(marker)", + ], + body: body, + path: "/mcp/\(marker)" + ) + } + + let responseRouter = Task { + let stream = await transport.receive() + var iterator = stream.makeAsyncIterator() + var responses: [Data] = [] + + for _ in 0..<2 { + let body = try #require(try await iterator.next()) + let json = try #require( + try JSONSerialization.jsonObject(with: body) as? [String: Any] + ) + let exchangeID = try #require(json["id"] as? String) + let params = try #require(json["params"] as? [String: Any]) + let marker = try #require(params["marker"] as? String) + let context = await transport.httpRequestContext(for: .string(exchangeID)) + #expect(context?.header("Authorization") == "Bearer \(marker)") + #expect(context?.path == "/mcp/\(marker)") + responses.append( + try JSONSerialization.data(withJSONObject: [ + "jsonrpc": "2.0", + "id": exchangeID, + "result": ["marker": marker], + ]) + ) + } + + #expect(await transport.routedRequestID(for: .number(7)) == nil) + + for response in responses { + try await transport.send(response) + } + } + + let firstTask = Task { await transport.handleRequest(request(marker: "first")) } + let secondTask = Task { await transport.handleRequest(request(marker: "second")) } + + try await responseRouter.value + let first = await firstTask.value + let second = await secondTask.value + + for (response, expectedMarker) in [(first, "first"), (second, "second")] { + #expect(response.statusCode == 200) + let body = try #require(response.bodyData) + let json = try #require( + try JSONSerialization.jsonObject(with: body) as? [String: Any] + ) + #expect(json["id"] as? Int == 7) + let result = try #require(json["result"] as? [String: Any]) + #expect(result["marker"] as? String == expectedMarker) + } + + await transport.disconnect() + } + + @Test("String and number wire IDs route to distinct HTTP exchanges") + func testConcurrentStringAndNumberWireIDs() async throws { + let transport = makeStatelessTransport() + try await transport.connect() + + func request(id: Any, marker: String) -> HTTPRequest { + let body = try! JSONSerialization.data(withJSONObject: [ + "jsonrpc": "2.0", + "id": id, + "method": "tools/list", + "params": ["marker": marker], + ]) + return HTTPRequest( + method: "POST", + headers: [ + "Content-Type": "application/json", + "Accept": "application/json", + "Authorization": "Bearer \(marker)", + ], + body: body, + path: "/mcp/\(marker)" + ) + } + + let responseRouter = Task { + let stream = await transport.receive() + var iterator = stream.makeAsyncIterator() + var responses: [Data] = [] + + for _ in 0..<2 { + let body = try #require(try await iterator.next()) + let json = try #require( + try JSONSerialization.jsonObject(with: body) as? [String: Any] + ) + let exchangeID = try #require(json["id"] as? String) + let params = try #require(json["params"] as? [String: Any]) + let marker = try #require(params["marker"] as? String) + responses.append( + try JSONSerialization.data(withJSONObject: [ + "jsonrpc": "2.0", + "id": exchangeID, + "result": ["marker": marker], + ]) + ) + } + + let stringContext = await transport.httpRequestContext(for: .string("1")) + #expect(stringContext?.header("Authorization") == "Bearer string") + #expect(stringContext?.path == "/mcp/string") + let numberContext = await transport.httpRequestContext(for: .number(1)) + #expect(numberContext?.header("Authorization") == "Bearer number") + #expect(numberContext?.path == "/mcp/number") + + for response in responses { + try await transport.send(response) + } + } + + let stringTask = Task { + await transport.handleRequest(request(id: "1", marker: "string")) + } + let numberTask = Task { + await transport.handleRequest(request(id: 1, marker: "number")) + } + + try await responseRouter.value + let stringResponse = await stringTask.value + let numberResponse = await numberTask.value + + let stringData = try #require(stringResponse.bodyData) + let stringJSON = try #require( + try JSONSerialization.jsonObject(with: stringData) as? [String: Any] + ) + #expect(stringJSON["id"] as? String == "1") + let stringResult = try #require(stringJSON["result"] as? [String: Any]) + #expect(stringResult["marker"] as? String == "string") + + let numberData = try #require(numberResponse.bodyData) + let numberJSON = try #require( + try JSONSerialization.jsonObject(with: numberData) as? [String: Any] + ) + #expect(numberJSON["id"] as? Int == 1) + let numberResult = try #require(numberJSON["result"] as? [String: Any]) + #expect(numberResult["marker"] as? String == "number") await transport.disconnect() } @@ -957,6 +1251,45 @@ struct StatelessHTTPServerTransportTests { await transport.disconnect() } + @Test("Cancellation lookup maps the wire ID to the active HTTP exchange") + func testCancellationLookupMapsToExchangeID() async throws { + let transport = makeStatelessTransport() + try await transport.connect() + + let stream = await transport.receive() + var iterator = stream.makeAsyncIterator() + let handleTask = Task { + await transport.handleRequest( + makeStatelessPOSTRequest(body: makeRequestBody(id: "cancel-me")) + ) + } + + let routedRequest = try #require(try await iterator.next()) + let requestJSON = try #require( + try JSONSerialization.jsonObject(with: routedRequest) as? [String: Any] + ) + let exchangeID = try #require(requestJSON["id"] as? String) + + #expect( + await transport.routedRequestID(for: .string("cancel-me")) + == .string(exchangeID) + ) + + let cancellationBody = makeCancelledNotificationBody(requestID: "cancel-me") + let cancelResponse = await transport.handleRequest( + makeStatelessPOSTRequest( + body: cancellationBody + ) + ) + #expect(cancelResponse.statusCode == 202) + + let routedCancellation = try #require(try await iterator.next()) + #expect(routedCancellation == cancellationBody) + + await transport.disconnect() + _ = await handleTask.value + } + // MARK: - Unsupported Methods @Test("GET returns 405 Method Not Allowed") @@ -1028,28 +1361,68 @@ struct StatelessHTTPServerTransportTests { #expect(response.statusCode == 404) } - @Test("Disconnect cancels waiting requests") - func testDisconnectCancelsWaitingRequests() async throws { + @Test("Shutdown completes waiting requests with 503") + func testShutdownCompletesWaitingRequests() async throws { let transport = makeStatelessTransport() try await transport.connect() + let stream = await transport.receive() + var iterator = stream.makeAsyncIterator() let requestBody = makeRequestBody(id: "cancel-test") - // Start a request that will block let handleTask = Task { await transport.handleRequest( makeStatelessPOSTRequest(body: requestBody) ) } + let routedBody = try #require(try await iterator.next()) + let routedJSON = try #require( + JSONSerialization.jsonObject(with: routedBody) as? [String: Any] + ) + let exchangeID = try #require(routedJSON["id"] as? String) - try await Task.sleep(for: .milliseconds(50)) - - // Disconnect while request is pending await transport.disconnect() - let response = await handleTask.value - // Should return error (500 or similar) since waiter was cancelled - #expect(response.statusCode == 500) + let response = try #require( + await raceAgainstTimeout(.seconds(1)) { await handleTask.value }, + "Shutdown left the HTTP request waiting" + ) + #expect(response.statusCode == 503) + let body = try #require(response.bodyData) + let json = try #require(JSONSerialization.jsonObject(with: body) as? [String: Any]) + let error = try #require(json["error"] as? [String: Any]) + #expect(error["code"] as? Int == MCPError.connectionClosed.code) + #expect(await transport.routedRequestID(for: .string("cancel-test")) == nil) + #expect(await transport.httpRequestContext(for: .string(exchangeID)) == nil) + } + + @Test("Cancelled HTTP caller releases its request mapping") + func testCancelledHTTPCallerClearsExchange() async throws { + let transport = makeStatelessTransport() + try await transport.connect() + let stream = await transport.receive() + var iterator = stream.makeAsyncIterator() + + let handleTask = Task { + await transport.handleRequest( + makeStatelessPOSTRequest(body: makeRequestBody(id: "gone")) + ) + } + let routedBody = try #require(try await iterator.next()) + let routedJSON = try #require( + JSONSerialization.jsonObject(with: routedBody) as? [String: Any] + ) + let exchangeID = try #require(routedJSON["id"] as? String) + + handleTask.cancel() + let completedInTime = await raceAgainstTimeout(.seconds(1)) { await handleTask.value } + if completedInTime == nil { await transport.disconnect() } + _ = try #require(completedInTime, "Cancelled HTTP caller kept its waiter until shutdown") + #expect(await transport.routedRequestID(for: .string("gone")) == nil) + #expect(await transport.httpRequestContext(for: .string(exchangeID)) == nil) + + try await transport.send(makeResponseBody(id: exchangeID)) + await transport.disconnect() } } @@ -1146,11 +1519,17 @@ struct ServerHandlerContextTests { path: "/mcp" ) + let stream = await transport.receive() + var iterator = stream.makeAsyncIterator() + // handleRequest blocks until a response is sent — run it concurrently. let handleTask = Task { await transport.handleRequest(httpRequest) } - // Give the handler time to register its waiter (and store the context). - try await Task.sleep(for: .milliseconds(50)) + let routedBody = try #require(try await iterator.next()) + let routedJSON = try #require( + try JSONSerialization.jsonObject(with: routedBody) as? [String: Any] + ) + let exchangeID = try #require(routedJSON["id"] as? String) let inFlight = await transport.httpRequestContext(for: .string("99")) #expect(inFlight != nil) @@ -1158,7 +1537,7 @@ struct ServerHandlerContextTests { #expect(inFlight?.path == "/mcp") // Send the response to unblock the waiter. - try await transport.send(makeResponseBody(id: "99")) + try await transport.send(makeResponseBody(id: exchangeID)) _ = await handleTask.value let afterResponse = await transport.httpRequestContext(for: .string("99")) @@ -1233,12 +1612,126 @@ struct ServerHandlerContextTests { #expect(await captured.id == .string("call-1")) } + @Test("Stateless cancellation finds the handler by its wire ID") + func testStatelessCancellationFindsHandlerByWireID() async throws { + let (entered, enteredContinuation) = AsyncStream.makeStream() + let (cancelled, cancelledContinuation) = AsyncStream.makeStream() + + let transport = makeStatelessTransport() + let server = Server(name: "TestServer", version: "1.0") + await server.withMethodHandler(CallTool.self) { _ in + enteredContinuation.yield(()) + do { + try await Task.sleep(for: .seconds(5)) + } catch is CancellationError { + cancelledContinuation.yield(()) + throw CancellationError() + } + return CallTool.Result(content: [.text(text: "late", annotations: nil, _meta: nil)]) + } + + try await server.start(transport: transport) + let requestBody = try JSONSerialization.data(withJSONObject: [ + "jsonrpc": "2.0", + "id": "cancel-handler", + "method": "tools/call", + "params": ["name": "slow-tool"] as [String: Any], + ]) + let handleTask = Task { + await transport.handleRequest(makeStatelessPOSTRequest(body: requestBody)) + } + + _ = try #require(await nextValue(from: entered)) + let cancelResponse = await transport.handleRequest( + makeStatelessPOSTRequest( + body: makeCancelledNotificationBody(requestID: "cancel-handler") + ) + ) + #expect(cancelResponse.statusCode == 202) + _ = try #require(await nextValue(from: cancelled)) + + await transport.disconnect() + _ = await handleTask.value + await server.stop() + } + + @Test("Ambiguous cancellation leaves both HTTP requests running") + func testAmbiguousCancellationKeepsBothRequests() async throws { + actor Barrier { + let entered: AsyncStream + private let enteredContinuation: AsyncStream.Continuation + private var waiters: [CheckedContinuation] = [] + + init() { + (entered, enteredContinuation) = AsyncStream.makeStream() + } + + func wait() async { + await withCheckedContinuation { continuation in + waiters.append(continuation) + enteredContinuation.yield(()) + } + } + + func release() { + for waiter in waiters { waiter.resume() } + waiters.removeAll() + } + } + + let barrier = Barrier() + let transport = makeStatelessTransport() + let server = Server(name: "TestServer", version: "1.0") + await server.withMethodHandler(CallTool.self) { _ in + await barrier.wait() + return CallTool.Result(content: [.text(text: "done", annotations: nil, _meta: nil)]) + } + try await server.start(transport: transport) + let (processed, processedContinuation) = AsyncStream.makeStream() + await server.onNotification(CancelledNotification.self) { _ in + processedContinuation.yield(()) + } + + let body = try JSONSerialization.data(withJSONObject: [ + "jsonrpc": "2.0", "id": "shared", "method": "tools/call", + "params": ["name": "slow-tool"], + ]) + let first = Task { await transport.handleRequest(makeStatelessPOSTRequest(body: body)) } + let second = Task { await transport.handleRequest(makeStatelessPOSTRequest(body: body)) } + var entered = barrier.entered.makeAsyncIterator() + _ = try #require(await entered.next()) + _ = try #require(await entered.next()) + + #expect(await transport.httpRequestContext(for: .string("shared")) == nil) + #expect(await transport.routedRequestID(for: .string("shared")) == nil) + + let cancellation = await transport.handleRequest( + makeStatelessPOSTRequest(body: makeCancelledNotificationBody(requestID: "shared")) + ) + #expect(cancellation.statusCode == 202) + var processedIterator = processed.makeAsyncIterator() + _ = try #require(await processedIterator.next()) + await barrier.release() + + for response in [await first.value, await second.value] { + #expect(response.statusCode == 200) + let data = try #require(response.bodyData) + let json = try #require(JSONSerialization.jsonObject(with: data) as? [String: Any]) + #expect(json["id"] as? String == "shared") + #expect(json["result"] != nil) + } + await server.stop() + } + @Test("Non-HTTP transport yields nil httpContext") func testNonHTTPTransportNilContext() async throws { actor Captured { var httpContext: HTTPRequest? var sawDispatch = false - func set(_ v: HTTPRequest?) { httpContext = v; sawDispatch = true } + func set(_ v: HTTPRequest?) { + httpContext = v + sawDispatch = true + } } let captured = Captured() @@ -1267,3 +1760,251 @@ struct ServerHandlerContextTests { await transport.disconnect() } } + +// MARK: - StatelessHTTPServerTransport Cancellation Tests (#255) +// +// Covers the fix for https://github.com/modelcontextprotocol/swift-sdk/issues/255: +// a request cancelled via `notifications/cancelled` gets no JSON-RPC response from +// `Server`. The transport completes the HTTP exchange with an empty SSE stream. +// +// Companion to fix/254-response-waiter-collision, which covers the separate +// "colliding JSON-RPC ids overwrite each other's waiter" gap (#254) — not an +// absence-of-response case, so not in scope here. +@Suite("StatelessHTTPServerTransport Cancellation Tests") +struct StatelessHTTPServerTransportCancellationTests { + + private func expectEmptySSE(_ response: HTTPResponse) async throws { + #expect(response.statusCode == 200) + #expect(response.headers["Content-Type"] == "text/event-stream") + guard case .stream(let stream, _) = response else { + Issue.record("Expected an SSE response") + return + } + var iterator = stream.makeAsyncIterator() + #expect(try await iterator.next() == nil) + } + + private struct SlowRequestHarness { + let transport: StatelessHTTPServerTransport + let server: Server + let inFlight: Task + /// Yields once the slow handler observes `CancellationError`. + let handlerCancelled: AsyncStream + } + + /// Starts a server whose `tools/call` handler sleeps long enough to be cancelled, + /// POSTs `requestBody`, and suspends until the handler has entered. + /// + /// Handler entry is deterministic evidence that the HTTP waiter is registered: + /// entry is causally downstream of the transport yielding the request body, and + /// the waiter is registered in the same synchronous actor turn as that yield, so + /// any cancellation POSTed after this returns is guaranteed to find the waiter. + private func startSlowRequest( + requestBody: Data, + accept: String = "application/json, text/event-stream" + ) async throws -> SlowRequestHarness { + let transport = makeStatelessTransport() + let server = Server(name: "TestServer", version: "1.0") + + let (handlerEntered, enteredContinuation) = AsyncStream.makeStream() + let (handlerCancelled, cancelledContinuation) = AsyncStream.makeStream() + + await server.withMethodHandler(CallTool.self) { _ in + enteredContinuation.yield(()) + do { + // `Task.sleep` is cancellation-aware: it throws `CancellationError` + // as soon as this handler's Task is cancelled. + try await Task.sleep(for: .seconds(5)) + } catch is CancellationError { + cancelledContinuation.yield(()) + throw CancellationError() + } + return CallTool.Result(content: [.text(text: "too slow", annotations: nil, _meta: nil)]) + } + + try await server.start(transport: transport) + + let inFlight = Task { + await transport.handleRequest(HTTPRequest( + method: "POST", + headers: ["Content-Type": "application/json", "Accept": accept], + body: requestBody + )) + } + + // Suspend until the handler has actually started (yields are buffered, so + // this cannot miss an entry that happened before we began iterating). + var entries = handlerEntered.makeAsyncIterator() + _ = await entries.next() + + return SlowRequestHarness( + transport: transport, + server: server, + inFlight: inFlight, + handlerCancelled: handlerCancelled + ) + } + + private func makeToolCallBody(id: Any) throws -> Data { + try JSONSerialization.data(withJSONObject: [ + "jsonrpc": "2.0", + "id": id, + "method": "tools/call", + "params": ["name": "slow-tool"] as [String: Any], + ]) + } + + @Test("Cancelled request completes its POST without a JSON-RPC response") + func testCancelledRequestCompletesWithoutResponse() async throws { + let harness = try await startSlowRequest(requestBody: makeToolCallBody(id: "slow-1")) + defer { Task { await harness.server.stop() } } + + let cancelResponse = await harness.transport.handleRequest( + makeStatelessPOSTRequest(body: makeCancelledNotificationBody(requestId: "slow-1")) + ) + #expect(cancelResponse.statusCode == 202) + + let completedInTime = await raceAgainstTimeout(.seconds(1)) { await harness.inFlight.value } + let response = try #require( + completedInTime, + "Original POST never completed after its request was cancelled" + ) + + try await expectEmptySSE(response) + + // The transport must still forward the notification to the Server: the slow + // handler observes CancellationError, not just the HTTP exchange completing. + let handlerSawCancellation = await raceAgainstTimeout(.seconds(1)) { + var cancellations = harness.handlerCancelled.makeAsyncIterator() + return await cancellations.next() != nil + } + #expect( + handlerSawCancellation == true, + "Slow handler never observed cancellation" + ) + } + + @Test("Cancellation with a reason still finishes without a response") + func testCancellationReasonDoesNotProduceResponse() async throws { + let harness = try await startSlowRequest(requestBody: makeToolCallBody(id: "slow-2")) + defer { Task { await harness.server.stop() } } + + let cancelResponse = await harness.transport.handleRequest( + makeStatelessPOSTRequest( + body: makeCancelledNotificationBody( + requestId: "slow-2", reason: "User requested cancellation")) + ) + #expect(cancelResponse.statusCode == 202) + + let completedInTime = await raceAgainstTimeout(.seconds(1)) { await harness.inFlight.value } + let response = try #require(completedInTime) + try await expectEmptySSE(response) + } + + @Test("Integer request id can be cancelled without a response") + func testIntegerRequestIDCancelled() async throws { + let harness = try await startSlowRequest(requestBody: makeToolCallBody(id: 7)) + defer { Task { await harness.server.stop() } } + + let cancelResponse = await harness.transport.handleRequest( + makeStatelessPOSTRequest(body: makeCancelledNotificationBody(requestId: 7)) + ) + #expect(cancelResponse.statusCode == 202) + + let completedInTime = await raceAgainstTimeout(.seconds(1)) { await harness.inFlight.value } + let response = try #require(completedInTime) + try await expectEmptySSE(response) + } + + @Test("JSON-only client gets a cancellation error with its original ID") + func testJSONOnlyCancellation() async throws { + let harness = try await startSlowRequest( + requestBody: makeToolCallBody(id: 7), accept: "application/json" + ) + defer { Task { await harness.server.stop() } } + + let cancelResponse = await harness.transport.handleRequest( + makeStatelessPOSTRequest(body: makeCancelledNotificationBody(requestId: 7)) + ) + #expect(cancelResponse.statusCode == 202) + let completed = try #require( + await raceAgainstTimeout(.seconds(1)) { await harness.inFlight.value } + ) + #expect(completed.statusCode == 200) + #expect(completed.headers[HTTPHeaderName.contentType] == ContentType.json) + let body = try #require(completed.bodyData) + let json = try #require(JSONSerialization.jsonObject(with: body) as? [String: Any]) + #expect(json["id"] as? Int == 7) + let error = try #require(json["error"] as? [String: Any]) + #expect(error["code"] as? Int == -32800) + } + + @Test("Cancellation for unknown or completed requests is ignored") + func testCancellationForUnknownOrCompletedRequestIsIgnored() async throws { + let transport = makeStatelessTransport() + let server = Server(name: "TestServer", version: "1.0") + + await server.withMethodHandler(CallTool.self) { _ in + CallTool.Result(content: [.text(text: "done", annotations: nil, _meta: nil)]) + } + + try await server.start(transport: transport) + defer { Task { await server.stop() } } + + // Complete a request normally. + let firstBody = try makeToolCallBody(id: "done-1") + let first = await raceAgainstTimeout(.seconds(1)) { + await transport.handleRequest(makeStatelessPOSTRequest(body: firstBody)) + } + let firstResponse = try #require(first) + #expect(firstResponse.statusCode == 200) + + // Cancelling the already-completed request is ignored ("Invalid cancellation + // notifications SHOULD be ignored": already completed requests). + let cancelCompleted = await transport.handleRequest( + makeStatelessPOSTRequest(body: makeCancelledNotificationBody(requestId: "done-1")) + ) + #expect(cancelCompleted.statusCode == 202) + + // Cancelling an id the transport has never seen is likewise ignored. + let cancelUnknown = await transport.handleRequest( + makeStatelessPOSTRequest(body: makeCancelledNotificationBody(requestId: "never-seen")) + ) + #expect(cancelUnknown.statusCode == 202) + + // The transport still serves requests normally afterwards. + let secondBody = try makeToolCallBody(id: "done-2") + let second = await raceAgainstTimeout(.seconds(1)) { + await transport.handleRequest(makeStatelessPOSTRequest(body: secondBody)) + } + let secondResponse = try #require(second) + #expect(secondResponse.statusCode == 200) + } + + @Test("Malformed cancellation without a requestId is ignored") + func testMalformedCancellationIsIgnored() async throws { + let harness = try await startSlowRequest(requestBody: makeToolCallBody(id: "slow-3")) + defer { Task { await harness.server.stop() } } + + // A cancellation with no requestId is ignored (202) and resumes nothing. + let malformed = try JSONSerialization.data(withJSONObject: [ + "jsonrpc": "2.0", + "method": "notifications/cancelled", + "params": [String: Any](), + ]) + let malformedResponse = await harness.transport.handleRequest( + makeStatelessPOSTRequest(body: malformed) + ) + #expect(malformedResponse.statusCode == 202) + + // The original request is still in flight; a valid cancellation completes it. + let cancelResponse = await harness.transport.handleRequest( + makeStatelessPOSTRequest(body: makeCancelledNotificationBody(requestId: "slow-3")) + ) + #expect(cancelResponse.statusCode == 202) + + let completedInTime = await raceAgainstTimeout(.seconds(1)) { await harness.inFlight.value } + let response = try #require(completedInTime) + try await expectEmptySSE(response) + } +} diff --git a/Tests/MCPTests/Helpers/MockTransport.swift b/Tests/MCPTests/Helpers/MockTransport.swift index 5979e9ad..0c8505d5 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)