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

[release/1.x] Release SSE response stream reference when GET request ends (#1628)

Browse source
This commit is contained in:
Jeff Handley 2026-07-07 15:36:02 -04:00 • committed by GitHub
commit 0021d8875d
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 141 additions and 29 deletions
Download patch file Download diff file Expand all files Collapse all files

81
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.
@ -137,33 +137,53 @@ public sealed partial class StreamableHttpServerTransport : ITransport
throw new InvalidOperationException("GET requests are not supported in stateless mode.");
}
using (await _unsolicitedMessageLock.LockAsync(cancellationToken).ConfigureAwait(false))
try
{
if (_getHttpRequestStarted)
using (await _unsolicitedMessageLock.LockAsync(cancellationToken).ConfigureAwait(false))
{
throw new InvalidOperationException("Session resumption is not yet supported. Please start a new session.");
if (_getHttpRequestStarted)
{
throw new InvalidOperationException("Session resumption is not yet supported. Please start a new session.");
}
_getHttpRequestStarted = true;
_httpSseWriter = new SseEventWriter(sseResponseStream);
_httpResponseTcs = new TaskCompletionSource<bool>(TaskCreationOptions.RunContinuationsAsynchronously);
_storeSseWriter = await TryCreateEventStreamAsync(streamId: UnsolicitedMessageStreamId, cancellationToken).ConfigureAwait(false);
if (_storeSseWriter is not null)
{
var primingItem = await _storeSseWriter.WriteEventAsync(SseItem.Prime<JsonRpcMessage>(), cancellationToken).ConfigureAwait(false);
await _httpSseWriter.WriteAsync(primingItem, cancellationToken).ConfigureAwait(false);
}
else
{
// If there's no priming write, flush the stream to ensure HTTP response headers are
// sent to the client now that the transport is ready to accept messages via SendMessageAsync.
await sseResponseStream.FlushAsync(cancellationToken).ConfigureAwait(false);
}
}
_getHttpRequestStarted = true;
_httpSseWriter = new SseEventWriter(sseResponseStream);
_httpResponseTcs = new TaskCompletionSource<bool>(TaskCreationOptions.RunContinuationsAsynchronously);
_storeSseWriter = await TryCreateEventStreamAsync(streamId: UnsolicitedMessageStreamId, cancellationToken).ConfigureAwait(false);
if (_storeSseWriter is not null)
// Wait for the response to be written before returning from the handler.
// This keeps the HTTP response open until the final response message is sent.
await _httpResponseTcs.Task.WaitAsync(cancellationToken).ConfigureAwait(false);
}
finally
{
// Release the SseEventWriter's reference to the response stream promptly when the GET
// request ends, regardless of how it exits. Otherwise the response stream (and the
// underlying Kestrel connection and associated memory pool buffers) remains pinned
// in memory until the session itself is disposed (via explicit DELETE or idle timeout).
// Clients that disconnect without sending DELETE — common with long-lived SSE — would
// otherwise accumulate significant unmanaged memory per session during that interval.
using (await _unsolicitedMessageLock.LockAsync(CancellationToken.None).ConfigureAwait(false))
{
var primingItem = await _storeSseWriter.WriteEventAsync(SseItem.Prime<JsonRpcMessage>(), cancellationToken).ConfigureAwait(false);
await _httpSseWriter.WriteAsync(primingItem, cancellationToken).ConfigureAwait(false);
}
else
{
// If there's no priming write, flush the stream to ensure HTTP response headers are
// sent to the client now that the transport is ready to accept messages via SendMessageAsync.
await sseResponseStream.FlushAsync(cancellationToken).ConfigureAwait(false);
if (_httpSseWriter is { } writer)
{
_httpSseWriter = null;
writer.Dispose();
}
}
}
// Wait for the response to be written before returning from the handler.
// This keeps the HTTP response open until the final response message is sent.
await _httpResponseTcs.Task.WaitAsync(cancellationToken).ConfigureAwait(false);
}
/// <summary>
@ -219,23 +239,22 @@ public sealed partial class StreamableHttpServerTransport : ITransport
return;
}
Debug.Assert(_httpSseWriter is not null);
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)
if (_httpSseWriter is { } writer)
{
// Only write the message to the response if the response has not completed.
try
{
await _httpSseWriter!.WriteAsync(item, cancellationToken).ConfigureAwait(false);
await writer.WriteAsync(item, cancellationToken).ConfigureAwait(false);
}
catch (Exception ex) when (!cancellationToken.IsCancellationRequested)
{
@ -249,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
{
@ -266,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)
{

89
tests/ModelContextProtocol.Tests/Transport/StreamableHttpServerTransportTests.cs Normal file
View file

@ -0,0 +1,89 @@
using ModelContextProtocol.Protocol;
using ModelContextProtocol.Server;
using ModelContextProtocol.Tests.Utils;
namespace ModelContextProtocol.Tests.Transport;
public class StreamableHttpServerTransportTests(ITestOutputHelper testOutputHelper) : LoggedTest(testOutputHelper)
{
[Fact]
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 responseStream = new RecordingStream();
using var cts = new CancellationTokenSource();
var getTask = transport.HandleGetRequestAsync(responseStream, cts.Token);
// 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();
}
public override void Flush() => _firstActivity.TrySetResult(true);
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;
}
}
}