diff --git a/README.md b/README.md index 1bf1670a..591a8976 100644 --- a/README.md +++ b/README.md @@ -42,6 +42,7 @@ of the MCP specification. - [Progress Tracking](#progress-tracking-1) - [Initialize Hook](#initialize-hook) - [HTTP Request Context in Handlers](#http-request-context-in-handlers) + - [Exact JSON in Method Handlers](#exact-json-in-method-handlers) - [Graceful Shutdown](#graceful-shutdown) - [Transports](#transports) - [Authentication](#authentication) @@ -1214,6 +1215,55 @@ Task.detached { await doWork(with: ctx?.httpContext) } Custom HTTP transports can opt in by conforming to `HTTPContextProviding` and returning the `HTTPRequest` for a given JSON-RPC id while it is in flight. +### Exact JSON in Method Handlers + +`Value` intentionally provides convenient semantic JSON values, including automatic data-URL +decoding and dictionary-backed objects. If a handler must instead inspect source member order, +duplicate names, Unicode-scalar-exact keys, or a data URL as a string, register the additive +raw-aware method-handler overload: + +```swift +await server.withMethodHandler(CallTool.self) { request, rawContext in + // `request` remains the usual validated, typed request. + guard let parameters = rawContext.uniqueParameters?.objectValue, + let arguments = parameters.uniqueValue(forExactKey: "arguments")?.objectValue + else { + throw MCPError.invalidParams("Expected one object-valued params member") + } + + // Every duplicate is retained in source order. Lookup compares Unicode + // scalar sequences exactly and never normalizes object names. + let payloads = arguments.values(forExactKey: "payload") + if case .string(let dataURL)? = payloads.first { + // A string remains a string even when it is a valid data URL. + print(dataURL) + } + + return .init(content: [.text("Handled \(request.params.name)")], isError: false) +} +``` + +`RawJSONValue` uses a dedicated parser and serializer rather than `Codable`, because keyed +decoding cannot retain duplicate object members. Raw-aware request parsing and encoding are +bounded by `RawJSONLimits`; pass custom limits when constructing the server if the defaults are +not suitable. These limits do not affect legacy handlers or inbound responses: + +```swift +let server = Server( + name: "MyModelServer", + version: "1.0.0", + rawJSONLimits: .init( + maximumDocumentBytes: 8 * 1024 * 1024, + maximumStringBytes: 6 * 1024 * 1024, + maximumDepth: 64, + maximumContainerEntries: 50_000 + ) +) +``` + +Inbound JSON strings always decode as `RawJSONValue.string`. The `data` case is reserved for +programmatically constructed raw values and is size-checked before base64 storage is allocated. + ### Graceful Shutdown We recommend using diff --git a/Sources/MCP/Base/Messages.swift b/Sources/MCP/Base/Messages.swift index e53bb729..3cb69925 100644 --- a/Sources/MCP/Base/Messages.swift +++ b/Sources/MCP/Base/Messages.swift @@ -154,6 +154,35 @@ extension Request { /// A type-erased request for request/response handling typealias AnyRequest = Request +/// The exact JSON representation retained alongside a typed inbound request. +/// +/// Register a raw-aware server method handler to receive this context when +/// object member order, duplicate keys, or data-URL strings must be preserved. +/// The ordinary typed request remains the source of the method's validated +/// parameters. +public struct RawRequestContext: Hashable, Sendable { + /// The complete inbound JSON-RPC request object. + public let request: ExactJSONObject + + /// Every exact `params` member in source order. + /// + /// Valid JSON-RPC requests normally contain at most one. Returning an + /// array keeps duplicate-member handling explicit instead of silently + /// selecting a first or last value. + public var parameters: [RawJSONValue] { + request.values(forExactKey: "params") + } + + /// The exact `params` value when it occurs once, otherwise `nil`. + public var uniqueParameters: RawJSONValue? { + request.uniqueValue(forExactKey: "params") + } + + public init(request: ExactJSONObject) { + self.request = request + } +} + extension AnyRequest { init(_ request: Request) throws { let encoder = JSONEncoder() @@ -166,21 +195,47 @@ extension AnyRequest { /// A box for request handlers that can be type-erased class RequestHandlerBox: @unchecked Sendable { - func callAsFunction(_ request: AnyRequest) async throws -> AnyResponse { + var requiresRawContext: Bool { false } + + func callAsFunction( + _ request: AnyRequest, + rawContext: RawRequestContext? = nil + ) async throws -> AnyResponse { fatalError("Must override") } } /// A typed request handler that can be used to handle requests of a specific type final class TypedRequestHandler: RequestHandlerBox, @unchecked Sendable { - private let _handle: @Sendable (Request) async throws -> Response + private let _handle: @Sendable (Request, RawRequestContext?) async throws -> Response + private let isRawAware: Bool + + override var requiresRawContext: Bool { isRawAware } init(_ handler: @escaping @Sendable (Request) async throws -> Response) { - self._handle = handler + self._handle = { request, _ in try await handler(request) } + self.isRawAware = false + super.init() + } + + init( + rawAware handler: + @escaping @Sendable (Request, RawRequestContext) async throws -> Response + ) { + self._handle = { request, rawContext in + guard let rawContext else { + throw MCPError.internalError("Raw request context is unavailable") + } + return try await handler(request, rawContext) + } + self.isRawAware = true super.init() } - override func callAsFunction(_ request: AnyRequest) async throws -> AnyResponse { + override func callAsFunction( + _ request: AnyRequest, + rawContext: RawRequestContext? = nil + ) async throws -> AnyResponse { let encoder = JSONEncoder() let decoder = JSONDecoder() @@ -189,7 +244,7 @@ final class TypedRequestHandler: RequestHandlerBox, @unchecked Sendab let request = try decoder.decode(Request.self, from: data) // Handle with concrete type - let response = try await _handle(request) + let response = try await _handle(request, rawContext) // Convert result to AnyMethod response switch response.result { diff --git a/Sources/MCP/Base/RawJSONValue.swift b/Sources/MCP/Base/RawJSONValue.swift new file mode 100644 index 00000000..e70faa02 --- /dev/null +++ b/Sources/MCP/Base/RawJSONValue.swift @@ -0,0 +1,1470 @@ +import Foundation + +/// Resource limits applied while decoding or encoding exact JSON values. +/// +/// The exact JSON parser uses these limits before constructing strings or +/// decoding programmatic data values, so callers can safely expose raw request +/// parameters to handlers without first passing unbounded input through +/// `JSONDecoder`. +public struct RawJSONLimits: Hashable, Sendable { + /// The maximum size of one JSON document, in bytes. + public var maximumDocumentBytes: Int + + /// The maximum UTF-8 byte length of one decoded JSON string. + public var maximumStringBytes: Int + + /// The maximum nesting depth of arrays and objects. + public var maximumDepth: Int + + /// The maximum total number of array elements and object members. + public var maximumContainerEntries: Int + + /// Conservative defaults suitable for MCP messages that may contain + /// image or audio payloads. + public static let `default` = RawJSONLimits( + maximumDocumentBytes: 64 * 1024 * 1024, + maximumStringBytes: 48 * 1024 * 1024, + maximumDepth: 128, + maximumContainerEntries: 250_000 + ) + + public init( + maximumDocumentBytes: Int, + maximumStringBytes: Int, + maximumDepth: Int, + maximumContainerEntries: Int + ) { + self.maximumDocumentBytes = maximumDocumentBytes + self.maximumStringBytes = maximumStringBytes + self.maximumDepth = maximumDepth + self.maximumContainerEntries = maximumContainerEntries + } +} + +/// An error produced while parsing or encoding an exact JSON value. +public struct RawJSONError: Error, Hashable, LocalizedError, Sendable { + /// The byte offset at which the error was detected, when parsing input. + public let byteOffset: Int? + + /// A human-readable description of the violated JSON or resource-limit rule. + public let message: String + + public var errorDescription: String? { + if let byteOffset { + return "Raw JSON error at byte \(byteOffset): \(message)" + } + return "Raw JSON error: \(message)" + } + + init(byteOffset: Int? = nil, message: String) { + self.byteOffset = byteOffset + self.message = message + } +} + +/// A JSON object that retains source member order and compares keys by their +/// exact Unicode scalar sequence. +/// +/// Swift `String` equality is canonically equivalent, which makes composed and +/// decomposed spellings compare equal. JSON object names do not have that +/// normalization rule. This representation therefore keeps members in an +/// array and performs scalar-exact lookup and equality. +public struct ExactJSONObject: Hashable, Sendable { + public struct Member: Hashable, Sendable { + public let key: String + public let value: RawJSONValue + + public init(key: String, value: RawJSONValue) { + self.key = key + self.value = value + } + + public static func == (lhs: Self, rhs: Self) -> Bool { + scalarExactEqual(lhs.key, rhs.key) && lhs.value == rhs.value + } + + public func hash(into hasher: inout Hasher) { + scalarExactHash(key, into: &hasher) + hasher.combine(value) + } + } + + /// Members in their original source order, including duplicate keys. + public let members: [Member] + + public init(members: [Member]) { + self.members = members + } + + /// Returns every value whose key has the same Unicode scalar sequence. + public func values(forExactKey key: String) -> [RawJSONValue] { + members.compactMap { member in + scalarExactEqual(member.key, key) ? member.value : nil + } + } + + /// Returns the value only when the exact key occurs once. + public func uniqueValue(forExactKey key: String) -> RawJSONValue? { + let matches = values(forExactKey: key) + guard matches.count == 1 else { return nil } + return matches[0] + } +} + +/// A lossless semantic JSON value for inspecting inbound MCP requests. +/// +/// Unlike ``Value``, this type preserves object member order and duplicate +/// keys, compares strings by their exact Unicode scalar sequence, and never +/// interprets JSON strings as data URLs. The `data` case exists only for +/// programmatically constructed values; parsing JSON always produces `string`. +public indirect enum RawJSONValue: Hashable, Sendable { + case null + case bool(Bool) + /// The exact valid JSON number spelling from the source document. + case number(String) + case string(String) + case data(mimeType: String? = nil, Data) + case array([RawJSONValue]) + case object(ExactJSONObject) + + /// Parses one complete UTF-8 JSON document without dictionary conversion, + /// string normalization, or data-URL coercion. + public static func decode( + _ data: Data, + limits: RawJSONLimits = .default + ) throws -> RawJSONValue { + try Task.checkCancellation() + try validateRawJSONLimits(limits) + guard data.count <= limits.maximumDocumentBytes else { + throw RawJSONError( + message: + "document exceeds maximumDocumentBytes (\(limits.maximumDocumentBytes))" + ) + } + + return try data.withUnsafeBytes { rawBuffer in + var parser = RawJSONParser( + bytes: rawBuffer.bindMemory(to: UInt8.self), + limits: limits + ) + return try parser.parse() + } + } + + /// Encodes the value as JSON after preflighting all configured limits. + /// + /// Programmatic `data` values are measured using the base64 length formula + /// before base64 bytes are allocated. + public func encodedData(limits: RawJSONLimits = .default) throws -> Data { + try validateRawJSONLimits(limits) + var measurement = RawJSONEncoder.Measurement(limits: limits) + let byteCount = try measurement.measure(self, depth: 0) + + var output = Data() + output.reserveCapacity(byteCount) + try RawJSONEncoder.append(self, to: &output, limits: limits, depth: 0) + return output + } + + public var stringValue: String? { + guard case .string(let value) = self else { return nil } + return value + } + + public var objectValue: ExactJSONObject? { + guard case .object(let value) = self else { return nil } + return value + } + + public static func == (lhs: Self, rhs: Self) -> Bool { + switch (lhs, rhs) { + case (.null, .null): + true + case (.bool(let lhs), .bool(let rhs)): + lhs == rhs + case (.number(let lhs), .number(let rhs)), (.string(let lhs), .string(let rhs)): + scalarExactEqual(lhs, rhs) + case (.data(let lhsType, let lhsData), .data(let rhsType, let rhsData)): + optionalScalarExactEqual(lhsType, rhsType) && lhsData == rhsData + case (.array(let lhs), .array(let rhs)): + lhs == rhs + case (.object(let lhs), .object(let rhs)): + lhs == rhs + default: + false + } + } + + public func hash(into hasher: inout Hasher) { + switch self { + case .null: + hasher.combine(0) + case .bool(let value): + hasher.combine(1) + hasher.combine(value) + case .number(let value): + hasher.combine(2) + scalarExactHash(value, into: &hasher) + case .string(let value): + hasher.combine(3) + scalarExactHash(value, into: &hasher) + case .data(let mimeType, let data): + hasher.combine(4) + if let mimeType { + hasher.combine(true) + scalarExactHash(mimeType, into: &hasher) + } else { + hasher.combine(false) + } + hasher.combine(data) + case .array(let values): + hasher.combine(5) + hasher.combine(values) + case .object(let object): + hasher.combine(6) + hasher.combine(object) + } + } + + static func batchElementRanges(in data: Data) throws -> [Range] { + try Task.checkCancellation() + let offsets = try data.withUnsafeBytes { rawBuffer in + let bytes = rawBuffer.bindMemory(to: UInt8.self) + var scanner = JSONContainerElementScanner( + bytes: bytes, + openingByte: 0x5B, + closingByte: 0x5D, + containerName: "batch" + ) + let ranges = try scanner.scan() + for range in ranges { + var parser = RawJSONParser( + bytes: UnsafeBufferPointer(rebasing: bytes[range]), + limits: .permittingDocument(range.count) + ) + try parser.validateSyntax() + } + return ranges + } + return offsets.map { range in + let lowerBound = data.index(data.startIndex, offsetBy: range.lowerBound) + let upperBound = data.index(data.startIndex, offsetBy: range.upperBound) + return lowerBound.. RawJSONRequestRoute? { + try Task.checkCancellation() + let maximumMethodBytes = rawAwareMethods.reduce(into: 0) { maximum, method in + maximum = max(maximum, max( + method.precomposedStringWithCanonicalMapping.utf8.count, + method.decomposedStringWithCanonicalMapping.utf8.count + )) + } + return try data.withUnsafeBytes { rawBuffer in + var parser = RawJSONParser( + bytes: rawBuffer.bindMemory(to: UInt8.self), + limits: .permittingDocument(data.count) + ) + return try parser.parseRequestRoute( + maximumMethodBytes: maximumMethodBytes, + rawAwareMethods: rawAwareMethods, + maximumIDBytes: maximumIDBytes + ) + } + } +} + +struct RawJSONRequestRoute { + let rawAwareMethod: String? + let methodMemberCount: Int + let hasID: Bool + let id: ID? +} + +extension RawJSONLimits { + static func permittingDocument(_ byteCount: Int) -> RawJSONLimits { + RawJSONLimits( + maximumDocumentBytes: byteCount, + maximumStringBytes: byteCount, + maximumDepth: min(max(byteCount, 1), 512), + maximumContainerEntries: byteCount + ) + } + +} + +private func scalarExactEqual(_ lhs: String, _ rhs: String) -> Bool { + lhs.unicodeScalars.elementsEqual(rhs.unicodeScalars) { $0.value == $1.value } +} + +private struct UTF8ValidationState { + private var remainingContinuationBytes = 0 + private var nextLowerBound: UInt8 = 0x80 + private var nextUpperBound: UInt8 = 0xBF + + var isComplete: Bool { + remainingContinuationBytes == 0 + } + + mutating func consume(_ byte: UInt8) -> Bool { + if remainingContinuationBytes > 0 { + guard byte >= nextLowerBound, byte <= nextUpperBound else { return false } + remainingContinuationBytes -= 1 + nextLowerBound = 0x80 + nextUpperBound = 0xBF + return true + } + + switch byte { + case 0x00...0x7F: + return true + case 0xC2...0xDF: + remainingContinuationBytes = 1 + case 0xE0: + remainingContinuationBytes = 2 + nextLowerBound = 0xA0 + case 0xE1...0xEC, 0xEE...0xEF: + remainingContinuationBytes = 2 + case 0xED: + remainingContinuationBytes = 2 + nextUpperBound = 0x9F + case 0xF0: + remainingContinuationBytes = 3 + nextLowerBound = 0x90 + case 0xF1...0xF3: + remainingContinuationBytes = 3 + case 0xF4: + remainingContinuationBytes = 3 + nextUpperBound = 0x8F + default: + return false + } + return true + } +} + +private struct JSONContainerElementScanner { + let bytes: UnsafeBufferPointer + let openingByte: UInt8 + let closingByte: UInt8 + let containerName: String + var index = 0 + var nextCancellationCheck = 4_096 + + mutating func scan() throws -> [Range] { + try skipWhitespace() + guard currentByte == openingByte else { + throw error("expected \(containerName)") + } + try advance() + try skipWhitespace() + if currentByte == closingByte { + try advance() + try finishDocument() + return [] + } + + var ranges: [Range] = [] + while true { + let start = index + let untrimmedEnd = try scanElementEnd() + var end = untrimmedEnd + while end > start, Self.isWhitespace(bytes[end - 1]) { + end -= 1 + } + guard end > start else { + throw error("batch item must not be empty") + } + ranges.append(start.. Int { + var nestedContainers = 0 + var inString = false + var escaped = false + + while let byte = currentByte { + if inString { + if escaped { + escaped = false + } else if byte == 0x5C { + escaped = true + } else if byte == 0x22 { + inString = false + } + try advance() + continue + } + + switch byte { + case 0x22: + inString = true + case 0x5B, 0x7B: + nestedContainers += 1 + case 0x5D, 0x7D: + if byte == closingByte, nestedContainers == 0 { + return index + } + if nestedContainers > 0 { + nestedContainers -= 1 + } + case 0x2C where nestedContainers == 0: + return index + default: + break + } + try advance() + } + + if inString { + throw error("unterminated string in batch") + } + throw error("unterminated \(containerName)") + } + + private mutating func finishDocument() throws { + try skipWhitespace() + guard index == bytes.count else { + throw error("unexpected trailing bytes after batch") + } + } + + private mutating func skipWhitespace() throws { + while let byte = currentByte, Self.isWhitespace(byte) { + try advance() + } + } + + private static func isWhitespace(_ byte: UInt8) -> Bool { + byte == 0x20 || byte == 0x09 || byte == 0x0A || byte == 0x0D + } + + private mutating func advance() throws { + index += 1 + if index >= nextCancellationCheck { + try Task.checkCancellation() + nextCancellationCheck = index + 4_096 + } + } + + private var currentByte: UInt8? { + index < bytes.count ? bytes[index] : nil + } + + private func error(_ message: String) -> RawJSONError { + RawJSONError(byteOffset: index, message: message) + } + +} + +private func optionalScalarExactEqual(_ lhs: String?, _ rhs: String?) -> Bool { + switch (lhs, rhs) { + case (.none, .none): true + case (.some(let lhs), .some(let rhs)): scalarExactEqual(lhs, rhs) + default: false + } +} + +private func scalarExactHash(_ value: String, into hasher: inout Hasher) { + hasher.combine(value.unicodeScalars.count) + for scalar in value.unicodeScalars { + hasher.combine(scalar.value) + } +} + +private func validateRawJSONLimits(_ limits: RawJSONLimits) throws { + guard limits.maximumDocumentBytes >= 0, + limits.maximumStringBytes >= 0, + limits.maximumDepth >= 0, + limits.maximumContainerEntries >= 0 + else { + throw RawJSONError(message: "JSON limits must not be negative") + } +} + +private struct RawJSONParser { + private let bytes: UnsafeBufferPointer + private let limits: RawJSONLimits + private var index = 0 + private var containerEntries = 0 + private var nextCancellationCheck = 4_096 + + init(bytes: UnsafeBufferPointer, limits: RawJSONLimits) { + self.bytes = bytes + self.limits = limits + } + + mutating func parse() throws -> RawJSONValue { + try checkCancellation() + try skipWhitespace() + let value = try parseValue(depth: 0) + try skipWhitespace() + guard index == bytes.count else { + throw error("unexpected trailing bytes") + } + return value + } + + mutating func validateSyntax() throws { + try checkCancellation() + try skipWhitespace() + try skipRoutingValue(depth: 0) + try skipWhitespace() + guard index == bytes.count else { + throw error("unexpected trailing bytes") + } + } + + mutating func parseRequestRoute( + maximumMethodBytes: Int, + rawAwareMethods: [String], + maximumIDBytes: Int + ) throws -> RawJSONRequestRoute? { + try checkCancellation() + try skipWhitespace() + guard currentByte == 0x7B else { return nil } + try consume(0x7B) + try skipWhitespace() + + var rawAwareMethod: String? + var methodMemberCount = 0 + var hasID = false + var id: ID? + + if currentByte != 0x7D { + while true { + guard currentByte == 0x22 else { + throw error("expected string object key") + } + let key = try parseRoutingString(maximumCapturedBytes: 6) + try skipWhitespace() + try consume(0x3A) + try skipWhitespace() + + if key.map({ scalarExactEqual($0, "method") }) == true { + methodMemberCount += 1 + if currentByte == 0x22 { + if let method = try parseRoutingString( + maximumCapturedBytes: maximumMethodBytes + ), rawAwareMethods.contains(method) { + rawAwareMethod = rawAwareMethod ?? method + } + } else { + try skipRoutingValue(depth: 1) + } + } else if key.map({ scalarExactEqual($0, "id") }) == true { + hasID = true + id = try parseRoutingID(maximumStringBytes: maximumIDBytes) + } else { + try skipRoutingValue(depth: 1) + } + + try skipWhitespace() + switch currentByte { + case 0x2C: + try advance() + try skipWhitespace() + case 0x7D: + break + case nil: + throw error("unterminated request object") + default: + throw error("expected ',' or '}' in request object") + } + if currentByte == 0x7D { break } + } + } + + try consume(0x7D) + try skipWhitespace() + guard index == bytes.count else { + throw error("unexpected trailing bytes") + } + return RawJSONRequestRoute( + rawAwareMethod: rawAwareMethod, + methodMemberCount: methodMemberCount, + hasID: hasID, + id: id + ) + } + + private mutating func parseRoutingID(maximumStringBytes: Int) throws -> ID? { + guard let byte = currentByte else { return nil } + switch byte { + case 0x22: + return try parseRoutingString(maximumCapturedBytes: maximumStringBytes).map(ID.string) + case 0x2D, 0x30...0x39: + guard let spelling = try parseRoutingNumber(maximumCapturedBytes: 32), + let value = Int(spelling) + else { + return nil + } + return .number(value) + default: + try skipRoutingValue(depth: 1) + return nil + } + } + + private mutating func parseRoutingString(maximumCapturedBytes: Int) throws -> String? { + try consume(0x22) + var captured: [UInt8] = [] + captured.reserveCapacity(min(maximumCapturedBytes, 64)) + var exceededLimit = false + var utf8Validation = UTF8ValidationState() + + while let byte = currentByte { + switch byte { + case 0x22: + guard utf8Validation.isComplete else { + throw error("invalid UTF-8 in string") + } + try advance() + guard !exceededLimit else { return nil } + guard let string = String(bytes: captured, encoding: .utf8) else { + throw error("invalid UTF-8 in string") + } + return string + case 0x5C: + guard utf8Validation.isComplete else { + throw error("invalid UTF-8 in string") + } + try advance() + let escapedBytes = try parseRoutingEscape() + appendRoutingBytes( + escapedBytes, + maximumCapturedBytes: maximumCapturedBytes, + captured: &captured, + exceededLimit: &exceededLimit + ) + case 0x00...0x1F: + throw error("unescaped control character in string") + default: + guard utf8Validation.consume(byte) else { + throw error("invalid UTF-8 in string") + } + appendRoutingBytes( + CollectionOfOne(byte), + maximumCapturedBytes: maximumCapturedBytes, + captured: &captured, + exceededLimit: &exceededLimit + ) + try advance() + } + } + throw error("unterminated string") + } + + private mutating func parseRoutingEscape() throws -> [UInt8] { + guard let byte = currentByte else { + throw error("unterminated string escape") + } + try advance() + switch byte { + case 0x22, 0x5C, 0x2F: + return [byte] + case 0x62: + return [0x08] + case 0x66: + return [0x0C] + case 0x6E: + return [0x0A] + case 0x72: + return [0x0D] + case 0x74: + return [0x09] + case 0x75: + let first = try parseHexQuad() + let scalarValue: UInt32 + if (0xD800...0xDBFF).contains(first) { + guard currentByte == 0x5C else { + throw error("high surrogate must be followed by a low surrogate") + } + try advance() + guard currentByte == 0x75 else { + throw error("high surrogate must be followed by a Unicode escape") + } + try advance() + let second = try parseHexQuad() + guard (0xDC00...0xDFFF).contains(second) else { + throw error("invalid low surrogate") + } + scalarValue = 0x10000 + ((first - 0xD800) << 10) + (second - 0xDC00) + } else { + guard !(0xDC00...0xDFFF).contains(first) else { + throw error("unpaired low surrogate") + } + scalarValue = first + } + guard let scalar = UnicodeScalar(scalarValue) else { + throw error("invalid Unicode scalar") + } + return Array(String(scalar).utf8) + default: + throw error("invalid string escape") + } + } + + private func appendRoutingBytes( + _ bytes: S, + maximumCapturedBytes: Int, + captured: inout [UInt8], + exceededLimit: inout Bool + ) where S.Element == UInt8 { + guard !exceededLimit else { return } + for byte in bytes { + if captured.count == maximumCapturedBytes { + exceededLimit = true + captured.removeAll(keepingCapacity: false) + return + } + captured.append(byte) + } + } + + private mutating func skipRoutingValue(depth: Int) throws { + guard depth <= 512 else { + throw error("routing depth exceeds 512") + } + guard let byte = currentByte else { + throw error("unexpected end of input") + } + switch byte { + case 0x6E: + try consumeKeyword("null") + case 0x74: + try consumeKeyword("true") + case 0x66: + try consumeKeyword("false") + case 0x22: + _ = try parseRoutingString(maximumCapturedBytes: 0) + case 0x5B: + try skipRoutingArray(depth: depth + 1) + case 0x7B: + try skipRoutingObject(depth: depth + 1) + case 0x2D, 0x30...0x39: + _ = try parseRoutingNumber(maximumCapturedBytes: 0) + default: + throw error("unexpected byte") + } + } + + private mutating func skipRoutingArray(depth: Int) throws { + try consume(0x5B) + try skipWhitespace() + if currentByte == 0x5D { + try advance() + return + } + while true { + try skipRoutingValue(depth: depth) + try skipWhitespace() + switch currentByte { + case 0x2C: + try advance() + try skipWhitespace() + case 0x5D: + try advance() + return + default: + throw error("expected ',' or ']' in array") + } + } + } + + private mutating func skipRoutingObject(depth: Int) throws { + try consume(0x7B) + try skipWhitespace() + if currentByte == 0x7D { + try advance() + return + } + while true { + guard currentByte == 0x22 else { + throw error("expected string object key") + } + _ = try parseRoutingString(maximumCapturedBytes: 0) + try skipWhitespace() + try consume(0x3A) + try skipWhitespace() + try skipRoutingValue(depth: depth) + try skipWhitespace() + switch currentByte { + case 0x2C: + try advance() + try skipWhitespace() + case 0x7D: + try advance() + return + default: + throw error("expected ',' or '}' in object") + } + } + } + + private mutating func parseRoutingNumber(maximumCapturedBytes: Int) throws -> String? { + let range = try scanNumber() + guard range.count <= maximumCapturedBytes else { return nil } + return String(decoding: bytes[range], as: UTF8.self) + } + + private mutating func parseValue(depth: Int) throws -> RawJSONValue { + guard let byte = currentByte else { + throw error("unexpected end of input") + } + + switch byte { + case 0x6E: + try consumeKeyword("null") + return .null + case 0x74: + try consumeKeyword("true") + return .bool(true) + case 0x66: + try consumeKeyword("false") + return .bool(false) + case 0x22: + return .string(try parseString()) + case 0x5B: + try checkContainerDepth(depth + 1) + return .array(try parseArray(depth: depth + 1)) + case 0x7B: + try checkContainerDepth(depth + 1) + return .object(try parseObject(depth: depth + 1)) + case 0x2D, 0x30...0x39: + return .number(try parseNumber()) + default: + throw error("unexpected byte") + } + } + + private mutating func parseArray(depth: Int) throws -> [RawJSONValue] { + try consume(0x5B) + try skipWhitespace() + if currentByte == 0x5D { + try advance() + return [] + } + + var values: [RawJSONValue] = [] + while true { + try countContainerEntry() + values.append(try parseValue(depth: depth)) + try skipWhitespace() + switch currentByte { + case 0x2C: + try advance() + try skipWhitespace() + case 0x5D: + try advance() + return values + case nil: + throw error("unterminated array") + default: + throw error("expected ',' or ']' in array") + } + } + } + + private mutating func parseObject(depth: Int) throws -> ExactJSONObject { + try consume(0x7B) + try skipWhitespace() + if currentByte == 0x7D { + try advance() + return ExactJSONObject(members: []) + } + + var members: [ExactJSONObject.Member] = [] + while true { + guard currentByte == 0x22 else { + throw error("expected string object key") + } + let key = try parseString() + try skipWhitespace() + try consume(0x3A) + try skipWhitespace() + try countContainerEntry() + let value = try parseValue(depth: depth) + members.append(.init(key: key, value: value)) + try skipWhitespace() + + switch currentByte { + case 0x2C: + try advance() + try skipWhitespace() + case 0x7D: + try advance() + return ExactJSONObject(members: members) + case nil: + throw error("unterminated object") + default: + throw error("expected ',' or '}' in object") + } + } + } + + private mutating func parseString() throws -> String { + try consume(0x22) + var segmentStart = index + var result = "" + var decodedByteCount = 0 + + while let byte = currentByte { + switch byte { + case 0x22: + try appendUTF8( + from: segmentStart, + to: index, + into: &result, + decodedByteCount: &decodedByteCount + ) + try advance() + return result + case 0x5C: + try appendUTF8( + from: segmentStart, + to: index, + into: &result, + decodedByteCount: &decodedByteCount + ) + try advance() + try appendEscape(into: &result, decodedByteCount: &decodedByteCount) + segmentStart = index + case 0x00...0x1F: + throw error("unescaped control character in string") + default: + try advance() + } + } + + throw error("unterminated string") + } + + private mutating func appendEscape( + into result: inout String, + decodedByteCount: inout Int + ) throws { + guard let byte = currentByte else { + throw error("unterminated string escape") + } + try advance() + + switch byte { + case 0x22: + try addDecodedBytes(1, to: &decodedByteCount) + result.append("\"") + case 0x5C: + try addDecodedBytes(1, to: &decodedByteCount) + result.append("\\") + case 0x2F: + try addDecodedBytes(1, to: &decodedByteCount) + result.append("/") + case 0x62: + try addDecodedBytes(1, to: &decodedByteCount) + result.append("\u{08}") + case 0x66: + try addDecodedBytes(1, to: &decodedByteCount) + result.append("\u{0C}") + case 0x6E: + try addDecodedBytes(1, to: &decodedByteCount) + result.append("\n") + case 0x72: + try addDecodedBytes(1, to: &decodedByteCount) + result.append("\r") + case 0x74: + try addDecodedBytes(1, to: &decodedByteCount) + result.append("\t") + case 0x75: + let first = try parseHexQuad() + let scalarValue: UInt32 + if (0xD800...0xDBFF).contains(first) { + guard currentByte == 0x5C else { + throw error("high surrogate must be followed by a low surrogate") + } + try advance() + guard currentByte == 0x75 else { + throw error("high surrogate must be followed by a Unicode escape") + } + try advance() + let second = try parseHexQuad() + guard (0xDC00...0xDFFF).contains(second) else { + throw error("invalid low surrogate") + } + scalarValue = 0x10000 + ((first - 0xD800) << 10) + (second - 0xDC00) + } else { + guard !(0xDC00...0xDFFF).contains(first) else { + throw error("unpaired low surrogate") + } + scalarValue = first + } + + guard let scalar = UnicodeScalar(scalarValue) else { + throw error("invalid Unicode scalar") + } + try addDecodedBytes(scalar.utf8ByteCount, to: &decodedByteCount) + result.unicodeScalars.append(scalar) + default: + throw error("invalid string escape") + } + + } + + private mutating func parseHexQuad() throws -> UInt32 { + var value: UInt32 = 0 + for _ in 0..<4 { + guard let byte = currentByte, let digit = hexDigit(byte) else { + throw error("invalid Unicode escape") + } + value = (value << 4) | digit + try advance() + } + return value + } + + private mutating func parseNumber() throws -> String { + let range = try scanNumber() + return String(decoding: bytes[range], as: UTF8.self) + } + + private mutating func scanNumber() throws -> Range { + let start = index + if currentByte == 0x2D { + try advance() + } + + guard let first = currentByte else { + throw error("incomplete number") + } + if first == 0x30 { + try advance() + if let byte = currentByte, (0x30...0x39).contains(byte) { + throw error("leading zero in number") + } + } else if (0x31...0x39).contains(first) { + try consumeDigits() + } else { + throw error("invalid number") + } + + if currentByte == 0x2E { + try advance() + guard let byte = currentByte, (0x30...0x39).contains(byte) else { + throw error("fraction requires a digit") + } + try consumeDigits() + } + + if currentByte == 0x65 || currentByte == 0x45 { + try advance() + if currentByte == 0x2B || currentByte == 0x2D { + try advance() + } + guard let byte = currentByte, (0x30...0x39).contains(byte) else { + throw error("exponent requires a digit") + } + try consumeDigits() + } + + return start..= nextCancellationCheck { + try checkCancellation() + nextCancellationCheck = index + 4_096 + } + } + + private func checkCancellation() throws { + try Task.checkCancellation() + } + + private var currentByte: UInt8? { + index < bytes.count ? bytes[index] : nil + } + + private func error(_ message: String) -> RawJSONError { + RawJSONError(byteOffset: index, message: message) + } +} + +private func hexDigit(_ byte: UInt8) -> UInt32? { + switch byte { + case 0x30...0x39: UInt32(byte - 0x30) + case 0x41...0x46: UInt32(byte - 0x41 + 10) + case 0x61...0x66: UInt32(byte - 0x61 + 10) + default: nil + } +} + +extension UnicodeScalar { + fileprivate var utf8ByteCount: Int { + switch value { + case 0...0x7F: 1 + case 0x80...0x7FF: 2 + case 0x800...0xFFFF: 3 + default: 4 + } + } +} + +private enum RawJSONEncoder { + struct Measurement { + let limits: RawJSONLimits + var containerEntries = 0 + + mutating func measure(_ value: RawJSONValue, depth: Int) throws -> Int { + try Task.checkCancellation() + + let count: Int + switch value { + case .null: + count = 4 + case .bool(let value): + count = value ? 4 : 5 + case .number(let value): + guard value.utf8.count <= limits.maximumDocumentBytes, + isValidJSONNumber(value) + else { + throw RawJSONError(message: "invalid programmatic JSON number") + } + count = value.utf8.count + case .string(let value): + try checkString(value) + count = try jsonStringByteCount(value) + case .data(let mimeType, let data): + let prefix = "data:\(mimeType ?? "text/plain");base64," + let base64Count = try base64EncodedLength(data.count) + let stringCount = try checkedAdd(prefix.utf8.count, base64Count) + guard stringCount <= limits.maximumStringBytes else { + throw RawJSONError( + message: + "encoded data string exceeds maximumStringBytes (\(limits.maximumStringBytes))" + ) + } + count = try checkedAdd(try jsonStringByteCount(prefix), base64Count) + case .array(let values): + try checkContainerDepth(depth + 1) + try countEntries(values.count) + var total = 2 + for (offset, child) in values.enumerated() { + if offset > 0 { total = try checkedAdd(total, 1) } + total = try checkedAdd(total, try measure(child, depth: depth + 1)) + } + count = total + case .object(let object): + try checkContainerDepth(depth + 1) + try countEntries(object.members.count) + var total = 2 + for (offset, member) in object.members.enumerated() { + try checkString(member.key) + if offset > 0 { total = try checkedAdd(total, 1) } + total = try checkedAdd(total, try jsonStringByteCount(member.key)) + total = try checkedAdd(total, 1) + total = try checkedAdd(total, try measure(member.value, depth: depth + 1)) + } + count = total + } + + guard count <= limits.maximumDocumentBytes else { + throw RawJSONError( + message: + "encoded value exceeds maximumDocumentBytes (\(limits.maximumDocumentBytes))" + ) + } + return count + } + + private func checkString(_ value: String) throws { + guard value.utf8.count <= limits.maximumStringBytes else { + throw RawJSONError( + message: "string exceeds maximumStringBytes (\(limits.maximumStringBytes))") + } + } + + private mutating func countEntries(_ count: Int) throws { + let (next, overflow) = containerEntries.addingReportingOverflow(count) + guard !overflow, next <= limits.maximumContainerEntries else { + throw RawJSONError( + message: + "container entries exceed maximumContainerEntries (\(limits.maximumContainerEntries))" + ) + } + containerEntries = next + } + + private func checkContainerDepth(_ depth: Int) throws { + guard depth <= limits.maximumDepth else { + throw RawJSONError( + message: "nesting exceeds maximumDepth (\(limits.maximumDepth))") + } + } + } + + static func append( + _ value: RawJSONValue, + to output: inout Data, + limits: RawJSONLimits, + depth: Int + ) throws { + try Task.checkCancellation() + switch value { + case .null: + output.append(contentsOf: "null".utf8) + case .bool(let value): + output.append(contentsOf: (value ? "true" : "false").utf8) + case .number(let value): + output.append(contentsOf: value.utf8) + case .string(let value): + try appendJSONString(value, to: &output) + case .data(let mimeType, let data): + let prefix = "data:\(mimeType ?? "text/plain");base64," + output.append(0x22) + try appendJSONStringContents(prefix, to: &output) + try Task.checkCancellation() + output.append(data.base64EncodedData()) + try Task.checkCancellation() + output.append(0x22) + case .array(let values): + output.append(0x5B) + for (offset, child) in values.enumerated() { + if offset > 0 { output.append(0x2C) } + try append(child, to: &output, limits: limits, depth: depth + 1) + } + output.append(0x5D) + case .object(let object): + output.append(0x7B) + for (offset, member) in object.members.enumerated() { + if offset > 0 { output.append(0x2C) } + try appendJSONString(member.key, to: &output) + output.append(0x3A) + try append(member.value, to: &output, limits: limits, depth: depth + 1) + } + output.append(0x7D) + } + } + + private static func appendJSONString(_ value: String, to output: inout Data) throws { + output.append(0x22) + try appendJSONStringContents(value, to: &output) + output.append(0x22) + } + + private static func appendJSONStringContents(_ value: String, to output: inout Data) throws { + var byteCount = 0 + for byte in value.utf8 { + switch byte { + case 0x08: output.append(contentsOf: "\\b".utf8) + case 0x09: output.append(contentsOf: "\\t".utf8) + case 0x0A: output.append(contentsOf: "\\n".utf8) + case 0x0C: output.append(contentsOf: "\\f".utf8) + case 0x0D: output.append(contentsOf: "\\r".utf8) + case 0x22: output.append(contentsOf: "\\\"".utf8) + case 0x5C: output.append(contentsOf: "\\\\".utf8) + case 0x00...0x1F: + output.append(contentsOf: [ + 0x5C, 0x75, 0x30, 0x30, + lowercaseHexDigits[Int((byte >> 4) & 0xF)], + lowercaseHexDigits[Int(byte & 0xF)], + ]) + default: + output.append(byte) + } + byteCount += 1 + if byteCount == 4_096 { + try Task.checkCancellation() + byteCount = 0 + } + } + } +} + +private func jsonStringByteCount(_ value: String) throws -> Int { + var count = 2 + var byteCount = 0 + for byte in value.utf8 { + let addition: Int + switch byte { + case 0x08, 0x09, 0x0A, 0x0C, 0x0D, 0x22, 0x5C: + addition = 2 + case 0x00...0x1F: + addition = 6 + default: + addition = 1 + } + count = try checkedAdd(count, addition) + byteCount += 1 + if byteCount == 4_096 { + try Task.checkCancellation() + byteCount = 0 + } + } + return count +} + +private let lowercaseHexDigits = Array("0123456789abcdef".utf8) + +private func base64EncodedLength(_ byteCount: Int) throws -> Int { + guard byteCount >= 0 else { + throw RawJSONError(message: "negative data length") + } + let completeGroups = byteCount / 3 + let (completeLength, multiplyOverflow) = completeGroups.multipliedReportingOverflow(by: 4) + guard !multiplyOverflow else { + throw RawJSONError(message: "encoded data length overflows Int") + } + guard byteCount % 3 != 0 else { return completeLength } + return try checkedAdd(completeLength, 4) +} + +private func checkedAdd(_ lhs: Int, _ rhs: Int) throws -> Int { + let (result, overflow) = lhs.addingReportingOverflow(rhs) + guard !overflow else { + throw RawJSONError(message: "encoded JSON length overflows Int") + } + return result +} + +private func isValidJSONNumber(_ value: String) -> Bool { + let bytes = Array(value.utf8) + guard !bytes.isEmpty else { return false } + var index = 0 + + if bytes[index] == 0x2D { + index += 1 + guard index < bytes.count else { return false } + } + + if bytes[index] == 0x30 { + index += 1 + if index < bytes.count, (0x30...0x39).contains(bytes[index]) { + return false + } + } else if (0x31...0x39).contains(bytes[index]) { + index += 1 + while index < bytes.count, (0x30...0x39).contains(bytes[index]) { + index += 1 + } + } else { + return false + } + + if index < bytes.count, bytes[index] == 0x2E { + index += 1 + guard index < bytes.count, (0x30...0x39).contains(bytes[index]) else { + return false + } + while index < bytes.count, (0x30...0x39).contains(bytes[index]) { + index += 1 + } + } + + if index < bytes.count, bytes[index] == 0x65 || bytes[index] == 0x45 { + index += 1 + if index < bytes.count, bytes[index] == 0x2B || bytes[index] == 0x2D { + index += 1 + } + guard index < bytes.count, (0x30...0x39).contains(bytes[index]) else { + return false + } + while index < bytes.count, (0x30...0x39).contains(bytes[index]) { + index += 1 + } + } + + return index == bytes.count +} diff --git a/Sources/MCP/Server/Server.swift b/Sources/MCP/Server/Server.swift index 060b51f8..85aa80b4 100644 --- a/Sources/MCP/Server/Server.swift +++ b/Sources/MCP/Server/Server.swift @@ -160,6 +160,8 @@ public actor Server { public var capabilities: Capabilities /// The server configuration public var configuration: Configuration + /// Bounds for exact inbound JSON retained by raw-aware handlers. + private let rawJSONLimits: RawJSONLimits /// Request handlers private var methodHandlers: [String: RequestHandlerBox] = [:] @@ -190,12 +192,14 @@ public actor Server { title: String? = nil, instructions: String? = nil, capabilities: Server.Capabilities = .init(), - configuration: Configuration = .default + configuration: Configuration = .default, + rawJSONLimits: RawJSONLimits = .default ) { self.serverInfo = Server.Info(name: name, version: version, title: title) self.capabilities = capabilities self.configuration = configuration self.instructions = instructions + self.rawJSONLimits = rawJSONLimits } /// Start the server @@ -224,33 +228,65 @@ public actor Server { var requestID: ID? do { - // Attempt to decode as batch first, then as individual response, request, or notification let decoder = JSONDecoder() - if let batch = try? decoder.decode(Server.Batch.self, from: data) { + if data.firstNonJSONWhitespaceByte == 0x5B { + let batch = try self.decodeBatch(data) try await handleBatch(batch) } 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) - } - } else if let message = try? decoder.decode(AnyMessage.self, from: data) { - try await handleMessage(message) } else { - // Try to extract request ID from raw JSON if possible - if let json = try? JSONDecoder().decode( - [String: Value].self, from: data), - let idValue = json["id"] + let route = try self.requestRoute(in: data) + requestID = route?.id + if route?.ambiguousRawMethod == true { + throw MCPError.invalidRequest( + "Raw-aware requests must contain exactly one method" + ) + } + if let route, route.hasID, let method = route.method, + self.methodHandlers[method]?.requiresRawContext == true { - if let strValue = idValue.stringValue { - requestID = .string(strValue) - } else if let intValue = idValue.intValue { - requestID = .number(intValue) + let rawValue = try RawJSONValue.decode( + data, + limits: self.rawJSONLimits + ) + guard case .object(let rawObject) = rawValue else { + throw MCPError.parseError("Invalid message format") + } + requestID = rawObject.jsonRPCRequestID ?? requestID + + guard let request = try? decoder.decode(AnyRequest.self, from: data) + else { + throw MCPError.invalidRequest("Invalid request format") + } + let rawContext = RawRequestContext(request: rawObject) + Task { + _ = try? await self.handleRequest( + request, + rawContext: rawContext, + sendResponse: true + ) } + } else if let request = try? decoder.decode(AnyRequest.self, from: data) { + requestID = request.id + Task { + _ = try? await self.handleRequest( + request, + rawContext: nil, + sendResponse: true + ) + } + } else if let message = try? decoder.decode(AnyMessage.self, from: data) { + try await handleMessage(message) + } else { + requestID = requestID + ?? (try? decoder.decode(InboundIDHeader.self, from: data).id) + throw MCPError.parseError("Invalid message format") } - throw MCPError.parseError("Invalid message format") } + } catch is CancellationError where Task.isCancelled { + break + } catch is CancellationError { + continue } catch let error where MCPError.isResourceTemporarilyUnavailable(error) { // Resource temporarily unavailable, retry after a short delay try? await Task.sleep(for: .milliseconds(10)) @@ -260,8 +296,7 @@ public actor Server { "Error processing message", metadata: ["error": "\(error)"]) let response = AnyMethod.response( id: requestID ?? .random, - error: error as? MCPError - ?? MCPError.internalError(error.localizedDescription) + error: error.asMCPError ) try? await send(response) } @@ -349,6 +384,26 @@ public actor Server { return self } + /// Register a method handler that receives both the validated typed request + /// and its exact inbound JSON representation. + /// + /// This overload is additive: existing parameter-only handlers continue to + /// use the same semantic `Value` decoding behavior. Use the raw context when + /// object order, duplicate keys, exact Unicode scalar spellings, or strings + /// that resemble data URLs are significant. + @discardableResult + public func withMethodHandler( + _ type: M.Type, + handler: + @escaping @Sendable (Request, RawRequestContext) async throws -> M.Result + ) -> Self { + methodHandlers[M.name] = TypedRequestHandler(rawAware: { request, rawContext in + let result = try await handler(request, rawContext) + return Response(id: request.id, result: result) + }) + return self + } + /// Register a notification handler @discardableResult public func onNotification( @@ -663,9 +718,9 @@ public actor Server { struct Batch: Sendable { /// An item in a JSON-RPC batch enum Item: Sendable { - case request(Request) + case request(Request, RawRequestContext?) case notification(Message) - + case invalid(ID, MCPError) } var items: [Item] @@ -675,6 +730,96 @@ public actor Server { } } + private func decodeBatch(_ data: Data) throws -> Batch { + let ranges = try RawJSONValue.batchElementRanges(in: data) + return Batch(items: try ranges.map { range in + try self.decodeBatchItem(data[range]) + }) + } + + private func decodeBatchItem(_ data: Data) throws -> Batch.Item { + let decoder = JSONDecoder() + let route: InboundRequestRoute? + do { + route = try self.requestRoute(in: data) + } catch is CancellationError { + throw CancellationError() + } catch { + return .invalid(.random, error.asMCPError) + } + if route?.ambiguousRawMethod == true { + return .invalid( + route?.id ?? .random, + MCPError.invalidRequest("Raw-aware requests must contain exactly one method") + ) + } + + if let route, route.hasID, let method = route.method, + self.methodHandlers[method]?.requiresRawContext == true + { + do { + let rawValue = try RawJSONValue.decode(data, limits: self.rawJSONLimits) + guard case .object(let object) = rawValue else { + return .invalid( + route.id ?? .random, + MCPError.invalidRequest("Invalid batch item format") + ) + } + let request = try decoder.decode(AnyRequest.self, from: data) + return .request(request, RawRequestContext(request: object)) + } catch is CancellationError { + throw CancellationError() + } catch is DecodingError { + return .invalid( + route.id ?? .random, + MCPError.invalidRequest("Invalid batch item format") + ) + } catch { + return .invalid(route.id ?? .random, error.asMCPError) + } + } + + if route?.hasID == true, let request = try? decoder.decode(AnyRequest.self, from: data) { + return .request(request, nil) + } + if route?.hasID == false, let message = try? decoder.decode(AnyMessage.self, from: data) { + return .notification(message) + } + + // Distinguish malformed JSON from a syntactically valid item with an + // invalid JSON-RPC shape, without imposing raw-aware limits on legacy traffic. + do { + _ = try RawJSONValue.decode(data, limits: .permittingDocument(data.count)) + return .invalid( + route?.id ?? .random, + MCPError.invalidRequest("Invalid batch item format") + ) + } catch is CancellationError { + throw CancellationError() + } catch { + return .invalid(route?.id ?? .random, error.asMCPError) + } + } + + private func requestRoute(in data: Data) throws -> InboundRequestRoute? { + let rawAwareMethods = methodHandlers.compactMap { method, handler in + handler.requiresRawContext ? method : nil + } + guard let route = try RawJSONValue.requestRoute( + in: data, + rawAwareMethods: rawAwareMethods, + maximumIDBytes: max(rawJSONLimits.maximumStringBytes, 1024) + ) else { + return nil + } + return InboundRequestRoute( + method: route.rawAwareMethod, + hasID: route.hasID, + id: route.id, + ambiguousRawMethod: route.rawAwareMethod != nil && route.methodMemberCount != 1 + ) + } + /// Process a batch of requests and/or notifications private func handleBatch(_ batch: Batch) async throws { await logger?.trace("Processing batch request", metadata: ["size": "\(batch.items.count)"]) @@ -693,22 +838,27 @@ public actor Server { for item in batch.items { do { switch item { - case .request(let request): + case .request(let request, let rawContext): // For batched requests, collect responses instead of sending immediately - if let response = try await handleRequest(request, sendResponse: false) { + if let response = try await handleRequest( + request, + rawContext: rawContext, + sendResponse: false + ) { responses.append(response) } case .notification(let notification): // Handle notification (no response needed) try await handleMessage(notification) + + case .invalid(let id, let error): + responses.append(AnyMethod.response(id: id, error: error)) } } catch { // Only add errors to response for requests (notifications don't have responses) - if case .request(let request) = item { - let mcpError = - error as? MCPError ?? MCPError.internalError(error.localizedDescription) - responses.append(AnyMethod.response(id: request.id, error: mcpError)) + if case .request(let request, _) = item { + responses.append(AnyMethod.response(id: request.id, error: error.asMCPError)) } } } @@ -735,7 +885,11 @@ 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) + private func handleRequest( + _ request: Request, + rawContext: RawRequestContext?, + sendResponse: Bool = true + ) async throws -> Response? { // Check if this is a pre-processed error request (empty method) @@ -795,7 +949,7 @@ public actor Server { try Task.checkCancellation() // Handle request and get response - let response = try await handler(request) + let response = try await handler(request, rawContext: rawContext) return response } catch is CancellationError { // Request was cancelled, don't send a response per MCP spec @@ -1040,49 +1194,52 @@ public actor Server { } } -extension Server.Batch: Codable { - init(from decoder: Decoder) throws { - let container = try decoder.singleValueContainer() +private struct InboundIDHeader: Decodable { + let id: ID - let encoder = JSONEncoder() - let decoder = JSONDecoder() - - var items: [Item] = [] - for item in try container.decode([Value].self) { - let data = try encoder.encode(item) - try items.append(decoder.decode(Item.self, from: data)) - } - - self.items = items + private enum CodingKeys: String, CodingKey { + case id } +} - func encode(to encoder: Encoder) throws { - var container = encoder.singleValueContainer() - try container.encode(items) - } +private struct InboundRequestRoute { + let method: String? + let hasID: Bool + let id: ID? + let ambiguousRawMethod: Bool } -extension Server.Batch.Item: Codable { - private enum CodingKeys: String, CodingKey { - case id +extension Data { + fileprivate var firstNonJSONWhitespaceByte: UInt8? { + first { byte in + byte != 0x20 && byte != 0x09 && byte != 0x0A && byte != 0x0D + } } +} - init(from decoder: Decoder) throws { - let container = try decoder.container(keyedBy: CodingKeys.self) - // Check if it's a request (has id) or notification (no id) - if container.contains(.id) { - self = .request(try Request(from: decoder)) - } else { - self = .notification(try Message(from: decoder)) +extension ExactJSONObject { + fileprivate var jsonRPCRequestID: ID? { + guard let rawID = uniqueValue(forExactKey: "id") else { return nil } + switch rawID { + case .string(let value): + return .string(value) + case .number(let value): + guard let number = Int(value) else { return nil } + return .number(number) + default: + return nil } } +} - func encode(to encoder: Encoder) throws { - switch self { - case .request(let request): - try request.encode(to: encoder) - case .notification(let notification): - try notification.encode(to: encoder) +extension Swift.Error { + fileprivate var asMCPError: MCPError { + if let error = self as? MCPError { + return error + } + if let error = self as? RawJSONError { + return .parseError(error.message) } + return .internalError(localizedDescription) } } diff --git a/Tests/MCPTests/RawJSONValueTests.swift b/Tests/MCPTests/RawJSONValueTests.swift new file mode 100644 index 00000000..24e72135 --- /dev/null +++ b/Tests/MCPTests/RawJSONValueTests.swift @@ -0,0 +1,235 @@ +import Foundation +import Testing + +@testable import MCP + +@Suite("Raw JSON Value Tests") +struct RawJSONValueTests { + @Test("Objects preserve order, duplicates, and scalar-exact keys") + func exactObjectMembers() throws { + let data = #"{"\u00e9":1,"é":2,"e\u0301":3,"same":4,"same":5}"#.data(using: .utf8)! + let value = try RawJSONValue.decode(data) + let object = try #require(value.objectValue) + + #expect(object.members.count == 5) + #expect(object.members.map(\.value) == [ + .number("1"), .number("2"), .number("3"), .number("4"), .number("5"), + ]) + #expect(object.members[0].key.unicodeScalars.map(\.value) == [0xE9]) + #expect(object.members[1].key.unicodeScalars.map(\.value) == [0xE9]) + #expect(object.members[2].key.unicodeScalars.map(\.value) == [0x65, 0x301]) + #expect(object.values(forExactKey: "é") == [.number("1"), .number("2")]) + #expect(object.values(forExactKey: "e\u{301}") == [.number("3")]) + #expect(object.values(forExactKey: "same") == [.number("4"), .number("5")]) + #expect(object.uniqueValue(forExactKey: "same") == nil) + #expect(RawJSONValue.string("é") != RawJSONValue.string("e\u{301}")) + } + + @Test("Escapes decode semantically without normalizing or coercing data URLs") + func stringsRemainStrings() throws { + let json = #""" + ["é","\u00e9","e\u0301","/","\/","data:,A%20brief%20note", + "d\u0061ta:text/plain;base64,SGVsbG8="] + """# + let data = json.data(using: .utf8)! + let value = try RawJSONValue.decode(data) + guard case .array(let values) = value else { + Issue.record("Expected an array") + return + } + + #expect(values[0] == values[1]) + #expect(values[0] != values[2]) + #expect(values[3] == values[4]) + #expect(values[5] == .string("data:,A%20brief%20note")) + #expect(values[6] == .string("data:text/plain;base64,SGVsbG8=")) + #expect(values.allSatisfy { value in + if case .data = value { return false } + return true + }) + } + + @Test("Numbers retain their exact JSON lexemes", arguments: [ + "0", "-0", "1", "1.0", "1e+10", "1E-10", "123456789012345678901234567890", + ]) + func numberLexemes(number: String) throws { + let value = try RawJSONValue.decode(Data(number.utf8)) + #expect(value == .number(number)) + #expect(try value.encodedData() == Data(number.utf8)) + } + + @Test("Malformed JSON is rejected", arguments: [ + "", "01", "-", "1.", ".1", "1e", "+1", "NaN", "Infinity", "[", "{\"a\":1", "true false", + "\"\\uD800\"", "\"\\uDC00\"", "\"\\uD800\\u0041\"", "\"\\uZZZZ\"", + ]) + func malformedJSON(json: String) { + #expect(throws: (any Error).self) { + try RawJSONValue.decode(Data(json.utf8)) + } + } + + @Test("Invalid UTF-8 and raw controls are rejected") + func invalidStrings() { + #expect(throws: RawJSONError.self) { + try RawJSONValue.decode(Data([0x22, 0xC3, 0x28, 0x22])) + } + #expect(throws: RawJSONError.self) { + try RawJSONValue.decode(Data([0x22, 0x0A, 0x22])) + } + + var malformedBatch = Data(#"[{"value":""#.utf8) + malformedBatch.append(contentsOf: [0xC3, 0x28]) + malformedBatch.append(contentsOf: Data(#""}]"#.utf8)) + #expect(throws: RawJSONError.self) { + try RawJSONValue.batchElementRanges(in: malformedBatch) + } + } + + @Test("Decoded string limits count Unicode output, not escape spelling") + func decodedStringBounds() throws { + let oneByte = limits(maximumStringBytes: 1) + #expect(try RawJSONValue.decode(Data(#""\u0061""#.utf8), limits: oneByte) == .string("a")) + + let twoBytes = limits(maximumStringBytes: 2) + #expect(try RawJSONValue.decode(Data(#""\u00e9""#.utf8), limits: twoBytes) == .string("é")) + + let fourBytes = limits(maximumStringBytes: 4) + #expect(try RawJSONValue.decode(Data(#""\uD83D\uDE00""#.utf8), limits: fourBytes) == .string("😀")) + + #expect(throws: RawJSONError.self) { + try RawJSONValue.decode(Data(#""\u00e9""#.utf8), limits: oneByte) + } + #expect(throws: RawJSONError.self) { + try RawJSONValue.decode(Data(#""e\u0301""#.utf8), limits: twoBytes) + } + } + + @Test("Container depth and entry limits are consistent") + func structuralBounds() throws { + let noContainers = limits(maximumDepth: 0) + #expect(try RawJSONValue.decode(Data("0".utf8), limits: noContainers) == .number("0")) + #expect(throws: RawJSONError.self) { + try RawJSONValue.decode(Data("[]".utf8), limits: noContainers) + } + + let oneContainer = limits(maximumDepth: 1) + #expect(try RawJSONValue.decode(Data("[]".utf8), limits: oneContainer) == .array([])) + #expect(try RawJSONValue.decode(Data("[0]".utf8), limits: oneContainer) == .array([.number("0")])) + #expect(throws: RawJSONError.self) { + try RawJSONValue.decode(Data("[[]]".utf8), limits: oneContainer) + } + + let oneEntry = limits(maximumContainerEntries: 1) + #expect(try RawJSONValue.decode(Data("[0]".utf8), limits: oneEntry) == .array([.number("0")])) + #expect(throws: RawJSONError.self) { + try RawJSONValue.decode(Data("[0,1]".utf8), limits: oneEntry) + } + } + + @Test("Document size is checked before parsing") + func documentBound() throws { + let data = Data(" [0] ".utf8) + let decoded = try RawJSONValue.decode( + data, + limits: limits(maximumDocumentBytes: data.count) + ) + #expect(decoded == .array([.number("0")])) + #expect(throws: RawJSONError.self) { + try RawJSONValue.decode(data, limits: limits(maximumDocumentBytes: data.count - 1)) + } + } + + @Test("A cancelled parse fails before copying or decoding input") + func cancellation() async { + let data = Data(repeating: 0x20, count: 8 * 1024 * 1024) + Data("null".utf8) + let task = Task { + while !Task.isCancelled { + await Task.yield() + } + return try RawJSONValue.decode(data) + } + task.cancel() + await #expect(throws: CancellationError.self) { + _ = try await task.value + } + } + + @Test("A cancelled routing scan propagates cancellation") + func routingCancellation() async { + let data = Data( + #"{"jsonrpc":"2.0","id":1,"params":{"payload":"value"},"method":"raw"}"#.utf8 + ) + let task = Task { + while !Task.isCancelled { + await Task.yield() + } + return try RawJSONValue.requestRoute( + in: data, + rawAwareMethods: ["raw"], + maximumIDBytes: 128 + ) + } + task.cancel() + await #expect(throws: CancellationError.self) { + _ = try await task.value + } + } + + @Test("Programmatic data is bounded before base64 encoding") + func programmaticDataBound() throws { + let value = RawJSONValue.data(mimeType: "application/octet-stream", Data(repeating: 0xAB, count: 100)) + #expect(throws: RawJSONError.self) { + try value.encodedData(limits: limits(maximumStringBytes: 32)) + } + + let encoded = try value.encodedData(limits: limits(maximumStringBytes: 512)) + let decoded = try RawJSONValue.decode(encoded, limits: limits(maximumStringBytes: 512)) + guard case .string(let string) = decoded else { + Issue.record("Programmatic data must encode as a JSON string") + return + } + #expect(string.hasPrefix("data:application/octet-stream;base64,")) + } + + @Test("Exact values round-trip without losing duplicate members") + func exactRoundTrip() throws { + let original = try RawJSONValue.decode( + Data(#"{"a":1,"a":2,"é":"literal","e\u0301":"decomposed","n":1e+02}"#.utf8) + ) + let encoded = try original.encodedData() + #expect(try RawJSONValue.decode(encoded) == original) + } + + @Test("Batch element splitting ignores delimiters inside nested strings") + func batchElementRanges() throws { + let data = Data(#"[ {"text":",]}","nested":[1,{"quote":"\""}]}, true, 1e2 ]"#.utf8) + let ranges = try RawJSONValue.batchElementRanges(in: data) + #expect(ranges.count == 3) + #expect(try RawJSONValue.decode(data.subdata(in: ranges[0])).objectValue != nil) + #expect(try RawJSONValue.decode(data.subdata(in: ranges[1])) == .bool(true)) + #expect(try RawJSONValue.decode(data.subdata(in: ranges[2])) == .number("1e2")) + } + + @Test("Batch splitting rejects malformed JSON", arguments: [ + "[1 2]", #"[{"x":"\q"}]"#, #"[{"x":1]]"#, + ]) + func malformedBatchSplitting(json: String) { + #expect(throws: RawJSONError.self) { + try RawJSONValue.batchElementRanges(in: Data(json.utf8)) + } + } + + private func limits( + maximumDocumentBytes: Int = 1024, + maximumStringBytes: Int = 128, + maximumDepth: Int = 8, + maximumContainerEntries: Int = 64 + ) -> RawJSONLimits { + RawJSONLimits( + maximumDocumentBytes: maximumDocumentBytes, + maximumStringBytes: maximumStringBytes, + maximumDepth: maximumDepth, + maximumContainerEntries: maximumContainerEntries + ) + } +} diff --git a/Tests/MCPTests/RawMethodHandlerTests.swift b/Tests/MCPTests/RawMethodHandlerTests.swift new file mode 100644 index 00000000..289ea6fb --- /dev/null +++ b/Tests/MCPTests/RawMethodHandlerTests.swift @@ -0,0 +1,485 @@ +import Foundation +import Testing + +@testable import MCP + +@Suite("Raw-aware method handler tests") +struct RawMethodHandlerTests { + private struct InspectRaw: MCP.Method { + struct Parameters: Codable, Hashable, Sendable { + let label: String + let payload: Value + } + + struct Result: Codable, Hashable, Sendable { + let accepted: Bool + } + + static let name = "test/inspect-raw" + } + + private struct LegacyEcho: MCP.Method { + struct Parameters: Codable, Hashable, Sendable { + let label: String + let payload: Value? + } + + struct Result: Codable, Hashable, Sendable { + let label: String + } + + static let name = "test/legacy-echo" + } + + private struct CanonicalRaw: MCP.Method { + struct Parameters: Codable, Hashable, Sendable { + let value: String + } + + struct Result: Codable, Hashable, Sendable { + let accepted: Bool + } + + static let name = "test/é" + } + + private actor Capture { + var rawRequests: [RawRequestContext] = [] + var typedPayloadWasData = false + var legacyLabels: [String] = [] + var legacyPayloadWasData = false + var rawHandlerCalls = 0 + + func recordRaw(_ context: RawRequestContext, typedPayloadWasData: Bool) { + rawRequests.append(context) + self.typedPayloadWasData = typedPayloadWasData + rawHandlerCalls += 1 + } + + func recordLegacy(_ label: String, payloadWasData: Bool) { + legacyLabels.append(label) + legacyPayloadWasData = payloadWasData + } + } + + @Test("Single requests expose exact params beside the typed request") + func singleRawAwareHandler() async throws { + let transport = MockTransport() + let server = Server(name: "raw-test", version: "1") + let capture = Capture() + + await server.withMethodHandler(InspectRaw.self) { request, rawContext in + await capture.recordRaw( + rawContext, + typedPayloadWasData: request.params.payload.dataValue != nil + ) + return .init(accepted: true) + } + try await server.start(transport: transport) + + let request = #""" + { + "jsonrpc" : "2.0", + "id" : 7, + "method" : "test/inspect-raw", + "params" : { "label" : "single", "payload" : "data:,one%20two" } + } + """# + await transport.queue(data: Data(request.utf8)) + try await waitUntil { await transport.sentData.count == 1 } + + #expect(await capture.rawHandlerCalls == 1) + #expect(await capture.typedPayloadWasData) + let context = try #require(await capture.rawRequests.first) + #expect(context.request.members.map(\.key) == ["jsonrpc", "id", "method", "params"]) + let parameters = try #require(context.uniqueParameters?.objectValue) + #expect(parameters.uniqueValue(forExactKey: "payload") == .string("data:,one%20two")) + + await server.stop() + } + + @Test("Raw-aware routing matches canonical Swift method identity") + func canonicalMethodIdentity() async throws { + let transport = MockTransport() + let server = Server(name: "raw-test", version: "1") + + actor State { + var called = false + func markCalled() { called = true } + } + let state = State() + await server.withMethodHandler(CanonicalRaw.self) { request, rawContext in + #expect(request.params.value == "ok") + #expect(rawContext.uniqueParameters?.objectValue != nil) + await state.markCalled() + return .init(accepted: true) + } + try await server.start(transport: transport) + + let request = #"{"jsonrpc":"2.0","id":9,"method":"test/e\u0301","params":{"value":"ok"}}"# + await transport.queue(data: Data(request.utf8)) + try await waitUntil { await transport.sentData.count == 1 } + + #expect(await state.called) + await server.stop() + } + + @Test("Mixed raw-aware and legacy handlers preserve batch item identity") + func mixedBatchHandlers() async throws { + let transport = MockTransport() + let limits = RawJSONLimits( + maximumDocumentBytes: 512, + maximumStringBytes: 32, + maximumDepth: 8, + maximumContainerEntries: 64 + ) + let server = Server(name: "raw-test", version: "1", rawJSONLimits: limits) + let capture = Capture() + let legacyLabel = String(repeating: "legacy", count: 20) + + await server.withMethodHandler(InspectRaw.self) { request, rawContext in + let payloadWasData: Bool + if case .data = request.params.payload { + payloadWasData = true + } else { + payloadWasData = false + } + await capture.recordRaw(rawContext, typedPayloadWasData: payloadWasData) + return .init(accepted: true) + } + await server.withMethodHandler(LegacyEcho.self) { parameters in + await capture.recordLegacy( + parameters.label, + payloadWasData: parameters.payload?.dataValue != nil + ) + return .init(label: parameters.label) + } + try await server.start(transport: transport) + + let batch = #""" + [ + { + "jsonrpc":"2.0","id":41,"method":"test/inspect-raw", + "params":{"label":"first","payload":"d\u0061ta:,A%20brief%20note","é":1,"e\u0301":2} + }, + { + "jsonrpc":"2.0","id":42,"method":"test/legacy-echo", + "params":{"label":"\#(legacyLabel)","payload":"data:text/plain;base64,SGVsbG8="} + } + ] + """# + await transport.queue(data: Data(batch.utf8)) + try await waitUntil { await transport.sentData.count == 1 } + + #expect(await capture.rawHandlerCalls == 1) + #expect(await capture.typedPayloadWasData) + #expect(await capture.legacyLabels == [legacyLabel]) + #expect(await capture.legacyPayloadWasData) + + let contexts = await capture.rawRequests + let context = try #require(contexts.first) + #expect(context.request.uniqueValue(forExactKey: "id") == .number("41")) + let rawParameters = try #require(context.parameters.first?.objectValue) + #expect( + rawParameters.uniqueValue(forExactKey: "payload") + == .string("data:,A%20brief%20note") + ) + #expect(rawParameters.values(forExactKey: "é") == [.number("1")]) + #expect(rawParameters.values(forExactKey: "e\u{301}") == [.number("2")]) + + let responseData = try #require(await transport.sentData.first) + let responses = try JSONDecoder().decode([AnyResponse].self, from: responseData) + #expect(responses.map(\.id) == [41, 42]) + + await server.stop() + } + + @Test("Oversized raw strings fail closed before the handler") + func oversizedStringRejected() async throws { + let limits = RawJSONLimits( + maximumDocumentBytes: 4096, + maximumStringBytes: 32, + maximumDepth: 8, + maximumContainerEntries: 64 + ) + let transport = MockTransport() + let server = Server(name: "raw-test", version: "1", rawJSONLimits: limits) + let capture = Capture() + + await server.withMethodHandler(InspectRaw.self) { request, rawContext in + await capture.recordRaw(rawContext, typedPayloadWasData: request.params.payload.dataValue != nil) + return .init(accepted: true) + } + try await server.start(transport: transport) + + let oversized = String(repeating: "A", count: 256) + let request = """ + { + "jsonrpc":"2.0","id":99,"method":"test/inspect-raw", + "params":{"label":"large","payload":"data:text/plain;base64,\(oversized)"} + } + """ + await transport.queue(data: Data(request.utf8)) + try await waitUntil { await transport.sentData.count == 1 } + + #expect(await capture.rawHandlerCalls == 0) + let responseData = try #require(await transport.sentData.first) + let response = try JSONDecoder().decode(AnyResponse.self, from: responseData) + guard case .failure(let error) = response.result else { + Issue.record("Expected a parse failure response") + await server.stop() + return + } + #expect(error.code == -32700) + + await server.stop() + } + + @Test("Oversized raw batch items retain IDs and do not suppress siblings") + func oversizedBatchItem() async throws { + let limits = RawJSONLimits( + maximumDocumentBytes: 4096, + maximumStringBytes: 32, + maximumDepth: 8, + maximumContainerEntries: 64 + ) + let transport = MockTransport() + let server = Server(name: "raw-test", version: "1", rawJSONLimits: limits) + let capture = Capture() + + await server.withMethodHandler(InspectRaw.self) { request, rawContext in + await capture.recordRaw( + rawContext, + typedPayloadWasData: request.params.payload.dataValue != nil + ) + return .init(accepted: true) + } + await server.withMethodHandler(LegacyEcho.self) { parameters in + await capture.recordLegacy( + parameters.label, + payloadWasData: parameters.payload?.dataValue != nil + ) + return .init(label: parameters.label) + } + try await server.start(transport: transport) + + let oversized = String(repeating: "A", count: 256) + let batch = """ + [ + {"jsonrpc":"2.0","id":99,"method":"test/inspect-raw", + "params":{"label":"large","payload":"\(oversized)"}}, + {"jsonrpc":"2.0","id":100,"method":"test/legacy-echo", + "params":{"label":"valid"}} + ] + """ + await transport.queue(data: Data(batch.utf8)) + try await waitUntil { await transport.sentData.count == 1 } + + let responseData = try #require(await transport.sentData.first) + let responses = try JSONDecoder().decode([AnyResponse].self, from: responseData) + #expect(responses.map(\.id) == [99, 100]) + guard case .failure(let error) = responses[0].result else { + Issue.record("Expected a bounded raw-item failure") + await server.stop() + return + } + #expect(error.code == -32700) + #expect(await capture.rawHandlerCalls == 0) + #expect(await capture.legacyLabels == ["valid"]) + + await server.stop() + } + + @Test("Malformed batch items are client errors") + func malformedBatchItem() async throws { + let transport = MockTransport() + let server = Server(name: "raw-test", version: "1") + let capture = Capture() + await server.withMethodHandler(LegacyEcho.self) { parameters in + await capture.recordLegacy( + parameters.label, + payloadWasData: parameters.payload?.dataValue != nil + ) + return .init(label: parameters.label) + } + try await server.start(transport: transport) + + let batch = #""" + [ + {"jsonrpc":"2.0","id":12,"method":42}, + {"jsonrpc":"2.0","id":13,"method":"test/legacy-echo","params":{"label":"valid"}} + ] + """# + await transport.queue(data: Data(batch.utf8)) + try await waitUntil { await transport.sentData.count == 1 } + + let responseData = try #require(await transport.sentData.first) + let responses = try JSONDecoder().decode([AnyResponse].self, from: responseData) + #expect(responses.map(\.id) == [12, 13]) + guard case .failure(let error) = responses[0].result else { + Issue.record("Expected an invalid-request response") + await server.stop() + return + } + #expect(error.code == -32600) + #expect(await capture.legacyLabels == ["valid"]) + + await server.stop() + } + + @Test("Malformed batch JSON remains a parse error") + func malformedBatchJSON() async throws { + let transport = MockTransport() + let server = Server(name: "raw-test", version: "1") + try await server.start(transport: transport) + + await transport.queue(data: Data(#"[{"jsonrpc":"2.0","id":12"#.utf8)) + try await waitUntil { await transport.sentData.count == 1 } + + let responseData = try #require(await transport.sentData.first) + let response = try JSONDecoder().decode(AnyResponse.self, from: responseData) + guard case .failure(let error) = response.result else { + Issue.record("Expected a parse-error response") + await server.stop() + return + } + #expect(error.code == -32700) + + await server.stop() + } + + @Test("Malformed legacy requests retain their request ID") + func malformedLegacyRequestID() async throws { + let transport = MockTransport() + let server = Server(name: "raw-test", version: "1") + try await server.start(transport: transport) + + let request = #"{"jsonrpc":"2.0","id":7,"method":42}"# + await transport.queue(data: Data(request.utf8)) + try await waitUntil { await transport.sentData.count == 1 } + + let responseData = try #require(await transport.sentData.first) + let response = try JSONDecoder().decode(AnyResponse.self, from: responseData) + #expect(response.id == 7) + guard case .failure(let error) = response.result else { + Issue.record("Expected a parse-error response") + await server.stop() + return + } + #expect(error.code == -32700) + + await server.stop() + } + + @Test("Ambiguous raw-aware method members are rejected") + func duplicateMethodMembers() async throws { + let transport = MockTransport() + let server = Server(name: "raw-test", version: "1") + let capture = Capture() + await server.withMethodHandler(InspectRaw.self) { request, rawContext in + await capture.recordRaw( + rawContext, + typedPayloadWasData: request.params.payload.dataValue != nil + ) + return .init(accepted: true) + } + await server.withMethodHandler(LegacyEcho.self) { parameters in + await capture.recordLegacy( + parameters.label, + payloadWasData: parameters.payload?.dataValue != nil + ) + return .init(label: parameters.label) + } + try await server.start(transport: transport) + + let request = #""" + { + "jsonrpc":"2.0","id":88, + "method":"test/legacy-echo","method":"test/inspect-raw", + "params":{"label":"ambiguous","payload":"plain"} + } + """# + await transport.queue(data: Data(request.utf8)) + try await waitUntil { await transport.sentData.count == 1 } + + let responseData = try #require(await transport.sentData.first) + let response = try JSONDecoder().decode(AnyResponse.self, from: responseData) + #expect(response.id == 88) + guard case .failure(let error) = response.result else { + Issue.record("Expected an invalid-request response") + await server.stop() + return + } + #expect(error.code == -32600) + #expect(await capture.rawHandlerCalls == 0) + #expect(await capture.legacyLabels.isEmpty) + + await server.stop() + } + + @Test("Raw limits do not change legacy handler traffic") + func legacyTrafficIgnoresRawLimits() async throws { + let limits = RawJSONLimits( + maximumDocumentBytes: 1, + maximumStringBytes: 1, + maximumDepth: 0, + maximumContainerEntries: 0 + ) + let transport = MockTransport() + let server = Server(name: "raw-test", version: "1", rawJSONLimits: limits) + let capture = Capture() + + await server.withMethodHandler(InspectRaw.self) { request, rawContext in + await capture.recordRaw( + rawContext, + typedPayloadWasData: request.params.payload.dataValue != nil + ) + return .init(accepted: true) + } + await server.withMethodHandler(LegacyEcho.self) { parameters in + await capture.recordLegacy( + parameters.label, + payloadWasData: parameters.payload?.dataValue != nil + ) + return .init(label: parameters.label) + } + try await server.start(transport: transport) + + let matchingNotification = #""" + { + "jsonrpc":"2.0","method":"test/inspect-raw", + "params":{"payload":"this notification is intentionally larger than raw limits"} + } + """# + await transport.queue(data: Data(matchingNotification.utf8)) + + let request = #""" + { + "jsonrpc":"2.0","id":77,"method":"test/legacy-echo", + "params":{"label":"unbounded-legacy","payload":"data:text/plain;base64,SGVsbG8="} + } + """# + await transport.queue(data: Data(request.utf8)) + try await waitUntil { await transport.sentData.count == 1 } + + #expect(await capture.legacyLabels == ["unbounded-legacy"]) + #expect(await capture.legacyPayloadWasData) + #expect(await capture.rawHandlerCalls == 0) + + await server.stop() + } + + private func waitUntil( + timeout: Duration = .seconds(2), + condition: @escaping @Sendable () async -> Bool + ) async throws { + let clock = ContinuousClock() + let deadline = clock.now.advanced(by: timeout) + while !(await condition()) { + guard clock.now < deadline else { + throw MCPError.internalError("Timed out waiting for handler") + } + try await Task.sleep(for: .milliseconds(10)) + } + } +}