modelcontextprotocol/csharp-sdk
Watch
1
0
Fork
You've already forked csharp-sdk
 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 1

Apply review feedback: replace _getHttpResponseCompleted with _disposed; rewrite test

Browse source 
- 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>
This commit is contained in:
Copilot 2026-05-28 12:03:58 -07:00 • committed by Stephen Halter
commit 25f080115d
2 changed files with 76 additions and 35 deletions
Download patch file Download diff file Expand all files Collapse all files

23
src/ModelContextProtocol.Core/Server/StreamableHttpServerTransport.cs
View file

@ -44,7 +44,7 @@ public sealed partial class StreamableHttpServerTransport : ITransport
private TaskCompletionSource<bool>? _httpResponseTcs;
private string? _negotiatedProtocolVersion;
private bool _getHttpRequestStarted;
private bool _getHttpResponseCompleted;
private bool _disposed;
/// <summary>
/// Initializes a new instance of the <see cref="StreamableHttpServerTransport"/> 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)
{

90
tests/ModelContextProtocol.Tests/Transport/StreamableHttpServerTransportTests.cs
View file

@ -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<OperationCanceledException>(() => getTask);
await transport.SendMessageAsync(
new JsonRpcNotification { Method = "test" },
TestContext.Current.CancellationToken);
Assert.Equal(writeCountBeforeCancel, responseStream.WriteCount);
}
private sealed class RecordingStream : Stream
{
private readonly TaskCompletionSource<bool> _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;
}
}
}