diff --git a/docs/concepts/transports/transports.md b/docs/concepts/transports/transports.md index 48d19986b..7bb689120 100644 --- a/docs/concepts/transports/transports.md +++ b/docs/concepts/transports/transports.md @@ -110,6 +110,16 @@ builder.Services.AddMcpServer() await builder.Build().RunAsync(); ``` +By default, the stdio server accepts a message of any size and buffers each line until its newline arrives. To bound the memory used to read one message, set : + +```csharp +builder.Services.AddMcpServer() + .WithStdioServerTransport(options => options.MaxMessageSize = 32 * 1024 * 1024) + .WithTools(); +``` + +A request over the limit is discarded without being buffered whole and is answered with a JSON-RPC `-32600` (invalid request) error for its id. A response over the limit, such as a large sampling result, fails the server's pending request instead. The same option applies to . + ### Streamable HTTP transport The [Streamable HTTP] transport uses HTTP for bidirectional communication with optional streaming. This is the recommended transport for remote servers. diff --git a/src/ModelContextProtocol.Core/Protocol/JsonRpcFrameScanner.cs b/src/ModelContextProtocol.Core/Protocol/JsonRpcFrameScanner.cs new file mode 100644 index 000000000..91e1062e9 --- /dev/null +++ b/src/ModelContextProtocol.Core/Protocol/JsonRpcFrameScanner.cs @@ -0,0 +1,481 @@ +using System.Buffers; +using System.Text.Json; + +namespace ModelContextProtocol.Protocol; + +/// +/// Inspects the top-level members of a JSON-RPC message that is fed in pieces, in constant memory. +/// +/// +/// +/// A transport uses this for a message it does not keep whole, such as one that exceeds the maximum +/// message size. The scanner records the top-level id and whether method, result, +/// or error is present, so the transport can tell a request it should answer from a response to +/// one of its own requests. The position of id is not fixed: this SDK writes it after params, +/// so it can follow megabytes of argument data. +/// +/// +/// This is best effort, not a validator. It tracks only nesting, strings, and escapes. Anything it +/// does not expect at the top level ends the scan, keeping what was found so far. The first top-level +/// id wins. A key longer than bytes is treated as an unknown key, and +/// an id value longer than bytes is not recovered. +/// +/// +internal sealed class JsonRpcFrameScanner +{ + /// A key longer than this, quotes included, is not one of the keys of interest, even escaped. + private const int MaxKeyLength = 64; + + /// An id value longer than this is not recovered. + private const int MaxIdLength = 1024; + + private readonly byte[] _key = new byte[MaxKeyLength]; + private readonly byte[] _idValue = new byte[MaxIdLength]; + + private State _state; + private int _keyLength; + private bool _keyOverflow; + private Key _currentKey; + private int _idValueLength; + private bool _escape; + private bool _inString; + private int _depth; + private RequestId _id; + private bool _hasId; + + private enum State + { + BeforeRoot, + ExpectKey, + InKey, + ExpectColon, + ExpectValue, + InIdString, + InIdLiteral, + SkipValue, + SkipLiteral, + AfterValue, + Done, + } + + private enum Key + { + Other, + Id, + Method, + Result, + Error, + } + + /// Gets a value indicating whether the scan has ended; further input is ignored. + public bool IsComplete => _state == State.Done; + + /// Gets a value indicating whether a top-level method member was seen. + public bool HasMethod { get; private set; } + + /// Gets a value indicating whether a top-level result or error member was seen. + public bool HasResultOrError { get; private set; } + + /// Starts a new scan. + public void Reset() + { + _state = State.BeforeRoot; + _keyLength = 0; + _keyOverflow = false; + _currentKey = Key.Other; + _idValueLength = 0; + _escape = false; + _inString = false; + _depth = 0; + _id = default; + _hasId = false; + HasMethod = false; + HasResultOrError = false; + } + + /// Feeds the next bytes of the message. + public void Feed(in ReadOnlySequence bytes) + { + foreach (ReadOnlyMemory segment in bytes) + { + if (IsComplete) + { + return; + } + + Feed(segment.Span); + } + } + + /// Feeds the next bytes of the message. + public void Feed(ReadOnlySpan bytes) + { + int i = 0; + while (i < bytes.Length && _state != State.Done) + { + byte b = bytes[i]; + switch (_state) + { + case State.BeforeRoot: + if (b == (byte)'{') + { + _state = State.ExpectKey; + } + else if (!IsWhiteSpace(b)) + { + Stop(); + } + + i++; + break; + + case State.ExpectKey: + if (b == (byte)'"') + { + _key[0] = b; + _keyLength = 1; + _keyOverflow = false; + _escape = false; + _state = State.InKey; + } + else if (!IsWhiteSpace(b)) + { + // '}' ends the object; anything else is not a key. + Stop(); + } + + i++; + break; + + case State.InKey: + AppendKey(b); + if (_escape) + { + _escape = false; + } + else if (b == (byte)'\\') + { + _escape = true; + } + else if (b == (byte)'"') + { + _currentKey = _keyOverflow ? Key.Other : ClassifyKey(_key.AsSpan(0, _keyLength)); + _state = State.ExpectColon; + } + + i++; + break; + + case State.ExpectColon: + if (b == (byte)':') + { + _state = State.ExpectValue; + } + else if (!IsWhiteSpace(b)) + { + Stop(); + } + + i++; + break; + + case State.ExpectValue: + i++; + if (IsWhiteSpace(b)) + { + break; + } + + switch (_currentKey) + { + case Key.Method: + HasMethod = true; + break; + case Key.Result or Key.Error: + HasResultOrError = true; + break; + case Key.Id when !_hasId: + if (b is (byte)'{' or (byte)'[') + { + // An object or array is not a valid JSON-RPC id. + Stop(); + return; + } + + _idValue[0] = b; + _idValueLength = 1; + _escape = false; + _state = b == (byte)'"' ? State.InIdString : State.InIdLiteral; + continue; + } + + _escape = false; + if (b == (byte)'"') + { + _inString = true; + _depth = 0; + _state = State.SkipValue; + } + else if (b is (byte)'{' or (byte)'[') + { + _inString = false; + _depth = 1; + _state = State.SkipValue; + } + else + { + _state = State.SkipLiteral; + } + + break; + + case State.InIdString: + if (!AppendIdValue(b)) + { + break; + } + + if (_escape) + { + _escape = false; + } + else if (b == (byte)'\\') + { + _escape = true; + } + else if (b == (byte)'"') + { + FinishId(); + } + + i++; + break; + + case State.InIdLiteral: + if (b is (byte)',' or (byte)'}' || IsWhiteSpace(b)) + { + // Leave the delimiter for AfterValue. + FinishId(); + break; + } + + if (AppendIdValue(b)) + { + i++; + } + + break; + + case State.SkipValue: + i = SkipValue(bytes, i); + break; + + case State.SkipLiteral: + if (b == (byte)',') + { + _state = State.ExpectKey; + } + else if (b == (byte)'}') + { + Stop(); + } + else if (IsWhiteSpace(b)) + { + _state = State.AfterValue; + } + + i++; + break; + + case State.AfterValue: + if (b == (byte)',') + { + _state = State.ExpectKey; + } + else if (!IsWhiteSpace(b)) + { + // '}' ends the object; anything else is malformed. + Stop(); + } + + i++; + break; + } + } + } + + /// Ends the input and returns the top-level id, if one was found. + public bool TryGetId(out RequestId id) + { + if (_state == State.InIdLiteral) + { + // A number at the very end of the input has no delimiter after it. + ParseId(); + } + + id = _id; + return _hasId; + } + + /// + /// Skips the body of a value that is not the id. Inside a string, it jumps to the next quote or + /// backslash with a vectorized search, so a multi-megabyte argument is not walked byte by byte. + /// + private int SkipValue(ReadOnlySpan bytes, int i) + { + while (i < bytes.Length) + { + if (_inString) + { + if (_escape) + { + _escape = false; + i++; + continue; + } + + int next = bytes.Slice(i).IndexOfAny((byte)'"', (byte)'\\'); + if (next < 0) + { + return bytes.Length; + } + + i += next; + if (bytes[i] == (byte)'\\') + { + _escape = true; + i++; + continue; + } + + // Closing quote. + _inString = false; + i++; + if (_depth == 0) + { + _state = State.AfterValue; + return i; + } + + continue; + } + + switch (bytes[i++]) + { + case (byte)'"': + _inString = true; + break; + case (byte)'{' or (byte)'[': + _depth++; + break; + case (byte)'}' or (byte)']': + if (--_depth == 0) + { + _state = State.AfterValue; + return i; + } + + break; + } + } + + return i; + } + + private void AppendKey(byte b) + { + if (_keyLength < MaxKeyLength) + { + _key[_keyLength++] = b; + } + else + { + _keyOverflow = true; + } + } + + private bool AppendIdValue(byte b) + { + if (_idValueLength >= MaxIdLength) + { + Stop(); + return false; + } + + _idValue[_idValueLength++] = b; + return true; + } + + private void FinishId() + { + ParseId(); + _state = State.AfterValue; + } + + private void ParseId() + { + try + { + var reader = new Utf8JsonReader(_idValue.AsSpan(0, _idValueLength), isFinalBlock: true, state: default); + if (!reader.Read()) + { + return; + } + + switch (reader.TokenType) + { + case JsonTokenType.String: + _id = new RequestId(reader.GetString()!); + _hasId = true; + break; + + case JsonTokenType.Number when reader.TryGetInt64(out long number): + _id = new RequestId(number); + _hasId = true; + break; + } + } + catch (Exception ex) when (ex is JsonException or InvalidOperationException) + { + // Not a valid string or integer, or a string with an invalid escape sequence: no id. + } + } + + private void Stop() => _state = State.Done; + + /// Classifies a quoted key, unescaping it first if needed (for example "id"). + private static Key ClassifyKey(ReadOnlySpan quotedKey) + { + if (quotedKey.IndexOf((byte)'\\') < 0) + { + return + quotedKey.SequenceEqual("\"id\""u8) ? Key.Id : + quotedKey.SequenceEqual("\"method\""u8) ? Key.Method : + quotedKey.SequenceEqual("\"result\""u8) ? Key.Result : + quotedKey.SequenceEqual("\"error\""u8) ? Key.Error : + Key.Other; + } + + try + { + var reader = new Utf8JsonReader(quotedKey, isFinalBlock: true, state: default); + if (!reader.Read() || reader.TokenType != JsonTokenType.String) + { + return Key.Other; + } + + return + reader.ValueTextEquals("id"u8) ? Key.Id : + reader.ValueTextEquals("method"u8) ? Key.Method : + reader.ValueTextEquals("result"u8) ? Key.Result : + reader.ValueTextEquals("error"u8) ? Key.Error : + Key.Other; + } + catch (JsonException) + { + return Key.Other; + } + } + + private static bool IsWhiteSpace(byte b) => b is (byte)' ' or (byte)'\t' or (byte)'\r' or (byte)'\n'; +} diff --git a/src/ModelContextProtocol.Core/Server/StdioServerTransport.cs b/src/ModelContextProtocol.Core/Server/StdioServerTransport.cs index f43426344..2b194667b 100644 --- a/src/ModelContextProtocol.Core/Server/StdioServerTransport.cs +++ b/src/ModelContextProtocol.Core/Server/StdioServerTransport.cs @@ -15,7 +15,19 @@ public sealed class StdioServerTransport : StreamServerTransport /// Optional logger factory used for logging employed by the transport. /// is or contains a null name. public StdioServerTransport(McpServerOptions serverOptions, ILoggerFactory? loggerFactory = null) - : this(GetServerName(serverOptions), loggerFactory: loggerFactory) + : this(GetServerName(serverOptions), loggerFactory, transportOptions: null) + { + } + + /// + /// Initializes a new instance of the class with transport options. + /// + /// The server options. + /// Optional logger factory used for logging employed by the transport. + /// Optional transport options, such as . + /// is or contains a null name. + public StdioServerTransport(McpServerOptions serverOptions, ILoggerFactory? loggerFactory, StreamServerTransportOptions? transportOptions) + : this(GetServerName(serverOptions), loggerFactory, transportOptions) { } @@ -26,10 +38,23 @@ public StdioServerTransport(McpServerOptions serverOptions, ILoggerFactory? logg /// Optional logger factory used for logging employed by the transport. /// is . public StdioServerTransport(string serverName, ILoggerFactory? loggerFactory = null) + : this(serverName, loggerFactory, transportOptions: null) + { + } + + /// + /// Initializes a new instance of the class with transport options. + /// + /// The name of the server. + /// Optional logger factory used for logging employed by the transport. + /// Optional transport options, such as . + /// is . + public StdioServerTransport(string serverName, ILoggerFactory? loggerFactory, StreamServerTransportOptions? transportOptions) : base(new CancellableStdinStream(Console.OpenStandardInput()), new BufferedStream(Console.OpenStandardOutput()), serverName ?? throw new ArgumentNullException(nameof(serverName)), - loggerFactory) + loggerFactory, + transportOptions) { } diff --git a/src/ModelContextProtocol.Core/Server/StreamServerTransport.cs b/src/ModelContextProtocol.Core/Server/StreamServerTransport.cs index cc9765066..fe2d485b1 100644 --- a/src/ModelContextProtocol.Core/Server/StreamServerTransport.cs +++ b/src/ModelContextProtocol.Core/Server/StreamServerTransport.cs @@ -1,8 +1,11 @@ using Microsoft.Extensions.Logging; using Microsoft.Extensions.Logging.Abstractions; using ModelContextProtocol.Protocol; +using System.Buffers; +using System.IO.Pipelines; using System.Text; using System.Text.Json; +using System.Text.Json.Serialization.Metadata; namespace ModelContextProtocol.Server; @@ -14,14 +17,30 @@ namespace ModelContextProtocol.Server; /// streams, allowing MCP communication with clients through various I/O channels such as network sockets, /// memory streams, or pipes. /// -public class StreamServerTransport : TransportBase +public partial class StreamServerTransport : TransportBase { + // Pooled read buffers. The 4 KiB default of StreamPipeReaderOptions makes a large message span too + // many segments; 64 KiB keeps the segment count low without holding much memory for small messages. + private const int ReadBufferSize = 64 * 1024; + private const int MinimumReadSize = 16 * 1024; + private static readonly byte[] s_newlineBytes = "\n"u8.ToArray(); + private static readonly JsonReaderOptions s_readerOptions = new() + { + // The same reader settings JsonSerializer.Deserialize(string, options) would use. + AllowTrailingCommas = McpJsonUtilities.ProtocolOptions.AllowTrailingCommas, + CommentHandling = McpJsonUtilities.ProtocolOptions.ReadCommentHandling, + MaxDepth = McpJsonUtilities.ProtocolOptions.MaxDepth, + }; + private readonly ILogger _logger; - private readonly TextReader _inputReader; + private readonly Stream _inputStream; + private readonly PipeReader _inputReader; private readonly Stream _outputStream; + private readonly long? _maxMessageSize; + private readonly JsonTypeInfo _messageTypeInfo = McpJsonUtilities.ProtocolOptions.GetTypeInfo(typeof(JsonRpcMessage)); private readonly SemaphoreSlim _sendLock = new(1, 1); private readonly CancellationTokenSource _shutdownCts = new(); @@ -29,6 +48,13 @@ public class StreamServerTransport : TransportBase private readonly Task _readLoopCompleted; private int _disposed = 0; + // State touched only by the read loop. + private long _searchedLength; + private bool _atStreamStart = true; + private bool _discarding; + private long _discardedLength; + private JsonRpcFrameScanner? _discardScanner; + /// /// Initializes a new instance of the class with explicit input/output streams. /// @@ -38,6 +64,26 @@ public class StreamServerTransport : TransportBase /// Optional logger factory used for logging employed by the transport. /// or is . public StreamServerTransport(Stream inputStream, Stream outputStream, string? serverName = null, ILoggerFactory? loggerFactory = null) + : this(inputStream, outputStream, serverName, loggerFactory, transportOptions: null) + { + } + + /// + /// Initializes a new instance of the class with explicit input/output streams + /// and transport options. + /// + /// The input to use as standard input. + /// The output to use as standard output. + /// Optional name of the server, used for diagnostic purposes, like logging. + /// Optional logger factory used for logging employed by the transport. + /// Optional transport options, such as . + /// or is . + public StreamServerTransport( + Stream inputStream, + Stream outputStream, + string? serverName, + ILoggerFactory? loggerFactory, + StreamServerTransportOptions? transportOptions) : base(serverName is not null ? $"Server (stream) ({serverName})" : "Server (stream)", loggerFactory) { Throw.IfNull(inputStream); @@ -45,12 +91,13 @@ public StreamServerTransport(Stream inputStream, Stream outputStream, string? se _logger = loggerFactory?.CreateLogger(GetType()) ?? NullLogger.Instance; -#if NET - _inputReader = new StreamReader(inputStream, Encoding.UTF8); -#else - _inputReader = new CancellableStreamReader(inputStream, Encoding.UTF8); -#endif + _inputStream = inputStream; + _inputReader = PipeReader.Create(inputStream, new StreamPipeReaderOptions( + bufferSize: ReadBufferSize, + minimumReadSize: MinimumReadSize, + leaveOpen: true)); _outputStream = outputStream; + _maxMessageSize = transportOptions?.MaxMessageSize; SetConnected(); _readLoopCompleted = Task.Run(ReadMessagesAsync, _shutdownCts.Token); @@ -95,75 +142,46 @@ private async Task ReadMessagesAsync() { CancellationToken shutdownToken = _shutdownCts.Token; Exception? error = null; + + // Messages produced from one read, in input order. Each is either handed to the session or, for an + // error reply to a request that could not be read, sent back to the peer. + List<(JsonRpcMessage Message, bool SendToPeer)> messages = []; try { LogTransportEnteringReadMessagesLoop(Name); while (!shutdownToken.IsCancellationRequested) { - var line = await _inputReader.ReadLineAsync(shutdownToken).ConfigureAwait(false); - if (string.IsNullOrWhiteSpace(line)) + ReadResult result = await _inputReader.ReadAsync(shutdownToken).ConfigureAwait(false); + if (result.IsCanceled) { - if (line is null) - { - LogTransportEndOfStream(Name); - break; - } - - continue; + break; } - LogTransportReceivedMessageSensitive(Name, line); + ReadOnlySequence buffer = result.Buffer; + SequencePosition consumed = ProcessBuffer(buffer, result.IsCompleted, messages); + _inputReader.AdvanceTo(consumed, buffer.End); - try + // Every message was deserialized into its own objects, so none of them refers to the pipe's + // buffers any more and they can be dispatched after AdvanceTo. + foreach ((JsonRpcMessage message, bool sendToPeer) in messages) { - if (JsonSerializer.Deserialize(line, McpJsonUtilities.ProtocolOptions.GetTypeInfo(typeof(JsonRpcMessage))) is JsonRpcMessage message) + if (sendToPeer) { - await WriteMessageAsync(message, shutdownToken).ConfigureAwait(false); + await TrySendErrorAsync(message, shutdownToken).ConfigureAwait(false); } else { - LogTransportMessageParseUnexpectedTypeSensitive(Name, line); + await WriteMessageAsync(message, shutdownToken).ConfigureAwait(false); } } - catch (JsonException ex) - { - if (Logger.IsEnabled(LogLevel.Trace)) - { - LogTransportMessageParseFailedSensitive(Name, line, ex); - } - else - { - LogTransportMessageParseFailed(Name, ex); - } - // Deserializing the full message failed, for example because the params object was nested - // more deeply than the JSON reader's MaxDepth allows. If the message still carried a request - // id, reply with a JSON-RPC parse error using that id so the caller's pending request - // completes instead of hanging until it times out. If no id can be recovered, the message - // was either a notification or too malformed to correlate, so we just continue reading. - if (TryRecoverRequestId(line, out RequestId id)) - { - var errorResponse = new JsonRpcError - { - Id = id, - Error = new JsonRpcErrorDetail - { - Code = (int)McpErrorCode.ParseError, - Message = "Failed to parse the JSON-RPC request.", - }, - }; - - try - { - await SendMessageAsync(errorResponse, shutdownToken).ConfigureAwait(false); - } - catch (Exception sendEx) when (sendEx is not OperationCanceledException) - { - // Swallow so a failed error-send does not tear down the read loop. No logging - // here because SendMessageAsync already logs send failures before it throws. - } - } + messages.Clear(); + + if (result.IsCompleted) + { + LogTransportEndOfStream(Name); + break; } } } @@ -178,10 +196,226 @@ private async Task ReadMessagesAsync() } finally { + try + { + await _inputReader.CompleteAsync().ConfigureAwait(false); + } + catch (Exception ex) + { + LogTransportCleanupReadTaskFailed(Name, ex); + } + SetDisconnected(error); } } + /// + /// Splits the buffered input into newline-delimited messages and processes each complete one. + /// + /// The position up to which the buffer was consumed. + /// + /// Bytes of an incomplete line that were already searched for a newline are not searched again on + /// the next read. A line that grows past the maximum message size is not kept: its bytes are fed to a + /// and released until the newline that ends it arrives. + /// + private SequencePosition ProcessBuffer( + ReadOnlySequence buffer, + bool isCompleted, + List<(JsonRpcMessage Message, bool SendToPeer)> messages) + { + while (buffer.Slice(_searchedLength).PositionOf((byte)'\n') is SequencePosition newline) + { + ReadOnlySequence line = buffer.Slice(0, newline); + buffer = buffer.Slice(buffer.GetPosition(1, newline)); + _searchedLength = 0; + + if (_discarding) + { + // The end of an oversized message whose start has already been released. + ContinueDiscarding(line); + EndDiscarding(messages); + } + else if (line.Length > _maxMessageSize) + { + // The whole oversized message arrived in this read. + StartDiscarding(line); + EndDiscarding(messages); + } + else + { + ProcessLine(line, messages); + } + } + + // What remains is an incomplete line. + if (_discarding) + { + ContinueDiscarding(buffer); + if (isCompleted) + { + EndDiscarding(messages); + } + + return buffer.End; + } + + if (buffer.Length > _maxMessageSize) + { + StartDiscarding(buffer); + if (isCompleted) + { + EndDiscarding(messages); + } + + return buffer.End; + } + + if (isCompleted) + { + // A final line without a trailing newline is still a message, as with TextReader.ReadLine. + if (!buffer.IsEmpty) + { + ProcessLine(buffer, messages); + } + + return buffer.End; + } + + _searchedLength = buffer.Length; + return buffer.Start; + } + + private void ProcessLine(ReadOnlySequence line, List<(JsonRpcMessage Message, bool SendToPeer)> messages) + { + if (_atStreamStart) + { + // Skip a UTF-8 byte order mark at the start of the stream, as StreamReader did. + _atStreamStart = false; + if (StartsWithUtf8Bom(line)) + { + line = line.Slice(3); + } + } + + if (IsWhiteSpace(line)) + { + return; + } + + if (Logger.IsEnabled(LogLevel.Trace)) + { + LogTransportReceivedMessageSensitive(Name, GetUtf8String(line)); + } + + try + { + var reader = new Utf8JsonReader(line, s_readerOptions); + if (JsonSerializer.Deserialize(ref reader, _messageTypeInfo) is JsonRpcMessage message) + { + // Deserialize(string) rejects anything after the value; keep that behavior. + if (reader.Read()) + { + throw new JsonException("The JSON-RPC message is followed by unexpected data."); + } + + messages.Add((message, false)); + } + else if (Logger.IsEnabled(LogLevel.Trace)) + { + LogTransportMessageParseUnexpectedTypeSensitive(Name, GetUtf8String(line)); + } + } + catch (JsonException ex) + { + if (Logger.IsEnabled(LogLevel.Trace)) + { + LogTransportMessageParseFailedSensitive(Name, GetUtf8String(line), ex); + } + else + { + LogTransportMessageParseFailed(Name, ex); + } + + // Deserializing the full message failed, for example because the params object was nested + // more deeply than the JSON reader's MaxDepth allows. If the message still carried a request + // id, reply with a JSON-RPC parse error using that id so the caller's pending request + // completes instead of hanging until it times out. If no id can be recovered, the message + // was either a notification or too malformed to correlate, so we just continue reading. + if (TryRecoverRequestId(line, out RequestId id)) + { + messages.Add((CreateError(id, McpErrorCode.ParseError, "Failed to parse the JSON-RPC request."), true)); + } + } + } + + private void StartDiscarding(in ReadOnlySequence head) + { + // Nothing of the line is kept in the pipe from here on, so there is nothing already searched. + _searchedLength = 0; + _atStreamStart = false; + _discarding = true; + _discardedLength = 0; + _discardScanner ??= new JsonRpcFrameScanner(); + _discardScanner.Reset(); + ContinueDiscarding(head); + } + + private void ContinueDiscarding(in ReadOnlySequence bytes) + { + _discardedLength += bytes.Length; + _discardScanner!.Feed(bytes); + } + + private void EndDiscarding(List<(JsonRpcMessage Message, bool SendToPeer)> messages) + { + _discarding = false; + LogTransportMessageTooLarge(Name, _discardedLength, _maxMessageSize.GetValueOrDefault()); + + JsonRpcFrameScanner scanner = _discardScanner!; + if (!scanner.TryGetId(out RequestId id)) + { + // A notification, or a message too malformed to correlate. + return; + } + + if (scanner.HasMethod) + { + // A request: tell the peer it was rejected, so its pending request completes. + messages.Add((CreateError(id, McpErrorCode.InvalidRequest, + $"The message exceeds the maximum message size of {_maxMessageSize} bytes."), true)); + } + else if (scanner.HasResultOrError) + { + // A response to a request this server sent, carrying this server's own request id. Sending + // an error to the peer would make no sense, so fail the pending request locally instead. + messages.Add((CreateError(id, McpErrorCode.InternalError, + $"The response exceeds the maximum message size of {_maxMessageSize} bytes and was discarded."), false)); + } + } + + private async Task TrySendErrorAsync(JsonRpcMessage errorResponse, CancellationToken cancellationToken) + { + try + { + await SendMessageAsync(errorResponse, cancellationToken).ConfigureAwait(false); + } + catch (Exception sendEx) when (sendEx is not OperationCanceledException) + { + // Swallow so a failed error-send does not tear down the read loop. No logging + // here because SendMessageAsync already logs send failures before it throws. + } + } + + private static JsonRpcError CreateError(RequestId id, McpErrorCode code, string message) => new() + { + Id = id, + Error = new JsonRpcErrorDetail + { + Code = (int)code, + Message = message, + }, + }; + /// /// Attempts to recover the JSON-RPC request id from a line that failed full deserialization. /// @@ -189,14 +423,13 @@ private async Task ReadMessagesAsync() /// This walks only the top-level object looking for an "id" property and skips every other value, /// using a large reader depth so a deeply nested "params" value cannot make recovery itself fail. /// - private static bool TryRecoverRequestId(string line, out RequestId id) + private static bool TryRecoverRequestId(in ReadOnlySequence line, out RequestId id) { id = default; try { - byte[] utf8 = Encoding.UTF8.GetBytes(line); - var reader = new Utf8JsonReader(utf8, new JsonReaderOptions + var reader = new Utf8JsonReader(line, new JsonReaderOptions { // Use the maximum reader depth so that an over-nested "params" value cannot make id // recovery throw for the same reason the original parse did. @@ -261,6 +494,41 @@ private static bool TryRecoverRequestId(string line, out RequestId id) return false; } + private static bool StartsWithUtf8Bom(in ReadOnlySequence line) + { + if (line.Length < 3) + { + return false; + } + + Span start = stackalloc byte[3]; + line.Slice(0, 3).CopyTo(start); + return start.SequenceEqual(""u8); + } + + private static bool IsWhiteSpace(in ReadOnlySequence line) + { + foreach (ReadOnlyMemory segment in line) + { + foreach (byte b in segment.Span) + { + if (b is not ((byte)' ' or (byte)'\t' or (byte)'\r' or (byte)'\n')) + { + return false; + } + } + } + + return true; + } + + private static string GetUtf8String(in ReadOnlySequence bytes) => +#if NET + Encoding.UTF8.GetString(bytes); +#else + Encoding.UTF8.GetString(bytes.ToArray()); +#endif + /// public override async ValueTask DisposeAsync() { @@ -279,7 +547,9 @@ public override async ValueTask DisposeAsync() // Dispose of stdin/out. Cancellation may not be able to wake up operations // synchronously blocked in a syscall; we need to forcefully close the handle / file descriptor. - _inputReader?.Dispose(); + // The read loop completes the PipeReader itself, since a PipeReader must not be completed + // while a read on it may still be in progress. + _inputStream?.Dispose(); _outputStream?.Dispose(); // Make sure the work has quiesced. @@ -303,4 +573,7 @@ public override async ValueTask DisposeAsync() GC.SuppressFinalize(this); } + + [LoggerMessage(Level = LogLevel.Warning, Message = "{EndpointName} transport discarded a message of {MessageSize} bytes, which exceeds the maximum message size of {MaxMessageSize} bytes.")] + private partial void LogTransportMessageTooLarge(string endpointName, long messageSize, long maxMessageSize); } diff --git a/src/ModelContextProtocol.Core/Server/StreamServerTransportOptions.cs b/src/ModelContextProtocol.Core/Server/StreamServerTransportOptions.cs new file mode 100644 index 000000000..3846079db --- /dev/null +++ b/src/ModelContextProtocol.Core/Server/StreamServerTransportOptions.cs @@ -0,0 +1,46 @@ +namespace ModelContextProtocol.Server; + +/// +/// Provides options for and . +/// +public class StreamServerTransportOptions +{ + private long? _maxMessageSize; + + /// + /// Gets or sets the maximum size, in bytes, of a single incoming JSON-RPC message. + /// + /// + /// The maximum number of UTF-8 bytes in one newline-delimited message, not counting the newline, + /// or for no limit. The default is . + /// + /// The value is zero or negative. + /// + /// + /// Without a limit, the transport buffers each incoming line until its newline arrives, however + /// long the line is. With a limit, the transport stops buffering a line as soon as it exceeds the + /// limit and discards the rest of it up to the next newline, so the memory used to read it stays + /// bounded by roughly the limit. + /// + /// + /// The transport scans a discarded message for its top-level id in constant memory. If the + /// message is a request, the transport replies with a JSON-RPC + /// error for that id. If it is a response to a request that this server sent, such as a large + /// sampling result, the pending request fails locally with an . A + /// notification, or a message whose id cannot be found, is dropped and logged. + /// + /// + public long? MaxMessageSize + { + get => _maxMessageSize; + set + { + if (value <= 0) + { + throw new ArgumentOutOfRangeException(nameof(value), value, "The maximum message size must be positive."); + } + + _maxMessageSize = value; + } + } +} diff --git a/src/ModelContextProtocol/McpServerBuilderExtensions.cs b/src/ModelContextProtocol/McpServerBuilderExtensions.cs index 7e24bc5c0..4273bc4fe 100644 --- a/src/ModelContextProtocol/McpServerBuilderExtensions.cs +++ b/src/ModelContextProtocol/McpServerBuilderExtensions.cs @@ -999,13 +999,35 @@ public static IMcpServerBuilder WithStdioServerTransport(this IMcpServerBuilder builder.Services.AddSingleton(sp => { var serverOptions = sp.GetRequiredService>(); + var transportOptions = sp.GetService>(); var loggerFactory = sp.GetService(); - return new StdioServerTransport(serverOptions.Value, loggerFactory); + return new StdioServerTransport(serverOptions.Value, loggerFactory, transportOptions?.Value); }); return builder; } + /// + /// Adds a server transport that uses standard input (stdin) and standard output (stdout) for communication, + /// configured with the specified transport options. + /// + /// The builder instance. + /// Configures the transport options, such as . + /// The builder provided in . + /// or is . + /// + /// This overload behaves like and registers + /// as a configuration of . + /// + public static IMcpServerBuilder WithStdioServerTransport(this IMcpServerBuilder builder, Action configureOptions) + { + Throw.IfNull(builder); + Throw.IfNull(configureOptions); + + builder.Services.Configure(configureOptions); + return builder.WithStdioServerTransport(); + } + /// /// Adds a server transport that uses the specified input and output streams for communication. /// @@ -1029,6 +1051,41 @@ public static IMcpServerBuilder WithStreamServerTransport( return builder; } + /// + /// Adds a server transport that uses the specified input and output streams for communication, + /// configured with the specified transport options. + /// + /// The builder instance. + /// The input to use as standard input. + /// The output to use as standard output. + /// Configures the transport options, such as . + /// The builder provided in . + /// + /// , , , or is . + /// + public static IMcpServerBuilder WithStreamServerTransport( + this IMcpServerBuilder builder, + Stream inputStream, + Stream outputStream, + Action configureOptions) + { + Throw.IfNull(builder); + Throw.IfNull(inputStream); + Throw.IfNull(outputStream); + Throw.IfNull(configureOptions); + + builder.Services.Configure(configureOptions); + + AddSingleSessionServerDependencies(builder.Services); + builder.Services.AddSingleton(sp => + { + var transportOptions = sp.GetRequiredService>(); + return new StreamServerTransport(inputStream, outputStream, serverName: null, loggerFactory: null, transportOptions.Value); + }); + + return builder; + } + private static void AddSingleSessionServerDependencies(IServiceCollection services) { services.AddHostedService(); diff --git a/tests/ModelContextProtocol.Tests/Configuration/McpServerBuilderExtensionsTransportsTests.cs b/tests/ModelContextProtocol.Tests/Configuration/McpServerBuilderExtensionsTransportsTests.cs index 6af33d8d2..299acd604 100644 --- a/tests/ModelContextProtocol.Tests/Configuration/McpServerBuilderExtensionsTransportsTests.cs +++ b/tests/ModelContextProtocol.Tests/Configuration/McpServerBuilderExtensionsTransportsTests.cs @@ -1,8 +1,11 @@ using Microsoft.Extensions.DependencyInjection; using Microsoft.Extensions.Hosting; +using Microsoft.Extensions.Options; using ModelContextProtocol.Protocol; using ModelContextProtocol.Server; using System.IO.Pipelines; +using System.Text; +using System.Text.Json; namespace ModelContextProtocol.Tests.Configuration; @@ -41,4 +44,48 @@ public async Task HostExecutionShutsDownWhenSingleSessionServerExits() clientToServerPipe.Writer.Complete(); await t; } + + [Fact] + public async Task WithStdioServerTransport_With_Options_Configures_Transport_Options() + { + var services = new ServiceCollection(); + services.AddMcpServer().WithStdioServerTransport(options => options.MaxMessageSize = 1234); + + // As above, don't resolve the stdio transport itself. + Assert.Single(services, s => s.ServiceType == typeof(ITransport)); + + await using var serviceProvider = services.BuildServiceProvider(); + Assert.Equal(1234, serviceProvider.GetRequiredService>().Value.MaxMessageSize); + } + + [Fact] + public async Task WithStreamServerTransport_With_Options_Applies_MaxMessageSize() + { + Pipe clientToServerPipe = new(), serverToClientPipe = new(); + + var services = new ServiceCollection(); + services.AddMcpServer().WithStreamServerTransport( + clientToServerPipe.Reader.AsStream(), + serverToClientPipe.Writer.AsStream(), + options => options.MaxMessageSize = 64); + + await using var serviceProvider = services.BuildServiceProvider(); + _ = serviceProvider.GetRequiredService(); + + string request = $$"""{"jsonrpc":"2.0","method":"ping","params":{"padding":"{{new string('x', 100)}}"},"id":3}"""; + await clientToServerPipe.Writer.WriteAsync(Encoding.UTF8.GetBytes(request + "\n"), TestContext.Current.CancellationToken); + + using var responseReader = new StreamReader(serverToClientPipe.Reader.AsStream(), Encoding.UTF8); + var responseLine = await responseReader.ReadLineAsync( +#if NET + TestContext.Current.CancellationToken +#endif + ); + + var error = Assert.IsType(JsonSerializer.Deserialize(responseLine!, McpJsonUtilities.DefaultOptions)); + Assert.Equal(new RequestId(3), error.Id); + Assert.Equal((int)McpErrorCode.InvalidRequest, error.Error.Code); + + clientToServerPipe.Writer.Complete(); + } } diff --git a/tests/ModelContextProtocol.Tests/ModelContextProtocol.Tests.csproj b/tests/ModelContextProtocol.Tests/ModelContextProtocol.Tests.csproj index 16398bed1..9beef20f7 100644 --- a/tests/ModelContextProtocol.Tests/ModelContextProtocol.Tests.csproj +++ b/tests/ModelContextProtocol.Tests/ModelContextProtocol.Tests.csproj @@ -41,6 +41,7 @@ + diff --git a/tests/ModelContextProtocol.Tests/Protocol/JsonRpcFrameScannerTests.cs b/tests/ModelContextProtocol.Tests/Protocol/JsonRpcFrameScannerTests.cs new file mode 100644 index 000000000..c067417d8 --- /dev/null +++ b/tests/ModelContextProtocol.Tests/Protocol/JsonRpcFrameScannerTests.cs @@ -0,0 +1,141 @@ +using ModelContextProtocol.Protocol; +using System.Text; + +namespace ModelContextProtocol.Tests.Protocol; + +public class JsonRpcFrameScannerTests +{ + public static TheoryData IdCases => new() + { + { """{"jsonrpc":"2.0","id":1,"method":"ping"}""", "1" }, + { """{"jsonrpc":"2.0","method":"ping","params":{"id":99,"a":[1,{"id":98}]},"id":2}""", "2" }, + { """{"jsonrpc":"2.0","method":"ping","params":{},"id":"abc"}""", "abc" }, + { """{"jsonrpc":"2.0","method":"ping","params":{"s":"\"id\":5,}{]["},"id":"x\"y"}""", "x\"y" }, + { """{"jsonrpc":"2.0","id":"escaped-key","method":"ping"}""", "escaped-key" }, + { """{ "jsonrpc" : "2.0" , "id" : 42 , "method" : "ping" }""", "42" }, + { """{"jsonrpc":"2.0","method":"ping","id":-7}""", "-7" }, + { """{"jsonrpc":"2.0","method":"ping","id":1,"id":2}""", "1" }, + { """{"jsonrpc":"2.0","method":"ping","id":null}""", null }, + { """{"jsonrpc":"2.0","method":"ping","id":1.5}""", null }, + { """{"jsonrpc":"2.0","method":"ping","id":{"a":1}}""", null }, + { """{"jsonrpc":"2.0","method":"ping","id":[1]}""", null }, + { """{"jsonrpc":"2.0","method":"notifications/initialized"}""", null }, + { """{"jsonrpc":"2.0","params":{"id":3}}""", null }, + { """[{"id":1}]""", null }, + { "not json", null }, + { "", null }, + }; + + [Theory] + [MemberData(nameof(IdCases))] + public void Finds_Top_Level_Id(string json, string? expectedId) + { + foreach (int chunkSize in new[] { 1, 2, 7, int.MaxValue }) + { + var scanner = Scan(json, chunkSize); + + Assert.Equal(expectedId is not null, scanner.TryGetId(out RequestId id)); + if (expectedId is not null) + { + Assert.Equal(expectedId, id.ToString()); + } + } + } + + [Fact] + public void Distinguishes_Numeric_And_String_Ids() + { + Assert.True(Scan("""{"id":5}""").TryGetId(out RequestId numeric)); + Assert.Equal(new RequestId(5), numeric); + + Assert.True(Scan("""{"id":"5"}""").TryGetId(out RequestId text)); + Assert.Equal(new RequestId("5"), text); + } + + [Fact] + public void Finds_Numeric_Id_At_The_End_Of_A_Truncated_Input() + { + // TryGetId ends the input, so a number with no delimiter after it is still complete. + Assert.True(Scan("""{"method":"ping","id":12""").TryGetId(out RequestId id)); + Assert.Equal(new RequestId(12), id); + } + + [Theory] + [InlineData("""{"jsonrpc":"2.0","id":1,"method":"tools/call","params":{}}""", true, false)] + [InlineData("""{"jsonrpc":"2.0","method":"notifications/message","params":{"result":1,"error":2}}""", true, false)] + [InlineData("""{"jsonrpc":"2.0","id":1,"result":{"method":"x"}}""", false, true)] + [InlineData("""{"jsonrpc":"2.0","error":{"code":-32603,"message":"m"},"id":1}""", false, true)] + [InlineData("""{"jsonrpc":"2.0","method":"ping","id":1}""", true, false)] + [InlineData("""{"jsonrpc":"2.0","id":1}""", false, false)] + public void Detects_Method_And_Result_Or_Error(string json, bool hasMethod, bool hasResultOrError) + { + foreach (int chunkSize in new[] { 1, 3, int.MaxValue }) + { + var scanner = Scan(json, chunkSize); + + Assert.Equal(hasMethod, scanner.HasMethod); + Assert.Equal(hasResultOrError, scanner.HasResultOrError); + } + } + + [Fact] + public void Scans_A_Large_Value_And_Finds_The_Id_After_It() + { + string value = string.Concat(Enumerable.Repeat("""a\\b\"c}{][ """, 100_000)); + string json = $$$"""{"jsonrpc":"2.0","method":"tools/call","params":{"arguments":{"data":"{{{value}}}"}},"id":"after"}"""; + + var scanner = Scan(json, chunkSize: 4096); + + Assert.True(scanner.IsComplete); + Assert.True(scanner.HasMethod); + Assert.True(scanner.TryGetId(out RequestId id)); + Assert.Equal("after", id.ToString()); + } + + [Fact] + public void Ignores_Input_After_The_Top_Level_Object_Ends() + { + var scanner = Scan("""{"method":"ping"} {"id":1}"""); + + Assert.True(scanner.IsComplete); + Assert.False(scanner.TryGetId(out _)); + } + + [Fact] + public void Ignores_An_Overly_Long_Id() + { + string json = $$$"""{"method":"ping","id":"{{{new string('a', 2048)}}}"}"""; + + Assert.False(Scan(json).TryGetId(out _)); + } + + [Fact] + public void Reset_Starts_A_New_Scan() + { + var scanner = new JsonRpcFrameScanner(); + scanner.Reset(); + scanner.Feed(Encoding.UTF8.GetBytes("""{"id":1,"method":"ping"}""")); + Assert.True(scanner.HasMethod); + + scanner.Reset(); + scanner.Feed(Encoding.UTF8.GetBytes("""{"result":{},"id":2}""")); + + Assert.False(scanner.HasMethod); + Assert.True(scanner.HasResultOrError); + Assert.True(scanner.TryGetId(out RequestId id)); + Assert.Equal(new RequestId(2), id); + } + + private static JsonRpcFrameScanner Scan(string json, int chunkSize = int.MaxValue) + { + byte[] bytes = Encoding.UTF8.GetBytes(json); + var scanner = new JsonRpcFrameScanner(); + scanner.Reset(); + for (int offset = 0; offset < bytes.Length; offset += chunkSize) + { + scanner.Feed(bytes.AsSpan(offset, Math.Min(chunkSize, bytes.Length - offset))); + } + + return scanner; + } +} diff --git a/tests/ModelContextProtocol.Tests/Transport/StdioServerTransportTests.cs b/tests/ModelContextProtocol.Tests/Transport/StdioServerTransportTests.cs index e287cf4d0..ed166da62 100644 --- a/tests/ModelContextProtocol.Tests/Transport/StdioServerTransportTests.cs +++ b/tests/ModelContextProtocol.Tests/Transport/StdioServerTransportTests.cs @@ -389,4 +389,106 @@ public async Task ReadMessagesAsync_Should_Log_Received_At_Trace_Level() Assert.NotEmpty(traceLogMessages); Assert.Contains(traceLogMessages, x => x.Message.Contains("\"method\":\"test\"") && x.Message.Contains("\"id\":99")); } + + [Fact] + public async Task ReadMessagesAsync_Should_Read_Message_Split_Across_Writes() + { + // Every byte arrives in its own write, so multi-byte UTF-8 characters are split across reads. + const string Line = """{"jsonrpc":"2.0","id":"zażółć","method":"test","params":{"text":"gęślą jaźń 🎉"}}""" + "\n"; + + Pipe pipe = new(); + using var input = pipe.Reader.AsStream(); + + await using var transport = new StreamServerTransport(input, Stream.Null, loggerFactory: LoggerFactory); + + foreach (byte b in Encoding.UTF8.GetBytes(Line)) + { + await pipe.Writer.WriteAsync(new[] { b }, TestContext.Current.CancellationToken); + } + + var request = Assert.IsType(await transport.MessageReader.ReadAsync(TestContext.Current.CancellationToken)); + Assert.Equal("zażółć", request.Id.ToString()); + Assert.Equal("gęślą jaźń 🎉", request.Params?["text"]?.GetValue()); + } + + [Fact] + public async Task ReadMessagesAsync_Should_Read_Several_Messages_From_One_Write_And_Skip_Blank_Lines() + { + Pipe pipe = new(); + using var input = pipe.Reader.AsStream(); + + await using var transport = new StreamServerTransport(input, Stream.Null, loggerFactory: LoggerFactory); + + await pipe.Writer.WriteAsync(Encoding.UTF8.GetBytes( + "\n \t\r\n" + + """{"jsonrpc":"2.0","id":1,"method":"first"}""" + "\r\n\n" + + """{"jsonrpc":"2.0","method":"notifications/second"}""" + "\n" + + """{"jsonrpc":"2.0","id":3,"method":"third"}""" + "\n"), + TestContext.Current.CancellationToken); + + Assert.Equal("first", Assert.IsType(await transport.MessageReader.ReadAsync(TestContext.Current.CancellationToken)).Method); + Assert.Equal("notifications/second", Assert.IsType(await transport.MessageReader.ReadAsync(TestContext.Current.CancellationToken)).Method); + Assert.Equal("third", Assert.IsType(await transport.MessageReader.ReadAsync(TestContext.Current.CancellationToken)).Method); + } + + [Fact] + public async Task ReadMessagesAsync_Should_Read_Final_Message_Without_Trailing_Newline() + { + Pipe pipe = new(); + using var input = pipe.Reader.AsStream(); + + await using var transport = new StreamServerTransport(input, Stream.Null, loggerFactory: LoggerFactory); + + await pipe.Writer.WriteAsync(Encoding.UTF8.GetBytes("""{"jsonrpc":"2.0","id":5,"method":"last"}"""), TestContext.Current.CancellationToken); + pipe.Writer.Complete(); + + var request = Assert.IsType(await transport.MessageReader.ReadAsync(TestContext.Current.CancellationToken)); + Assert.Equal("last", request.Method); + + // End of input completes the transport. + Assert.False(await transport.MessageReader.WaitToReadAsync(TestContext.Current.CancellationToken)); + Assert.False(transport.IsConnected); + } + + [Fact] + public async Task ReadMessagesAsync_Should_Skip_Utf8_Byte_Order_Mark_At_Stream_Start() + { + Pipe pipe = new(); + using var input = pipe.Reader.AsStream(); + + await using var transport = new StreamServerTransport(input, Stream.Null, loggerFactory: LoggerFactory); + + byte[] message = Encoding.UTF8.GetBytes("""{"jsonrpc":"2.0","id":1,"method":"bom"}""" + "\n"); + await pipe.Writer.WriteAsync(new byte[] { 0xEF, 0xBB, 0xBF }.Concat(message).ToArray(), TestContext.Current.CancellationToken); + + var request = Assert.IsType(await transport.MessageReader.ReadAsync(TestContext.Current.CancellationToken)); + Assert.Equal("bom", request.Method); + } + + [Fact] + public async Task ReadMessagesAsync_Should_Respond_With_ParseError_For_Trailing_Data() + { + Pipe inputPipe = new(); + Pipe outputPipe = new(); + using var input = inputPipe.Reader.AsStream(); + using var output = outputPipe.Writer.AsStream(); + + await using var transport = new StreamServerTransport(input, output, loggerFactory: LoggerFactory); + + await inputPipe.Writer.WriteAsync( + Encoding.UTF8.GetBytes("""{"jsonrpc":"2.0","method":"ping","id":12} {"extra":true}""" + "\n"), + TestContext.Current.CancellationToken); + + using var responseReader = new StreamReader(outputPipe.Reader.AsStream(), Encoding.UTF8); + var responseLine = await responseReader.ReadLineAsync( +#if NET + TestContext.Current.CancellationToken +#endif + ); + + var error = Assert.IsType(JsonSerializer.Deserialize(responseLine!, McpJsonUtilities.DefaultOptions)); + Assert.Equal(new RequestId(12), error.Id); + Assert.Equal((int)McpErrorCode.ParseError, error.Error.Code); + Assert.True(transport.IsConnected); + } } diff --git a/tests/ModelContextProtocol.Tests/Transport/StreamServerTransportMaxMessageSizeTests.cs b/tests/ModelContextProtocol.Tests/Transport/StreamServerTransportMaxMessageSizeTests.cs new file mode 100644 index 000000000..f5fb30a3b --- /dev/null +++ b/tests/ModelContextProtocol.Tests/Transport/StreamServerTransportMaxMessageSizeTests.cs @@ -0,0 +1,360 @@ +using Microsoft.Extensions.DependencyInjection; +using Microsoft.Extensions.Logging; +using ModelContextProtocol.Client; +using ModelContextProtocol.Protocol; +using ModelContextProtocol.Server; +using ModelContextProtocol.Tests.Utils; +using System.ComponentModel; +using System.IO.Pipelines; +using System.Text; +using System.Text.Json; + +namespace ModelContextProtocol.Tests.Transport; + +public class StreamServerTransportMaxMessageSizeTests(ITestOutputHelper testOutputHelper) : LoggedTest(testOutputHelper) +{ + private const int MaxMessageSize = 1024; + + // A string value long enough to push a message over MaxMessageSize. It contains the characters a + // naive scanner could trip over: quotes, backslashes, braces, and brackets inside a string. + private static readonly string s_largeValue = string.Concat(Enumerable.Repeat("""abc\"}{][\\ """, 1_000)); + + [Theory] + [InlineData(1)] + [InlineData(100)] + [InlineData(int.MaxValue)] + public async Task Oversized_Request_With_Id_After_Params_Is_Answered_With_InvalidRequest(int chunkSize) + { + // This SDK's client writes "id" after "params", so the id of an oversized request can come + // at its very end, long after the transport stopped buffering the message. + string oversized = $$$"""{"jsonrpc":"2.0","method":"tools/call","params":{"name":"echo","arguments":{"text":"{{{s_largeValue}}}"}},"id":"req-1"}"""; + Assert.True(Encoding.UTF8.GetByteCount(oversized) > MaxMessageSize); + + await using var harness = new TransportHarness(MaxMessageSize, LoggerFactory); + + await harness.WriteAsync(oversized + "\n", chunkSize); + await harness.WriteAsync("""{"jsonrpc":"2.0","id":2,"method":"ping"}""" + "\n"); + + var error = Assert.IsType(await harness.ReadOutputAsync()); + Assert.Equal("req-1", error.Id.ToString()); + Assert.Equal((int)McpErrorCode.InvalidRequest, error.Error.Code); + Assert.Contains("maximum message size", error.Error.Message); + + // The oversized request never reaches the session; the next message does. + var next = Assert.IsType(await harness.ReadMessageAsync()); + Assert.Equal("2", next.Id.ToString()); + Assert.Equal(RequestMethods.Ping, next.Method); + Assert.True(harness.Transport.IsConnected); + } + + [Fact] + public async Task Oversized_Request_With_Numeric_Id_Before_Params_Is_Answered_With_InvalidRequest() + { + string oversized = $$$"""{"jsonrpc":"2.0","id":41,"method":"tools/call","params":{"text":"{{{s_largeValue}}}"}}"""; + + await using var harness = new TransportHarness(MaxMessageSize, LoggerFactory); + + await harness.WriteAsync(oversized + "\n", chunkSize: 256); + + var error = Assert.IsType(await harness.ReadOutputAsync()); + Assert.Equal(new RequestId(41), error.Id); + Assert.Equal((int)McpErrorCode.InvalidRequest, error.Error.Code); + } + + [Fact] + public async Task Oversized_Request_At_End_Of_Stream_Without_Newline_Is_Answered() + { + string oversized = $$$"""{"jsonrpc":"2.0","method":"tools/call","params":{"text":"{{{s_largeValue}}}"},"id":7}"""; + + await using var harness = new TransportHarness(MaxMessageSize, LoggerFactory); + + await harness.WriteAsync(oversized, chunkSize: 300); + harness.CompleteInput(); + + var error = Assert.IsType(await harness.ReadOutputAsync()); + Assert.Equal(new RequestId(7), error.Id); + Assert.Equal((int)McpErrorCode.InvalidRequest, error.Error.Code); + } + + [Theory] + [InlineData("""{"jsonrpc":"2.0","result":{"text":"{{VALUE}}"},"id":17}""")] + [InlineData("""{"jsonrpc":"2.0","id":17,"error":{"code":-32603,"message":"{{VALUE}}"}}""")] + public async Task Oversized_Response_Fails_The_Pending_Request_Locally(string template) + { + // A response carries the id of a request this server sent. Answering it with an error would + // send the server's own id back to the client, so the transport fails the request locally. + string oversized = template.Replace("{{VALUE}}", s_largeValue); + + await using var harness = new TransportHarness(MaxMessageSize, LoggerFactory); + + await harness.WriteAsync(oversized + "\n", chunkSize: 512); + await harness.WriteAsync("""{"jsonrpc":"2.0","id":2,"method":"ping"}""" + "\n"); + + var error = Assert.IsType(await harness.ReadMessageAsync()); + Assert.Equal(new RequestId(17), error.Id); + Assert.Equal((int)McpErrorCode.InternalError, error.Error.Code); + Assert.Contains("maximum message size", error.Error.Message); + + Assert.IsType(await harness.ReadMessageAsync()); + + // By now any reply to the peer would already have been written, since replies are sent in input order. + Assert.Equal(0, harness.OutputLength); + } + + [Fact] + public async Task Oversized_Notification_Is_Dropped() + { + string oversized = $$$"""{"jsonrpc":"2.0","method":"notifications/message","params":{"level":"info","data":"{{{s_largeValue}}}"}}"""; + + await using var harness = new TransportHarness(MaxMessageSize, LoggerFactory); + + await harness.WriteAsync(oversized + "\n", chunkSize: 512); + await harness.WriteAsync("""{"jsonrpc":"2.0","id":2,"method":"ping"}""" + "\n"); + + var next = Assert.IsType(await harness.ReadMessageAsync()); + Assert.Equal(new RequestId(2), next.Id); + Assert.Equal(0, harness.OutputLength); + Assert.Contains(MockLoggerProvider.LogMessages, m => m.LogLevel == LogLevel.Warning && m.Message.Contains("exceeds the maximum message size")); + } + + [Fact] + public async Task Message_At_The_Limit_Is_Accepted_And_One_Byte_Over_Is_Rejected() + { + const string AtLimit = """{"jsonrpc":"2.0","id":1,"method":"ping"}"""; + const string OverLimit = """{"jsonrpc":"2.0","id":22,"method":"ping"}"""; + int limit = Encoding.UTF8.GetByteCount(AtLimit); + Assert.Equal(limit + 1, Encoding.UTF8.GetByteCount(OverLimit)); + + await using var harness = new TransportHarness(limit, LoggerFactory); + + await harness.WriteAsync(AtLimit + "\n" + OverLimit + "\n"); + + var accepted = Assert.IsType(await harness.ReadMessageAsync()); + Assert.Equal(new RequestId(1), accepted.Id); + + var error = Assert.IsType(await harness.ReadOutputAsync()); + Assert.Equal(new RequestId(22), error.Id); + Assert.Equal((int)McpErrorCode.InvalidRequest, error.Error.Code); + } + + [Fact] + public async Task Without_A_Limit_A_Large_Message_Is_Accepted() + { + string large = $$$"""{"jsonrpc":"2.0","method":"tools/call","params":{"text":"{{{s_largeValue}}}"},"id":3}"""; + + await using var harness = new TransportHarness(maxMessageSize: null, LoggerFactory); + + await harness.WriteAsync(large + "\n", chunkSize: 100); + + var request = Assert.IsType(await harness.ReadMessageAsync()); + Assert.Equal(new RequestId(3), request.Id); + Assert.Equal(JsonSerializer.Deserialize($"\"{s_largeValue}\""), request.Params?["text"]?.GetValue()); + } + + [Theory] + [InlineData(0)] + [InlineData(-1)] + public void MaxMessageSize_Must_Be_Positive(long value) + { + var options = new StreamServerTransportOptions(); + + Assert.Throws("value", () => options.MaxMessageSize = value); + Assert.Null(options.MaxMessageSize); + + options.MaxMessageSize = 1; + options.MaxMessageSize = null; + Assert.Null(options.MaxMessageSize); + } + + [Fact] + public async Task Oversized_Sampling_Result_Fails_The_Servers_Request_And_Keeps_The_Session() + { + Pipe clientToServer = new(), serverToClient = new(); + using var cts = CancellationTokenSource.CreateLinkedTokenSource(TestContext.Current.CancellationToken); + + var services = new ServiceCollection(); + services.AddLogging(); + services.AddSingleton(XunitLoggerProvider); + services + .AddMcpServer() + .WithStreamServerTransport( + clientToServer.Reader.AsStream(), + serverToClient.Writer.AsStream(), + options => options.MaxMessageSize = 64 * 1024) + .WithTools(); + + await using var serviceProvider = services.BuildServiceProvider(validateScopes: true); + var server = serviceProvider.GetRequiredService(); + Task serverTask = server.RunAsync(cts.Token); + + string sampledText = "x"; + var clientOptions = new McpClientOptions + { + // Exercise the server-to-client sampling request, which the 2026-07-28 protocol replaces with MRTR. + ProtocolVersion = "2025-11-25", + DiscoverProbeTimeout = TestConstants.DefaultTimeout, + Capabilities = new() { Sampling = new() }, + Handlers = new() + { + SamplingHandler = (_, _, _) => new(new CreateMessageResult + { + Content = [new TextContentBlock { Text = sampledText }], + Model = "test-model", + }), + }, + }; + + await using (McpClient client = await McpClient.CreateAsync( + new StreamClientTransport(clientToServer.Writer.AsStream(), serverToClient.Reader.AsStream(), LoggerFactory), + clientOptions, + LoggerFactory, + TestContext.Current.CancellationToken)) + { + sampledText = new string('x', 256 * 1024); + var tooLarge = await client.CallToolAsync("sample", cancellationToken: TestContext.Current.CancellationToken); + Assert.Contains("maximum message size", Assert.IsType(Assert.Single(tooLarge.Content)).Text); + + // The session is still usable afterwards. + sampledText = "small"; + var ok = await client.CallToolAsync("sample", cancellationToken: TestContext.Current.CancellationToken); + Assert.Equal("Sampled 5 characters.", Assert.IsType(Assert.Single(ok.Content)).Text); + } + + await cts.CancelAsync(); + clientToServer.Writer.Complete(); + serverToClient.Writer.Complete(); + await serverTask; + } + + public sealed class SamplingTool + { + [McpServerTool(Name = "sample"), Description("Samples once and reports the result.")] + public static async Task SampleAsync(McpServer server, CancellationToken cancellationToken) + { + try + { + var result = await server.SampleAsync(new CreateMessageRequestParams + { + Messages = [new SamplingMessage { Role = Role.User, Content = [new TextContentBlock { Text = "hi" }] }], + MaxTokens = 10, + }, cancellationToken); + + return $"Sampled {Assert.IsType(Assert.Single(result.Content)).Text.Length} characters."; + } + catch (McpProtocolException ex) + { + return ex.Message; + } + } + } + + /// A over in-memory pipes. + private sealed class TransportHarness : IAsyncDisposable + { + private readonly Pipe _input = new(); + private readonly Pipe _output = new(); + private readonly StreamReader _outputReader; + private long _outputLength; + + public TransportHarness(long? maxMessageSize, ILoggerFactory loggerFactory) + { + Transport = new StreamServerTransport( + _input.Reader.AsStream(), + new CountingStream(_output.Writer.AsStream(), this), + serverName: null, + loggerFactory, + new StreamServerTransportOptions { MaxMessageSize = maxMessageSize }); + _outputReader = new StreamReader(_output.Reader.AsStream(), Encoding.UTF8); + } + + public StreamServerTransport Transport { get; } + + public long OutputLength => Interlocked.Read(ref _outputLength); + + public async Task WriteAsync(string text, int chunkSize = int.MaxValue) + { + byte[] bytes = Encoding.UTF8.GetBytes(text); + for (int offset = 0; offset < bytes.Length; offset += chunkSize) + { + int count = Math.Min(chunkSize, bytes.Length - offset); + await _input.Writer.WriteAsync(bytes.AsMemory(offset, count), TestContext.Current.CancellationToken); + } + } + + public void CompleteInput() => _input.Writer.Complete(); + + public async Task ReadMessageAsync() + { + using var cts = CancellationTokenSource.CreateLinkedTokenSource(TestContext.Current.CancellationToken); + cts.CancelAfter(TestConstants.DefaultTimeout); + return await Transport.MessageReader.ReadAsync(cts.Token); + } + + public async Task ReadOutputAsync() + { + using var cts = CancellationTokenSource.CreateLinkedTokenSource(TestContext.Current.CancellationToken); + cts.CancelAfter(TestConstants.DefaultTimeout); + var line = await _outputReader.ReadLineAsync( +#if NET + cts.Token +#endif + ); + + Assert.NotNull(line); + return JsonSerializer.Deserialize(line, McpJsonUtilities.DefaultOptions); + } + + public async ValueTask DisposeAsync() + { + await Transport.DisposeAsync(); + _input.Writer.Complete(); + _output.Writer.Complete(); + _outputReader.Dispose(); + } + + private sealed class CountingStream(Stream inner, TransportHarness owner) : Stream + { + public override bool CanRead => false; + public override bool CanSeek => false; + public override bool CanWrite => true; + public override long Length => throw new NotSupportedException(); + public override long Position { get => throw new NotSupportedException(); set => throw new NotSupportedException(); } + + public override void Write(byte[] buffer, int offset, int count) + { + Interlocked.Add(ref owner._outputLength, count); + inner.Write(buffer, offset, count); + } + + public override Task WriteAsync(byte[] buffer, int offset, int count, CancellationToken cancellationToken) + { + Interlocked.Add(ref owner._outputLength, count); + return inner.WriteAsync(buffer, offset, count, cancellationToken); + } + +#if NET + public override ValueTask WriteAsync(ReadOnlyMemory buffer, CancellationToken cancellationToken = default) + { + Interlocked.Add(ref owner._outputLength, buffer.Length); + return inner.WriteAsync(buffer, cancellationToken); + } +#endif + + public override void Flush() => inner.Flush(); + public override Task FlushAsync(CancellationToken cancellationToken) => inner.FlushAsync(cancellationToken); + public override int Read(byte[] buffer, int offset, int count) => throw new NotSupportedException(); + public override long Seek(long offset, SeekOrigin origin) => throw new NotSupportedException(); + public override void SetLength(long value) => throw new NotSupportedException(); + + protected override void Dispose(bool disposing) + { + if (disposing) + { + inner.Dispose(); + } + + base.Dispose(disposing); + } + } + } +}