From 432d1d0451d4e866f16947e2bd7fcd9060e546c2 Mon Sep 17 00:00:00 2001 From: King Star Date: Sun, 26 Jul 2026 22:20:19 +0800 Subject: [PATCH 1/7] fix: isolate stateless HTTP request exchanges Signed-off-by: King Star --- .../HTTPServer/HTTPServerTypes.swift | 8 + .../StatelessHTTPServerTransport.swift | 167 +++++++++- Sources/MCP/Server/Server.swift | 15 +- Tests/MCPTests/HTTPServerTransportTests.swift | 315 +++++++++++++++++- 4 files changed, 477 insertions(+), 28 deletions(-) diff --git a/Sources/MCP/Base/Transports/HTTPServer/HTTPServerTypes.swift b/Sources/MCP/Base/Transports/HTTPServer/HTTPServerTypes.swift index bd56191c..7f84af90 100644 --- a/Sources/MCP/Base/Transports/HTTPServer/HTTPServerTypes.swift +++ b/Sources/MCP/Base/Transports/HTTPServer/HTTPServerTypes.swift @@ -225,6 +225,14 @@ 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? +} + // 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..6457e72e 100644 --- a/Sources/MCP/Base/Transports/HTTPServer/StatelessHTTPServerTransport.swift +++ b/Sources/MCP/Base/Transports/HTTPServer/StatelessHTTPServerTransport.swift @@ -29,7 +29,9 @@ import Logging /// - Session management is handled externally or not needed /// /// For full streaming and session support, use ``StatefulHTTPServerTransport`` instead. -public actor StatelessHTTPServerTransport: Transport, HTTPContextProviding { +public actor StatelessHTTPServerTransport: + Transport, HTTPContextProviding, RoutedRequestIDProviding +{ public nonisolated let logger: Logger // MARK: - Dependencies @@ -48,15 +50,23 @@ 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 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 +127,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: response) + } catch { + waiter.continuation.resume(throwing: error) + throw error + } case .notification(let method): logger.debug( @@ -221,32 +237,146 @@ 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 do { responseData = try await withCheckedThrowingContinuation { continuation in - responseWaiters[requestID] = continuation + responseWaiters[exchangeID] = ResponseWaiter( + originalID: originalID, + continuation: continuation + ) } } catch { - httpRequestContexts.removeValue(forKey: requestID) + removeHTTPContext(exchangeID: exchangeID, requestID: originalID) return .error( statusCode: 500, .internalError("Error processing request: \(error.localizedDescription)") ) } - httpRequestContexts.removeValue(forKey: requestID) + removeHTTPContext(exchangeID: exchangeID, requestID: originalID) return .data(responseData, 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) { + guard var json = try JSONSerialization.jsonObject(with: body) as? [String: Any] else { + throw MCPError.parseError("Invalid JSON-RPC request") + } + + guard let originalID = requestID(from: json["id"]) else { + throw MCPError.parseError("Invalid JSON-RPC request id") + } + + json["id"] = exchangeID + let routedBody = try JSONSerialization.data( + withJSONObject: json, + options: [.sortedKeys, .withoutEscapingSlashes] + ) + return (routedBody, originalID) + } + + private func requestID(from value: Any?) -> ID? { + if let stringID = value as? String { + return .string(stringID) + } + if let numberID = value as? Int { + return .number(numberID) + } + return nil + } + + private func restoringResponseID(in data: Data, to originalID: ID) throws -> Data { + guard var json = try JSONSerialization.jsonObject(with: data) as? [String: Any] else { + throw MCPError.parseError("Invalid JSON-RPC response") + } + + switch originalID { + case .string(let value): + json["id"] = value + case .number(let value): + json["id"] = value + } + + return try JSONSerialization.data( + withJSONObject: json, + options: [.sortedKeys, .withoutEscapingSlashes] + ) + } + + private func registerHTTPContext( + _ request: HTTPRequest, + exchangeID: String, + requestID: ID + ) { + httpRequestContexts[exchangeID] = request + exchangeIDsByRequestID[requestID, default: []].append(exchangeID) + } + + private func removeHTTPContext(exchangeID: String, requestID: ID) { + httpRequestContexts.removeValue(forKey: exchangeID) + guard var exchangeIDs = exchangeIDsByRequestID[requestID] else { return } + exchangeIDs.removeAll { $0 == exchangeID } + if exchangeIDs.isEmpty { + exchangeIDsByRequestID.removeValue(forKey: requestID) + } else { + exchangeIDsByRequestID[requestID] = exchangeIDs + } + } + // MARK: - HTTPContextProviding public func httpRequestContext(for id: ID) -> HTTPRequest? { - httpRequestContexts[id.description] + if case .string(let exchangeID) = id, + let request = httpRequestContexts[exchangeID] + { + return request + } + guard let exchangeID = exchangeIDsByRequestID[id]?.last 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 +388,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..cecbf009 100644 --- a/Sources/MCP/Server/Server.swift +++ b/Sources/MCP/Server/Server.swift @@ -773,7 +773,9 @@ 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) // Create a task to handle the request with cancellation support. // Set currentHandlerContext as a task local so handlers see it. @@ -984,7 +986,9 @@ public actor Server { } // Cancel the pending request task if it exists and remove from tracking - if let task = await self.removePendingRequest(id: requestId) { + if let pendingRequestID = await self.routedRequestID(for: requestId), + let task = await self.removePendingRequest(id: pendingRequestID) + { task.cancel() await self.logger?.debug( "Cancelled request", @@ -1001,6 +1005,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/HTTPServerTransportTests.swift b/Tests/MCPTests/HTTPServerTransportTests.swift index 8e1c89b4..b97516c1 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", @@ -158,6 +167,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( @@ -895,7 +923,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 +942,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 +1162,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") @@ -1146,11 +1390,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 +1408,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,6 +1483,49 @@ 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("Non-HTTP transport yields nil httpContext") func testNonHTTPTransportNilContext() async throws { actor Captured { From f853eed575982376ed4a775f9d1b481dd0ff4255 Mon Sep 17 00:00:00 2001 From: Ian Gordon Date: Mon, 20 Jul 2026 14:26:36 -0400 Subject: [PATCH 2/7] fix: complete the HTTP exchange when a request is cancelled in StatelessHTTPServerTransport A request cancelled via notifications/cancelled correctly receives no JSON-RPC response from Server, but nothing resumed the transport's HTTP response waiter, so the original POST hung until transport termination. The transport now detects notifications/cancelled, forwards it to the Server unchanged (its existing cancellation path still cancels the handler task), and resumes the matching waiter with a synthesized JSON-RPC error response (code -32002, "Request cancelled"), satisfying the Streamable HTTP requirement that a request POST receive one JSON object. Unknown, already-completed, and requestId-less cancellations are ignored per spec. Fixes #255 Co-Authored-By: Claude Fable 5 --- .../StatelessHTTPServerTransport.swift | 38 ++- Tests/MCPTests/HTTPServerTransportTests.swift | 287 ++++++++++++++++++ 2 files changed, 324 insertions(+), 1 deletion(-) diff --git a/Sources/MCP/Base/Transports/HTTPServer/StatelessHTTPServerTransport.swift b/Sources/MCP/Base/Transports/HTTPServer/StatelessHTTPServerTransport.swift index 6457e72e..d24406ad 100644 --- a/Sources/MCP/Base/Transports/HTTPServer/StatelessHTTPServerTransport.swift +++ b/Sources/MCP/Base/Transports/HTTPServer/StatelessHTTPServerTransport.swift @@ -222,7 +222,15 @@ public actor StatelessHTTPServerTransport: // Handle by message type switch messageKind { - case .notification, .response: + case .notification(let method): + // Yield to server and return 202 Accepted + incomingContinuation.yield(body) + if method == CancelledNotification.name { + completeExchangeForCancelledRequest(body) + } + return .accepted() + + case .response: // Yield to server and return 202 Accepted incomingContinuation.yield(body) return .accepted() @@ -350,6 +358,34 @@ public actor StatelessHTTPServerTransport: } } + private func completeExchangeForCancelledRequest(_ body: Data) { + guard let notification = try? JSONDecoder().decode( + Message.self, from: body), + let requestID = notification.params.requestId, + let exchangeIDs = exchangeIDsByRequestID[requestID], + exchangeIDs.count == 1, + let exchangeID = exchangeIDs.first, + let waiter = responseWaiters.removeValue(forKey: exchangeID) + else { return } + + removeHTTPContext(exchangeID: exchangeID, requestID: waiter.originalID) + var message = "Request cancelled" + if let reason = notification.params.reason { + message += ": \(reason)" + } + let response = AnyMethod.response( + id: waiter.originalID, + error: .serverError(code: -32002, message: message) + ) + do { + let encoder = JSONEncoder() + encoder.outputFormatting = [.sortedKeys, .withoutEscapingSlashes] + waiter.continuation.resume(returning: try encoder.encode(response)) + } catch { + waiter.continuation.resume(throwing: error) + } + } + // MARK: - HTTPContextProviding public func httpRequestContext(for id: ID) -> HTTPRequest? { diff --git a/Tests/MCPTests/HTTPServerTransportTests.swift b/Tests/MCPTests/HTTPServerTransportTests.swift index b97516c1..301d6a6d 100644 --- a/Tests/MCPTests/HTTPServerTransportTests.swift +++ b/Tests/MCPTests/HTTPServerTransportTests.swift @@ -104,6 +104,18 @@ private func makeStatelessPOSTRequest(body: Data) -> HTTPRequest { ) } +/// `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 { @@ -148,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, @@ -1560,3 +1609,241 @@ 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` (per the cancellation spec), so the transport itself must complete the +// hung HTTP exchange with a synthesized "Request cancelled" JSON-RPC error to +// satisfy the Streamable HTTP requirement that a request POST receive one JSON +// object or an 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 { + + /// The implementation-specific JSON-RPC error code the transport synthesizes + /// for cancelled requests. + private static let requestCancelledCode = -32002 + + 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) 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(makeStatelessPOSTRequest(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 with a Request cancelled error") + func testCancelledRequestCompletesWithError() 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" + ) + + #expect(response.statusCode == 200) + #expect(response.headers["Content-Type"] == "application/json") + + let body = try #require(response.bodyData) + let json = try #require(try JSONSerialization.jsonObject(with: body) as? [String: Any]) + #expect(json["jsonrpc"] as? String == "2.0") + #expect(json["id"] as? String == "slow-1") + #expect(json["result"] == nil) + + let error = try #require(json["error"] as? [String: Any]) + #expect(error["code"] as? Int == Self.requestCancelledCode) + let message = try #require(error["message"] as? String) + #expect(message.contains("Request cancelled")) + + // 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 reason is included in the synthesized error message") + func testCancellationReasonIncludedInErrorMessage() 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) + let body = try #require(response.bodyData) + let json = try #require(try JSONSerialization.jsonObject(with: body) as? [String: Any]) + let error = try #require(json["error"] as? [String: Any]) + let message = try #require(error["message"] as? String) + #expect(message.contains("User requested cancellation")) + } + + @Test("Integer request id is preserved in the synthesized error response") + func testIntegerRequestIDPreserved() 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) + let body = try #require(response.bodyData) + let json = try #require(try JSONSerialization.jsonObject(with: body) as? [String: Any]) + // The id must round-trip as a JSON number, not the string "7". + #expect(json["id"] as? Int == 7) + let error = try #require(json["error"] as? [String: Any]) + #expect(error["code"] as? Int == Self.requestCancelledCode) + } + + @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) + let body = try #require(response.bodyData) + let json = try #require(try JSONSerialization.jsonObject(with: body) as? [String: Any]) + let error = try #require(json["error"] as? [String: Any]) + #expect(error["code"] as? Int == Self.requestCancelledCode) + } +} From b0a9f33ae95273bb3c803d2df43a9270ce7b9e11 Mon Sep 17 00:00:00 2001 From: EhsanAzish80 Date: Sat, 26 Sep 2026 11:04:59 +0300 Subject: [PATCH 3/7] fix: preserve HTTP request identity during cancellation --- .../HTTPServer/HTTPServerTypes.swift | 4 + .../StatelessHTTPServerTransport.swift | 47 ++-- Sources/MCP/Server/Server.swift | 204 +++++++++++++++--- Tests/MCPTests/CancellationTests.swift | 168 ++++++++++++++- Tests/MCPTests/HTTPServerTransportTests.swift | 112 +++++++++- Tests/MCPTests/Helpers/MockTransport.swift | 32 ++- 6 files changed, 498 insertions(+), 69 deletions(-) diff --git a/Sources/MCP/Base/Transports/HTTPServer/HTTPServerTypes.swift b/Sources/MCP/Base/Transports/HTTPServer/HTTPServerTypes.swift index 7f84af90..7e14fa70 100644 --- a/Sources/MCP/Base/Transports/HTTPServer/HTTPServerTypes.swift +++ b/Sources/MCP/Base/Transports/HTTPServer/HTTPServerTypes.swift @@ -233,6 +233,10 @@ package protocol RoutedRequestIDProviding: Sendable { 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 d24406ad..c69c03c6 100644 --- a/Sources/MCP/Base/Transports/HTTPServer/StatelessHTTPServerTransport.swift +++ b/Sources/MCP/Base/Transports/HTTPServer/StatelessHTTPServerTransport.swift @@ -30,7 +30,7 @@ import Logging /// /// For full streaming and session support, use ``StatefulHTTPServerTransport`` instead. public actor StatelessHTTPServerTransport: - Transport, HTTPContextProviding, RoutedRequestIDProviding + Transport, HTTPContextProviding, RoutedRequestIDProviding, CancelledHTTPExchangeCompleting { public nonisolated let logger: Logger @@ -222,12 +222,9 @@ public actor StatelessHTTPServerTransport: // Handle by message type switch messageKind { - case .notification(let method): + case .notification: // Yield to server and return 202 Accepted incomingContinuation.yield(body) - if method == CancelledNotification.name { - completeExchangeForCancelledRequest(body) - } return .accepted() case .response: @@ -264,11 +261,15 @@ public actor StatelessHTTPServerTransport: // Wait for the server to process and send a response let responseData: Data do { - responseData = try await withCheckedThrowingContinuation { continuation in - responseWaiters[exchangeID] = ResponseWaiter( - originalID: originalID, - continuation: continuation - ) + responseData = try await withTaskCancellationHandler { + try await withCheckedThrowingContinuation { continuation in + responseWaiters[exchangeID] = ResponseWaiter( + originalID: originalID, + continuation: continuation + ) + } + } onCancel: { + Task { await self.cancelWaitingExchange(exchangeID) } } } catch { removeHTTPContext(exchangeID: exchangeID, requestID: originalID) @@ -358,24 +359,23 @@ public actor StatelessHTTPServerTransport: } } - private func completeExchangeForCancelledRequest(_ body: Data) { - guard let notification = try? JSONDecoder().decode( - Message.self, from: body), - let requestID = notification.params.requestId, - let exchangeIDs = exchangeIDsByRequestID[requestID], - exchangeIDs.count == 1, - let exchangeID = exchangeIDs.first, + private func cancelWaitingExchange(_ exchangeID: String) { + guard let waiter = responseWaiters.removeValue(forKey: exchangeID) else { return } + removeHTTPContext(exchangeID: exchangeID, requestID: waiter.originalID) + waiter.continuation.resume(throwing: CancellationError()) + } + + package func completeCancelledExchange(for routedID: ID, reason: String?) { + guard case .string(let exchangeID) = routedID, let waiter = responseWaiters.removeValue(forKey: exchangeID) else { return } removeHTTPContext(exchangeID: exchangeID, requestID: waiter.originalID) var message = "Request cancelled" - if let reason = notification.params.reason { - message += ": \(reason)" - } + if let reason { message += ": \(reason)" } let response = AnyMethod.response( id: waiter.originalID, - error: .serverError(code: -32002, message: message) + error: .serverError(code: -32800, message: message) ) do { let encoder = JSONEncoder() @@ -394,7 +394,10 @@ public actor StatelessHTTPServerTransport: { return request } - guard let exchangeID = exchangeIDsByRequestID[id]?.last else { + guard let exchangeIDs = exchangeIDsByRequestID[id], + exchangeIDs.count == 1, + let exchangeID = exchangeIDs.first + else { return nil } return httpRequestContexts[exchangeID] diff --git a/Sources/MCP/Server/Server.swift b/Sources/MCP/Server/Server.swift index cecbf009..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,9 +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 handlerRequestID = await (connection as? any RoutedRequestIDProviding)? + 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. @@ -804,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) @@ -830,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)"]) @@ -958,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)", @@ -985,18 +1101,38 @@ public actor Server { return } - // Cancel the pending request task if it exists and remove from tracking - if let pendingRequestID = await self.routedRequestID(for: requestId), - let task = await self.removePendingRequest(id: pendingRequestID) - { - 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)"] 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 301d6a6d..a4736bf9 100644 --- a/Tests/MCPTests/HTTPServerTransportTests.swift +++ b/Tests/MCPTests/HTTPServerTransportTests.swift @@ -255,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", @@ -1344,6 +1348,33 @@ struct StatelessHTTPServerTransportTests { // Should return error (500 or similar) since waiter was cancelled #expect(response.statusCode == 500) } + + @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() + _ = await handleTask.value + #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() + } } // MARK: - HTTPContextProviding / Server.currentHandlerContext @@ -1575,12 +1606,83 @@ struct ServerHandlerContextTests { 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() @@ -1627,7 +1729,7 @@ struct StatelessHTTPServerTransportCancellationTests { /// The implementation-specific JSON-RPC error code the transport synthesizes /// for cancelled requests. - private static let requestCancelledCode = -32002 + private static let requestCancelledCode = -32800 private struct SlowRequestHarness { let transport: StatelessHTTPServerTransport 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) From 8469c92b3dd623b9222f529e3082d4ab832114dc Mon Sep 17 00:00:00 2001 From: EhsanAzish80 Date: Sat, 26 Sep 2026 23:24:45 +0300 Subject: [PATCH 4/7] test: require prompt HTTP waiter cleanup on caller cancellation --- Tests/MCPTests/HTTPServerTransportTests.swift | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/Tests/MCPTests/HTTPServerTransportTests.swift b/Tests/MCPTests/HTTPServerTransportTests.swift index a4736bf9..b1e591f6 100644 --- a/Tests/MCPTests/HTTPServerTransportTests.swift +++ b/Tests/MCPTests/HTTPServerTransportTests.swift @@ -1368,7 +1368,9 @@ struct StatelessHTTPServerTransportTests { let exchangeID = try #require(routedJSON["id"] as? String) handleTask.cancel() - _ = await handleTask.value + 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) From 4ae4b3e68bbdba7ab579a897947caff07011b460 Mon Sep 17 00:00:00 2001 From: EhsanAzish80 Date: Sun, 27 Sep 2026 08:55:18 +0300 Subject: [PATCH 5/7] fix: report transport shutdown separately from HTTP cancellation --- .../StatelessHTTPServerTransport.swift | 6 ++++ Tests/MCPTests/HTTPServerTransportTests.swift | 29 +++++++++++++------ 2 files changed, 26 insertions(+), 9 deletions(-) diff --git a/Sources/MCP/Base/Transports/HTTPServer/StatelessHTTPServerTransport.swift b/Sources/MCP/Base/Transports/HTTPServer/StatelessHTTPServerTransport.swift index c69c03c6..284bf2e2 100644 --- a/Sources/MCP/Base/Transports/HTTPServer/StatelessHTTPServerTransport.swift +++ b/Sources/MCP/Base/Transports/HTTPServer/StatelessHTTPServerTransport.swift @@ -169,6 +169,9 @@ public actor StatelessHTTPServerTransport: /// - **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( @@ -271,6 +274,9 @@ public actor StatelessHTTPServerTransport: } onCancel: { Task { await self.cancelWaitingExchange(exchangeID) } } + } catch MCPError.connectionClosed { + removeHTTPContext(exchangeID: exchangeID, requestID: originalID) + return .error(statusCode: 503, .connectionClosed) } catch { removeHTTPContext(exchangeID: exchangeID, requestID: originalID) return .error( diff --git a/Tests/MCPTests/HTTPServerTransportTests.swift b/Tests/MCPTests/HTTPServerTransportTests.swift index b1e591f6..fca442e3 100644 --- a/Tests/MCPTests/HTTPServerTransportTests.swift +++ b/Tests/MCPTests/HTTPServerTransportTests.swift @@ -1325,28 +1325,39 @@ 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") From adcb3baa92167ac1c94f5b4769a217b02f543e43 Mon Sep 17 00:00:00 2001 From: EhsanAzish80 Date: Sat, 3 Oct 2026 08:59:19 +0300 Subject: [PATCH 6/7] Complete cancelled HTTP requests without JSON-RPC replies --- .../StatelessHTTPServerTransport.swift | 39 +++++------ Tests/MCPTests/HTTPServerTransportTests.swift | 64 +++++++------------ 2 files changed, 43 insertions(+), 60 deletions(-) diff --git a/Sources/MCP/Base/Transports/HTTPServer/StatelessHTTPServerTransport.swift b/Sources/MCP/Base/Transports/HTTPServer/StatelessHTTPServerTransport.swift index 284bf2e2..f170e852 100644 --- a/Sources/MCP/Base/Transports/HTTPServer/StatelessHTTPServerTransport.swift +++ b/Sources/MCP/Base/Transports/HTTPServer/StatelessHTTPServerTransport.swift @@ -29,6 +29,7 @@ import Logging /// - Session management is handled externally or not needed /// /// For full streaming and session support, use ``StatefulHTTPServerTransport`` instead. +/// A cancelled request ends with an empty SSE response so no JSON-RPC reply is sent. public actor StatelessHTTPServerTransport: Transport, HTTPContextProviding, RoutedRequestIDProviding, CancelledHTTPExchangeCompleting { @@ -50,9 +51,14 @@ public actor StatelessHTTPServerTransport: // MARK: - Response waiters + private enum ResponseOutcome { + case data(Data) + case cancelled + } + private struct ResponseWaiter { let originalID: ID - let continuation: CheckedContinuation + let continuation: CheckedContinuation } /// Maps a transport-private exchange ID to the matching HTTP response waiter. @@ -136,7 +142,7 @@ public actor StatelessHTTPServerTransport: } do { let response = try restoringResponseID(in: data, to: waiter.originalID) - waiter.continuation.resume(returning: response) + waiter.continuation.resume(returning: .data(response)) } catch { waiter.continuation.resume(throwing: error) throw error @@ -262,9 +268,9 @@ public actor StatelessHTTPServerTransport: incomingContinuation.yield(routedBody) // Wait for the server to process and send a response - let responseData: Data + let outcome: ResponseOutcome do { - responseData = try await withTaskCancellationHandler { + outcome = try await withTaskCancellationHandler { try await withCheckedThrowingContinuation { continuation in responseWaiters[exchangeID] = ResponseWaiter( originalID: originalID, @@ -286,7 +292,15 @@ public actor StatelessHTTPServerTransport: } removeHTTPContext(exchangeID: exchangeID, requestID: originalID) - return .data(responseData, headers: [HTTPHeaderName.contentType: ContentType.json]) + switch outcome { + case .data(let responseData): + return .data(responseData, headers: [HTTPHeaderName.contentType: ContentType.json]) + case .cancelled: + return .stream( + AsyncThrowingStream { $0.finish() }, + headers: [HTTPHeaderName.contentType: ContentType.sse] + ) + } } private func makeExchangeID(excluding requestID: String) -> String { @@ -377,19 +391,8 @@ public actor StatelessHTTPServerTransport: else { return } removeHTTPContext(exchangeID: exchangeID, requestID: waiter.originalID) - var message = "Request cancelled" - if let reason { message += ": \(reason)" } - let response = AnyMethod.response( - id: waiter.originalID, - error: .serverError(code: -32800, message: message) - ) - do { - let encoder = JSONEncoder() - encoder.outputFormatting = [.sortedKeys, .withoutEscapingSlashes] - waiter.continuation.resume(returning: try encoder.encode(response)) - } catch { - waiter.continuation.resume(throwing: error) - } + logger.debug("Completed cancelled HTTP exchange", metadata: ["reason": "\(reason ?? "none")"]) + waiter.continuation.resume(returning: .cancelled) } // MARK: - HTTPContextProviding diff --git a/Tests/MCPTests/HTTPServerTransportTests.swift b/Tests/MCPTests/HTTPServerTransportTests.swift index fca442e3..4e28ec82 100644 --- a/Tests/MCPTests/HTTPServerTransportTests.swift +++ b/Tests/MCPTests/HTTPServerTransportTests.swift @@ -98,7 +98,7 @@ private func makeStatelessPOSTRequest(body: Data) -> HTTPRequest { method: "POST", headers: [ "Content-Type": "application/json", - "Accept": "application/json", + "Accept": "application/json, text/event-stream", ], body: body ) @@ -1729,10 +1729,7 @@ struct ServerHandlerContextTests { // // 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` (per the cancellation spec), so the transport itself must complete the -// hung HTTP exchange with a synthesized "Request cancelled" JSON-RPC error to -// satisfy the Streamable HTTP requirement that a request POST receive one JSON -// object or an SSE stream. +// `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 @@ -1740,9 +1737,16 @@ struct ServerHandlerContextTests { @Suite("StatelessHTTPServerTransport Cancellation Tests") struct StatelessHTTPServerTransportCancellationTests { - /// The implementation-specific JSON-RPC error code the transport synthesizes - /// for cancelled requests. - private static let requestCancelledCode = -32800 + 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 @@ -1807,8 +1811,8 @@ struct StatelessHTTPServerTransportCancellationTests { ]) } - @Test("Cancelled request completes its POST with a Request cancelled error") - func testCancelledRequestCompletesWithError() async throws { + @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() } } @@ -1823,19 +1827,7 @@ struct StatelessHTTPServerTransportCancellationTests { "Original POST never completed after its request was cancelled" ) - #expect(response.statusCode == 200) - #expect(response.headers["Content-Type"] == "application/json") - - let body = try #require(response.bodyData) - let json = try #require(try JSONSerialization.jsonObject(with: body) as? [String: Any]) - #expect(json["jsonrpc"] as? String == "2.0") - #expect(json["id"] as? String == "slow-1") - #expect(json["result"] == nil) - - let error = try #require(json["error"] as? [String: Any]) - #expect(error["code"] as? Int == Self.requestCancelledCode) - let message = try #require(error["message"] as? String) - #expect(message.contains("Request 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. @@ -1849,8 +1841,8 @@ struct StatelessHTTPServerTransportCancellationTests { ) } - @Test("Cancellation reason is included in the synthesized error message") - func testCancellationReasonIncludedInErrorMessage() async throws { + @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() } } @@ -1863,15 +1855,11 @@ struct StatelessHTTPServerTransportCancellationTests { let completedInTime = await raceAgainstTimeout(.seconds(1)) { await harness.inFlight.value } let response = try #require(completedInTime) - let body = try #require(response.bodyData) - let json = try #require(try JSONSerialization.jsonObject(with: body) as? [String: Any]) - let error = try #require(json["error"] as? [String: Any]) - let message = try #require(error["message"] as? String) - #expect(message.contains("User requested cancellation")) + try await expectEmptySSE(response) } - @Test("Integer request id is preserved in the synthesized error response") - func testIntegerRequestIDPreserved() async throws { + @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() } } @@ -1882,12 +1870,7 @@ struct StatelessHTTPServerTransportCancellationTests { let completedInTime = await raceAgainstTimeout(.seconds(1)) { await harness.inFlight.value } let response = try #require(completedInTime) - let body = try #require(response.bodyData) - let json = try #require(try JSONSerialization.jsonObject(with: body) as? [String: Any]) - // The id must round-trip as a JSON number, not the string "7". - #expect(json["id"] as? Int == 7) - let error = try #require(json["error"] as? [String: Any]) - #expect(error["code"] as? Int == Self.requestCancelledCode) + try await expectEmptySSE(response) } @Test("Cancellation for unknown or completed requests is ignored") @@ -1956,9 +1939,6 @@ struct StatelessHTTPServerTransportCancellationTests { let completedInTime = await raceAgainstTimeout(.seconds(1)) { await harness.inFlight.value } let response = try #require(completedInTime) - let body = try #require(response.bodyData) - let json = try #require(try JSONSerialization.jsonObject(with: body) as? [String: Any]) - let error = try #require(json["error"] as? [String: Any]) - #expect(error["code"] as? Int == Self.requestCancelledCode) + try await expectEmptySSE(response) } } From 206bc64c04ebec4b2af0a7f82ca0c5fbd91bdc19 Mon Sep 17 00:00:00 2001 From: EhsanAzish80 Date: Sat, 24 Oct 2026 21:08:56 +0300 Subject: [PATCH 7/7] Preserve JSON payloads when routing HTTP request IDs --- .../StatelessHTTPServerTransport.swift | 143 +++++++++++++----- Tests/MCPTests/HTTPServerTransportTests.swift | 70 ++++++++- 2 files changed, 175 insertions(+), 38 deletions(-) diff --git a/Sources/MCP/Base/Transports/HTTPServer/StatelessHTTPServerTransport.swift b/Sources/MCP/Base/Transports/HTTPServer/StatelessHTTPServerTransport.swift index f170e852..d31a0db6 100644 --- a/Sources/MCP/Base/Transports/HTTPServer/StatelessHTTPServerTransport.swift +++ b/Sources/MCP/Base/Transports/HTTPServer/StatelessHTTPServerTransport.swift @@ -215,6 +215,12 @@ public actor StatelessHTTPServerTransport: .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( @@ -296,10 +302,20 @@ public actor StatelessHTTPServerTransport: case .data(let responseData): return .data(responseData, headers: [HTTPHeaderName.contentType: ContentType.json]) case .cancelled: - return .stream( - AsyncThrowingStream { $0.finish() }, - headers: [HTTPHeaderName.contentType: ContentType.sse] + 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]) } } @@ -315,48 +331,103 @@ public actor StatelessHTTPServerTransport: } private func routingRequest(_ body: Data, as exchangeID: String) throws -> (Data, ID) { - guard var json = try JSONSerialization.jsonObject(with: body) as? [String: Any] else { - throw MCPError.parseError("Invalid JSON-RPC request") - } - - guard let originalID = requestID(from: json["id"]) else { - throw MCPError.parseError("Invalid JSON-RPC request id") - } + try replacingTopLevelID(in: body, with: .string(exchangeID)) + } - json["id"] = exchangeID - let routedBody = try JSONSerialization.data( - withJSONObject: json, - options: [.sortedKeys, .withoutEscapingSlashes] - ) - return (routedBody, originalID) + private func restoringResponseID(in data: Data, to originalID: ID) throws -> Data { + try replacingTopLevelID(in: data, with: originalID).0 } - private func requestID(from value: Any?) -> ID? { - if let stringID = value as? String { - return .string(stringID) + /// 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") } - if let numberID = value as? Int { - return .number(numberID) - } - return nil - } - private func restoringResponseID(in data: Data, to originalID: ID) throws -> Data { - guard var json = try JSONSerialization.jsonObject(with: data) as? [String: Any] else { - throw MCPError.parseError("Invalid JSON-RPC response") + 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") } - switch originalID { - case .string(let value): - json["id"] = value - case .number(let value): - json["id"] = value + 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.. SlowRequestHarness { + 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") @@ -1786,7 +1825,11 @@ struct StatelessHTTPServerTransportCancellationTests { try await server.start(transport: transport) let inFlight = Task { - await transport.handleRequest(makeStatelessPOSTRequest(body: requestBody)) + 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 @@ -1873,6 +1916,29 @@ struct StatelessHTTPServerTransportCancellationTests { 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()