From 25f080115da01e36f8af67dc56bbb2675e4e85ba Mon Sep 17 00:00:00 2001 From: Copilot <223556219+Copilot@users.noreply.github.com> Date: Thu, 28 May 2026 12:03:58 -0700 Subject: [PATCH] Apply review feedback: replace _getHttpResponseCompleted with _disposed; rewrite test - Remove _getHttpResponseCompleted and rely on _httpSseWriter being null as the signal that the GET response stream has been released. Null out _httpSseWriter in DisposeAsync so SendMessageAsync correctly skips writing post-dispose. - Replace the dispose-idempotency role of the old flag with a dedicated _disposed bool. - Move the explanatory comment in SendMessageAsync onto the event store branch (which intentionally still runs when the response stream is gone, to support Last-Event-ID replay). - Rewrite the unit test to assert observable behavior via a recording Stream instead of reflecting on private fields. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../Server/StreamableHttpServerTransport.cs | 23 +++-- .../StreamableHttpServerTransportTests.cs | 90 ++++++++++++++----- 2 files changed, 77 insertions(+), 36 deletions(-) diff --git a/src/ModelContextProtocol.Core/Server/StreamableHttpServerTransport.cs b/src/ModelContextProtocol.Core/Server/StreamableHttpServerTransport.cs index 64a160c..cd39b26 100644 --- a/src/ModelContextProtocol.Core/Server/StreamableHttpServerTransport.cs +++ b/src/ModelContextProtocol.Core/Server/StreamableHttpServerTransport.cs @@ -44,7 +44,7 @@ public sealed partial class StreamableHttpServerTransport : ITransport private TaskCompletionSource? _httpResponseTcs; private string? _negotiatedProtocolVersion; private bool _getHttpRequestStarted; - private bool _getHttpResponseCompleted; + private bool _disposed; /// /// Initializes a new instance of the class. @@ -177,7 +177,6 @@ public sealed partial class StreamableHttpServerTransport : ITransport // otherwise accumulate significant unmanaged memory per session during that interval. using (await _unsolicitedMessageLock.LockAsync(CancellationToken.None).ConfigureAwait(false)) { - _getHttpResponseCompleted = true; if (_httpSseWriter is { } writer) { _httpSseWriter = null; @@ -240,23 +239,19 @@ public sealed partial class StreamableHttpServerTransport : ITransport return; } - // _httpSseWriter may be null here if the GET request has already ended (e.g. client - // disconnected). _getHttpResponseCompleted is set to true in that case, so the write - // below is correctly skipped. The event store writer (if configured) still captures - // the message so a reconnecting client can replay it via Last-Event-ID. Debug.Assert(_httpResponseTcs is not null); var item = SseItem.Message(message); if (_storeSseWriter is not null) { + // Always record the message in the event store (if configured) — even when the GET + // response stream is gone — so a reconnecting client can replay it via Last-Event-ID. item = await _storeSseWriter.WriteEventAsync(item, cancellationToken).ConfigureAwait(false); } - if (!_getHttpResponseCompleted && _httpSseWriter is { } writer) + if (_httpSseWriter is { } writer) { - // Only write the message to the response if the response has not completed. - try { await writer.WriteAsync(item, cancellationToken).ConfigureAwait(false); @@ -273,12 +268,12 @@ public sealed partial class StreamableHttpServerTransport : ITransport { using var _ = await _unsolicitedMessageLock.LockAsync().ConfigureAwait(false); - if (_getHttpResponseCompleted) + if (_disposed) { return; } - _getHttpResponseCompleted = true; + _disposed = true; try { @@ -290,7 +285,11 @@ public sealed partial class StreamableHttpServerTransport : ITransport try { _httpResponseTcs?.TrySetResult(true); - _httpSseWriter?.Dispose(); + if (_httpSseWriter is { } writer) + { + _httpSseWriter = null; + writer.Dispose(); + } if (_storeSseWriter is not null) { diff --git a/tests/ModelContextProtocol.Tests/Transport/StreamableHttpServerTransportTests.cs b/tests/ModelContextProtocol.Tests/Transport/StreamableHttpServerTransportTests.cs index 77f322b..ce2147e 100644 --- a/tests/ModelContextProtocol.Tests/Transport/StreamableHttpServerTransportTests.cs +++ b/tests/ModelContextProtocol.Tests/Transport/StreamableHttpServerTransportTests.cs @@ -1,47 +1,89 @@ +using ModelContextProtocol.Protocol; using ModelContextProtocol.Server; using ModelContextProtocol.Tests.Utils; -using System.IO.Pipelines; -using System.Reflection; namespace ModelContextProtocol.Tests.Transport; public class StreamableHttpServerTransportTests(ITestOutputHelper testOutputHelper) : LoggedTest(testOutputHelper) { [Fact] - public async Task HandleGetRequestAsync_ReleasesStreamReference_AfterRequestEnds() + public async Task SendMessageAsync_AfterGetRequestEnds_DoesNotWriteToResponseStream() { + // Regression test for the SSE response stream being retained after the GET request + // handler returns. Without releasing the stream reference, the Kestrel connection + // and its associated memory pool buffers (~20MiB per SSE session) stay pinned in + // unmanaged memory until the session is eventually disposed (via explicit DELETE or + // idle timeout), causing steady memory growth for servers whose clients disconnect + // without sending DELETE. After the GET handler returns, SendMessageAsync must not + // attempt to write to the (now released) response stream. + await using var transport = new StreamableHttpServerTransport() { SessionId = "test-session", }; - var pipe = new Pipe(); - var responseStream = pipe.Writer.AsStream(); + var responseStream = new RecordingStream(); using var cts = new CancellationTokenSource(); - cts.CancelAfter(TimeSpan.FromMilliseconds(100)); + var getTask = transport.HandleGetRequestAsync(responseStream, cts.Token); - try - { - await transport.HandleGetRequestAsync(responseStream, cts.Token); - } - catch (OperationCanceledException) + // Wait until the GET handler has finished initialization (signaled by the initial + // flush that sends HTTP response headers) so we know _httpSseWriter is set. + await responseStream.FirstActivity.WaitAsync(TestConstants.DefaultTimeout, TestContext.Current.CancellationToken); + + var writeCountBeforeCancel = responseStream.WriteCount; + + cts.Cancel(); + await Assert.ThrowsAnyAsync(() => getTask); + + await transport.SendMessageAsync( + new JsonRpcNotification { Method = "test" }, + TestContext.Current.CancellationToken); + + Assert.Equal(writeCountBeforeCancel, responseStream.WriteCount); + } + + private sealed class RecordingStream : Stream + { + private readonly TaskCompletionSource _firstActivity = new(TaskCreationOptions.RunContinuationsAsynchronously); + private int _writeCount; + + public Task FirstActivity => _firstActivity.Task; + public int WriteCount => Volatile.Read(ref _writeCount); + + 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(); } - // After the GET request handler returns, the transport must not retain a reference to the - // response stream via _httpSseWriter. Otherwise the Kestrel connection and its associated - // memory pool buffers (which can be ~20MiB per SSE session) stay pinned in unmanaged memory - // until the session is eventually disposed (via explicit DELETE or idle timeout), causing - // steady memory growth for servers whose clients disconnect without sending DELETE. - var httpSseWriterField = typeof(StreamableHttpServerTransport).GetField( - "_httpSseWriter", - BindingFlags.Instance | BindingFlags.NonPublic); - Assert.NotNull(httpSseWriterField); - var httpSseWriterValue = httpSseWriterField.GetValue(transport); - Assert.Null(httpSseWriterValue); + public override void Flush() => _firstActivity.TrySetResult(true); - await pipe.Reader.CompleteAsync(); - await pipe.Writer.CompleteAsync(); + public override Task FlushAsync(CancellationToken cancellationToken) + { + _firstActivity.TrySetResult(true); + return Task.CompletedTask; + } + + 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(); + + public override void Write(byte[] buffer, int offset, int count) + { + Interlocked.Increment(ref _writeCount); + _firstActivity.TrySetResult(true); + } + + public override Task WriteAsync(byte[] buffer, int offset, int count, CancellationToken cancellationToken) + { + Interlocked.Increment(ref _writeCount); + _firstActivity.TrySetResult(true); + return Task.CompletedTask; + } } }