From 11f2d18a3b8c3091215bda4f243ec1e95ff2bf41 Mon Sep 17 00:00:00 2001 From: Artur Sawicki Date: Mon, 9 Nov 2026 01:07:59 +0100 Subject: [PATCH] Read stdio messages as UTF-8 and add StreamServerTransportOptions.MaxMessageSize StreamServerTransport now reads through a PipeReader and deserializes each JsonRpcMessage directly from the UTF-8 bytes of the line, without an intermediate string. For a 20 MiB base64 tools/call argument this removes the 128 MiB ArrayPool rent made by the StreamReader line buffer and cuts allocation per message from 368 MiB to 86 MiB on the first call. The new optional MaxMessageSize (default: no limit) bounds the memory used to read one message. An oversized line is not buffered: the transport scans it in constant memory for its top-level id, which this SDK's client writes after params. An oversized request is answered with -32600 for its id. An oversized response to a request the server sent (for example a sampling result) fails that pending request locally instead of sending the server's own id back to the peer. Oversized notifications are dropped and logged. The option is exposed through new StreamServerTransport and StdioServerTransport constructor overloads and new WithStdioServerTransport and WithStreamServerTransport builder overloads. Fixes #2004 --- docs/concepts/transports/transports.md | 10 + .../Protocol/JsonRpcFrameScanner.cs | 481 ++++++++++++++++++ .../Server/StdioServerTransport.cs | 29 +- .../Server/StreamServerTransport.cs | 397 ++++++++++++--- .../Server/StreamServerTransportOptions.cs | 46 ++ .../McpServerBuilderExtensions.cs | 59 ++- ...pServerBuilderExtensionsTransportsTests.cs | 47 ++ .../ModelContextProtocol.Tests.csproj | 1 + .../Protocol/JsonRpcFrameScannerTests.cs | 141 +++++ .../Transport/StdioServerTransportTests.cs | 102 ++++ ...treamServerTransportMaxMessageSizeTests.cs | 360 +++++++++++++ 11 files changed, 1608 insertions(+), 65 deletions(-) create mode 100644 src/ModelContextProtocol.Core/Protocol/JsonRpcFrameScanner.cs create mode 100644 src/ModelContextProtocol.Core/Server/StreamServerTransportOptions.cs create mode 100644 tests/ModelContextProtocol.Tests/Protocol/JsonRpcFrameScannerTests.cs create mode 100644 tests/ModelContextProtocol.Tests/Transport/StreamServerTransportMaxMessageSizeTests.cs 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); + } + } + } +}