Release SSE response stream reference when GET request ends (#1519)

This commit is contained in:
Joel Forsyth 2026-06-05 20:14:49 -04:00 • committed by GitHub
commit acb4cc9fd1
No known key found for this signature in database
GPG key ID: B5690EEEBB952194

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)
{