Streamable HTTP resumability + redelivery + SSE polling via server-side disconnect (#1077)

Co-authored-by: Stephen Halter <halter73@gmail.com>

src/ModelContextProtocol.Core/Server/ISseEventStreamReader.cs

@@ -0,0 +1,36 @@
+using ModelContextProtocol.Protocol;
+using System.Net.ServerSentEvents;
+
+namespace ModelContextProtocol.Server;
+
+/// <summary>
+/// Provides read access to an SSE event stream, allowing events to be consumed asynchronously.
+/// </summary>
+public interface ISseEventStreamReader
+{
+    /// <summary>
+    /// Gets the session ID associated with the stream being read.
+    /// </summary>
+    string SessionId { get; }
+
+    /// <summary>
+    /// Gets the ID of the stream.
+    /// </summary>
+    /// <remarks>
+    /// This value is guaranteed to be unique on a per-session basis.
+    /// </remarks>
+    string StreamId { get; }
+
+    /// <summary>
+    /// Gets the messages from the stream as an <see cref="IAsyncEnumerable{T}"/>.
+    /// </summary>
+    /// <param name="cancellationToken">A token to cancel the operation.</param>
+    /// <returns>An <see cref="IAsyncEnumerable{T}"/> of <see cref="SseItem{T}"/> containing JSON-RPC messages.</returns>
+    /// <remarks>
+    /// If the stream's mode is set to <see cref="SseEventStreamMode.Polling"/>, the returned
+    /// messages will only include the currently-available events starting at the last event ID specified
+    /// when the reader was created. Otherwise, the returned messages will continue until the associated
+    /// <see cref="ISseEventStreamWriter"/> is disposed.
+    /// </remarks>
+    IAsyncEnumerable<SseItem<JsonRpcMessage?>> ReadEventsAsync(CancellationToken cancellationToken = default);
+}

src/ModelContextProtocol.Core/Server/ISseEventStreamStore.cs

@@ -0,0 +1,23 @@
+namespace ModelContextProtocol.Server;
+
+/// <summary>
+/// Provides storage and retrieval of SSE event streams, enabling resumability and redelivery of events.
+/// </summary>
+public interface ISseEventStreamStore
+{
+    /// <summary>
+    /// Creates a new SSE event stream with the specified options.
+    /// </summary>
+    /// <param name="options">The configuration options for the new stream.</param>
+    /// <param name="cancellationToken">A token to cancel the operation.</param>
+    /// <returns>A writer for the newly created event stream.</returns>
+    ValueTask<ISseEventStreamWriter> CreateStreamAsync(SseEventStreamOptions options, CancellationToken cancellationToken = default);
+
+    /// <summary>
+    /// Gets a reader for an existing event stream based on the last event ID.
+    /// </summary>
+    /// <param name="lastEventId">The ID of the last event received by the client, used to resume from that point.</param>
+    /// <param name="cancellationToken">A token to cancel the operation.</param>
+    /// <returns>A reader for the event stream, or <c>null</c> if no matching stream is found.</returns>
+    ValueTask<ISseEventStreamReader?> GetStreamReaderAsync(string lastEventId, CancellationToken cancellationToken = default);
+}

src/ModelContextProtocol.Core/Server/ISseEventStreamWriter.cs

@@ -0,0 +1,30 @@
+using ModelContextProtocol.Protocol;
+using System.Net.ServerSentEvents;
+
+namespace ModelContextProtocol.Server;
+
+/// <summary>
+/// Provides write access to an SSE event stream, allowing events to be written and tracked with unique IDs.
+/// </summary>
+public interface ISseEventStreamWriter : IAsyncDisposable
+{
+    /// <summary>
+    /// Sets the mode of the event stream.
+    /// </summary>
+    /// <param name="mode">The new mode to set for the event stream.</param>
+    /// <param name="cancellationToken">A token to cancel the operation.</param>
+    /// <returns>A task that represents the asynchronous operation.</returns>
+    ValueTask SetModeAsync(SseEventStreamMode mode, CancellationToken cancellationToken = default);
+
+    /// <summary>
+    /// Writes an event to the stream.
+    /// </summary>
+    /// <param name="sseItem">The original <see cref="SseItem{T}"/>.</param>
+    /// <param name="cancellationToken">A token to cancel the operation.</param>
+    /// <returns>A new <see cref="SseItem{T}"/> with a populated event ID.</returns>
+    /// <remarks>
+    /// If the provided <paramref name="sseItem"/> already has an event ID, this method skips writing the event.
+    /// Otherwise, an event ID unique to all sessions and streams is generated and assigned to the event.
+    /// </remarks>
+    ValueTask<SseItem<JsonRpcMessage?>> WriteEventAsync(SseItem<JsonRpcMessage?> sseItem, CancellationToken cancellationToken = default);
+}

src/ModelContextProtocol.Core/Server/McpServerImpl.cs

@@ -561,7 +561,7 @@ internal sealed partial class McpServerImpl : McpServer
 
                 try
                 {
-                    return await handler(request, cancellationToken);
+                    return await handler(request, cancellationToken).ConfigureAwait(false);
                 }
                 catch (Exception e) when (e is not OperationCanceledException and not McpProtocolException)
                 {

src/ModelContextProtocol.Core/Server/McpSseEventWriterExtensions.cs

@@ -0,0 +1,70 @@
+using ModelContextProtocol.Protocol;
+using System.Buffers;
+using System.Net.ServerSentEvents;
+using System.Text.Json;
+
+namespace ModelContextProtocol.Server;
+
+/// <summary>
+/// Provides MCP extension methods for <see cref="SseEventWriter"/>.
+/// </summary>
+internal static class McpSseEventWriterExtensions
+{
+    [ThreadStatic]
+    private static Utf8JsonWriter? _jsonWriter;
+
+    /// <summary>
+    /// Writes an SSE item containing a <see cref="JsonRpcMessage"/>.
+    /// </summary>
+    /// <param name="writer">The <see cref="SseEventWriter"/>.</param>
+    /// <param name="item">The SSE item containing the <see cref="JsonRpcMessage"/>.</param>
+    /// <param name="cancellationToken">The token to monitor for cancellation requests.</param>
+    /// <returns>A task representing the asynchronous write operation.</returns>
+    public static ValueTask WriteAsync(this SseEventWriter writer, SseItem<JsonRpcMessage?> item, CancellationToken cancellationToken = default)
+        => writer.WriteAsync(item, FormatJsonRpcMessage, cancellationToken);
+
+    /// <summary>
+    /// Writes an SSE item containing a <see cref="string"/>.
+    /// </summary>
+    /// <param name="writer">The <see cref="SseEventWriter"/>.</param>
+    /// <param name="item">The SSE item containing the string.</param>
+    /// <param name="cancellationToken">The token to monitor for cancellation requests.</param>
+    /// <returns>A task representing the asynchronous write operation.</returns>
+    public static ValueTask WriteAsync(this SseEventWriter writer, SseItem<string> item, CancellationToken cancellationToken = default)
+        => writer.WriteAsync(item, FormatString, cancellationToken);
+
+    /// <summary>
+    /// Formats a <see cref="JsonRpcMessage"/> message by writing it as JSON to the buffer writer.
+    /// </summary>
+    private static void FormatJsonRpcMessage(SseItem<JsonRpcMessage?> item, IBufferWriter<byte> writer)
+    {
+        if (item.Data is null)
+        {
+            return;
+        }
+
+        if (_jsonWriter is null)
+        {
+            _jsonWriter = new Utf8JsonWriter(writer);
+        }
+        else
+        {
+            _jsonWriter.Reset(writer);
+        }
+
+        JsonSerializer.Serialize(_jsonWriter, item.Data, McpJsonUtilities.JsonContext.Default.JsonRpcMessage!);
+    }
+
+    /// <summary>
+    /// Formats a string by writing it as UTF-8 to the buffer writer.
+    /// </summary>
+    private static void FormatString(SseItem<string> item, IBufferWriter<byte> writer)
+    {
+        if (item.Data is null)
+        {
+            return;
+        }
+
+        writer.WriteUtf8String(item.Data);
+    }
+}

src/ModelContextProtocol.Core/Server/RequestContext.cs

@@ -82,4 +82,21 @@ public sealed class RequestContext<TParams>
     /// including the method name, parameters, request ID, and associated transport and user information.
     /// </remarks>
     public JsonRpcRequest JsonRpcRequest { get; }
+
+    /// <summary>
+    /// Ends the current response and enables polling for updates from the server.
+    /// </summary>
+    /// <param name="retryInterval">The interval at which the client should poll for updates.</param>
+    /// <param name="cancellationToken">The cancellation token.</param>
+    /// <returns>A <see cref="ValueTask"/> that completes when polling has been enabled.</returns>
+    /// <exception cref="InvalidOperationException">Thrown when the transport does not support polling.</exception>
+    public async ValueTask EnablePollingAsync(TimeSpan retryInterval, CancellationToken cancellationToken = default)
+    {
+        if (JsonRpcRequest.Context?.RelatedTransport is not StreamableHttpPostTransport transport)
+        {
+            throw new InvalidOperationException("Polling is only supported for Streamable HTTP transports.");
+        }
+
+        await transport.EnablePollingAsync(retryInterval, cancellationToken).ConfigureAwait(false);
+    }
 }

src/ModelContextProtocol.Core/Server/SseEventStreamMode.cs

@@ -0,0 +1,20 @@
+namespace ModelContextProtocol.Server;
+
+/// <summary>
+/// Represents the mode of an SSE event stream.
+/// </summary>
+public enum SseEventStreamMode
+{
+    /// <summary>
+    /// Causes the event stream returned by <see cref="ISseEventStreamReader.ReadEventsAsync(System.Threading.CancellationToken)"/> to only end when
+    /// the associated <see cref="ISseEventStreamWriter"/> gets disposed.
+    /// </summary>
+    Streaming = 0,
+
+    /// <summary>
+    /// Causes the event stream returned by <see cref="ISseEventStreamReader.ReadEventsAsync(System.Threading.CancellationToken)"/> to end
+    /// after the most recent event has been consumed. This forces clients to keep making new requests in order to receive
+    /// the latest messages.
+    /// </summary>
+    Polling = 1,
+}

src/ModelContextProtocol.Core/Server/SseEventStreamOptions.cs

@@ -0,0 +1,22 @@
+namespace ModelContextProtocol.Server;
+
+/// <summary>
+/// Configuration options for creating an SSE event stream.
+/// </summary>
+public sealed class SseEventStreamOptions
+{
+    /// <summary>
+    /// Gets or sets the session ID associated with the event stream.
+    /// </summary>
+    public required string SessionId { get; set; }
+
+    /// <summary>
+    /// Gets or sets the stream ID that uniquely identifies this stream within a session.
+    /// </summary>
+    public required string StreamId { get; set; }
+
+    /// <summary>
+    /// Gets or sets the mode of the event stream. Defaults to <see cref="SseEventStreamMode.Streaming"/>.
+    /// </summary>
+    public SseEventStreamMode Mode { get; set; }
+}

src/ModelContextProtocol.Core/Server/SseResponseStreamTransport.cs

@@ -1,4 +1,5 @@
 using ModelContextProtocol.Protocol;
+using System.Net.ServerSentEvents;
 using System.Security.Claims;
 using System.Threading.Channels;
 
@@ -27,14 +28,17 @@ namespace ModelContextProtocol.Server;
 /// <param name="sessionId">The identifier corresponding to the current MCP session.</param>
 public sealed class SseResponseStreamTransport(Stream sseResponseStream, string? messageEndpoint = "/message", string? sessionId = null) : ITransport
 {
-    private readonly SseWriter _sseWriter = new(messageEndpoint);
     private readonly Channel<JsonRpcMessage> _incomingChannel = Channel.CreateBounded<JsonRpcMessage>(new BoundedChannelOptions(1)
     {
         SingleReader = true,
         SingleWriter = false,
     });
+    private readonly SemaphoreSlim _lock = new(1, 1);
+    private readonly TaskCompletionSource<bool> _completedTcs = new(TaskCreationOptions.RunContinuationsAsynchronously);
+    private readonly SseEventWriter _sseWriter = new(sseResponseStream);
 
     private bool _isConnected;
+    private bool _disposed;
 
     /// <summary>
     /// Starts the transport and writes the JSON-RPC messages sent via <see cref="SendMessageAsync"/>
@@ -45,7 +49,14 @@ public sealed class SseResponseStreamTransport(Stream sseResponseStream, string?
     public async Task RunAsync(CancellationToken cancellationToken = default)
     {
         _isConnected = true;
-        await _sseWriter.WriteAllAsync(sseResponseStream, cancellationToken).ConfigureAwait(false);
+
+        // Write the endpoint event first
+        if (messageEndpoint is not null)
+        {
+            await _sseWriter.WriteAsync(SseItem.Endpoint(messageEndpoint), cancellationToken).ConfigureAwait(false);
+        }
+
+        await _completedTcs.Task.WaitAsync(cancellationToken).ConfigureAwait(false);
     }
 
     /// <inheritdoc/>
@@ -57,17 +68,34 @@ public sealed class SseResponseStreamTransport(Stream sseResponseStream, string?
     /// <inheritdoc/>
     public async ValueTask DisposeAsync()
     {
+        using var _ = await _lock.LockAsync().ConfigureAwait(false);
+
+        if (_disposed)
+        {
+            return;
+        }
+
+        _disposed = true;
         _isConnected = false;
         _incomingChannel.Writer.TryComplete();
-        await _sseWriter.DisposeAsync().ConfigureAwait(false);
+        _completedTcs.TrySetResult(true);
+        _sseWriter.Dispose();
     }
 
     /// <inheritdoc/>
     public async Task SendMessageAsync(JsonRpcMessage message, CancellationToken cancellationToken = default)
     {
         Throw.IfNull(message);
-        // If the underlying writer has been disposed, just drop the message.
-        await _sseWriter.SendMessageAsync(message, cancellationToken).ConfigureAwait(false);
+
+        using var _ = await _lock.LockAsync(cancellationToken).ConfigureAwait(false);
+
+        // If disposed, just drop the message.
+        if (_disposed)
+        {
+            return;
+        }
+
+        await _sseWriter.WriteAsync(SseItem.Message(message), cancellationToken).ConfigureAwait(false);
     }
 
     /// <summary>

src/ModelContextProtocol.Core/Server/SseWriter.cs

@@ -1,120 +0,0 @@
-using ModelContextProtocol.Protocol;
-using System.Buffers;
-using System.Net.ServerSentEvents;
-using System.Text;
-using System.Text.Json;
-using System.Threading.Channels;
-
-namespace ModelContextProtocol.Server;
-
-internal sealed class SseWriter(string? messageEndpoint = null, BoundedChannelOptions? channelOptions = null) : IAsyncDisposable
-{
-    private readonly Channel<SseItem<JsonRpcMessage?>> _messages = Channel.CreateBounded<SseItem<JsonRpcMessage?>>(channelOptions ?? new BoundedChannelOptions(1)
-    {
-        SingleReader = true,
-        SingleWriter = false,
-    });
-
-    private Utf8JsonWriter? _jsonWriter;
-    private Task? _writeTask;
-    private CancellationToken? _writeCancellationToken;
-
-    private readonly SemaphoreSlim _disposeLock = new(1, 1);
-    private bool _disposed;
-
-    public Func<IAsyncEnumerable<SseItem<JsonRpcMessage?>>, CancellationToken, IAsyncEnumerable<SseItem<JsonRpcMessage?>>>? MessageFilter { get; set; }
-
-    public Task WriteAllAsync(Stream sseResponseStream, CancellationToken cancellationToken)
-    {
-        Throw.IfNull(sseResponseStream);
-
-        // When messageEndpoint is set, the very first SSE event isn't really an IJsonRpcMessage, but there's no API to write a single
-        // item of a different type, so we fib and special-case the "endpoint" event type in the formatter.
-        if (messageEndpoint is not null && !_messages.Writer.TryWrite(new SseItem<JsonRpcMessage?>(null, "endpoint")))
-        {
-            throw new InvalidOperationException("You must call RunAsync before calling SendMessageAsync.");
-        }
-
-        _writeCancellationToken = cancellationToken;
-
-        var messages = _messages.Reader.ReadAllAsync(cancellationToken);
-        if (MessageFilter is not null)
-        {
-            messages = MessageFilter(messages, cancellationToken);
-        }
-
-        _writeTask = SseFormatter.WriteAsync(messages, sseResponseStream, WriteJsonRpcMessageToBuffer, cancellationToken);
-        return _writeTask;
-    }
-
-    public async Task<bool> SendMessageAsync(JsonRpcMessage message, CancellationToken cancellationToken = default)
-    {
-        Throw.IfNull(message);
-
-        using var _ = await _disposeLock.LockAsync(cancellationToken).ConfigureAwait(false);
-
-        if (_disposed)
-        {
-            // Don't throw ObjectDisposedException here; just return false to indicate the message wasn't sent.
-            // The calling transport can determine what to do in this case (drop the message, or fall back to another transport).
-            return false;
-        }
-
-        // Emit redundant "event: message" lines for better compatibility with other SDKs.
-        await _messages.Writer.WriteAsync(new SseItem<JsonRpcMessage?>(message, SseParser.EventTypeDefault), cancellationToken).ConfigureAwait(false);
-        return true;
-    }
-
-    public async ValueTask DisposeAsync()
-    {
-        using var _ = await _disposeLock.LockAsync().ConfigureAwait(false);
-
-        if (_disposed)
-        {
-            return;
-        }
-
-        _messages.Writer.Complete();
-        try
-        {
-            if (_writeTask is not null)
-            {
-                await _writeTask.ConfigureAwait(false);
-            }
-        }
-        catch (OperationCanceledException) when (_writeCancellationToken?.IsCancellationRequested == true)
-        {
-            // Ignore exceptions caused by intentional cancellation during shutdown.
-        }
-        finally
-        {
-            _jsonWriter?.Dispose();
-            _disposed = true;
-        }
-    }
-
-    private void WriteJsonRpcMessageToBuffer(SseItem<JsonRpcMessage?> item, IBufferWriter<byte> writer)
-    {
-        if (item.EventType == "endpoint" && messageEndpoint is not null)
-        {
-            writer.Write(Encoding.UTF8.GetBytes(messageEndpoint));
-            return;
-        }
-
-        JsonSerializer.Serialize(GetUtf8JsonWriter(writer), item.Data, McpJsonUtilities.JsonContext.Default.JsonRpcMessage!);
-    }
-
-    private Utf8JsonWriter GetUtf8JsonWriter(IBufferWriter<byte> writer)
-    {
-        if (_jsonWriter is null)
-        {
-            _jsonWriter = new Utf8JsonWriter(writer);
-        }
-        else
-        {
-            _jsonWriter.Reset(writer);
-        }
-
-        return _jsonWriter;
-    }
-}

src/ModelContextProtocol.Core/Server/StreamableHttpPostTransport.cs

@@ -1,9 +1,6 @@
-using ModelContextProtocol.Protocol;
+using ModelContextProtocol.Protocol;
 using System.Diagnostics;
-using System.IO.Pipelines;
 using System.Net.ServerSentEvents;
-using System.Runtime.CompilerServices;
-using System.Security.Claims;
 using System.Text.Json;
 using System.Threading.Channels;
 
@@ -13,17 +10,25 @@ namespace ModelContextProtocol.Server;
 /// Handles processing the request/response body pairs for the Streamable HTTP transport.
 /// This is typically used via <see cref="JsonRpcMessageContext.RelatedTransport"/>.
 /// </summary>
-internal sealed class StreamableHttpPostTransport(StreamableHttpServerTransport parentTransport, Stream responseStream) : ITransport
+internal sealed class StreamableHttpPostTransport(StreamableHttpServerTransport parentTransport, Stream responseStream, CancellationToken sessionCancellationToken) : ITransport
 {
-    private readonly SseWriter _sseWriter = new();
+    private readonly SemaphoreSlim _messageLock = new(1, 1);
+    private readonly TaskCompletionSource<bool> _httpResponseTcs = new(TaskCreationOptions.RunContinuationsAsynchronously);
+    private readonly SseEventWriter _httpSseWriter = new(responseStream);
+
+    private TaskCompletionSource<bool>? _storeStreamTcs;
+    private ISseEventStreamWriter? _storeSseWriter;
+
     private RequestId _pendingRequest;
+    private bool _finalResponseMessageSent;
+    private bool _httpResponseCompleted;
 
     public ChannelReader<JsonRpcMessage> MessageReader => throw new NotSupportedException("JsonRpcMessage.Context.RelatedTransport should only be used for sending messages.");
 
     string? ITransport.SessionId => parentTransport.SessionId;
 
     /// <returns>
-    /// True, if data was written to the respond body.
+    /// True, if data was written to the response body.
     /// False, if nothing was written because the request body did not contain any <see cref="JsonRpcRequest"/> messages to respond to.
     /// The HTTP application should typically respond with an empty "202 Accepted" response in this scenario.
     /// </returns>
@@ -35,11 +40,11 @@ internal sealed class StreamableHttpPostTransport(StreamableHttpServerTransport
         {
             _pendingRequest = request.Id;
 
-            // Invoke the initialize request callback if applicable.
-            if (parentTransport.OnInitRequestReceived is { } onInitRequest && request.Method == RequestMethods.Initialize)
+            // Invoke the initialize request handler if applicable.
+            if (request.Method == RequestMethods.Initialize)
             {
                 var initializeRequest = JsonSerializer.Deserialize(request.Params, McpJsonUtilities.JsonContext.Default.InitializeRequestParams);
-                await onInitRequest(initializeRequest).ConfigureAwait(false);
+                await parentTransport.HandleInitRequestAsync(initializeRequest).ConfigureAwait(false);
             }
         }
 
@@ -51,15 +56,28 @@ internal sealed class StreamableHttpPostTransport(StreamableHttpServerTransport
             message.Context.ExecutionContext = ExecutionContext.Capture();
         }
 
-        await parentTransport.MessageWriter.WriteAsync(message, cancellationToken).ConfigureAwait(false);
-
         if (_pendingRequest.Id is null)
         {
+            await parentTransport.MessageWriter.WriteAsync(message, cancellationToken).ConfigureAwait(false);
             return false;
         }
 
-        _sseWriter.MessageFilter = StopOnFinalResponseFilter;
-        await _sseWriter.WriteAllAsync(responseStream, cancellationToken).ConfigureAwait(false);
+        using (await _messageLock.LockAsync(cancellationToken).ConfigureAwait(false))
+        {
+            var primingItem = await TryStartSseEventStreamAsync(_pendingRequest).ConfigureAwait(false);
+            if (primingItem.HasValue)
+            {
+                await _httpSseWriter.WriteAsync(primingItem.Value, cancellationToken).ConfigureAwait(false);
+            }
+
+            // Ensure that we've sent the priming event before processing the incoming request.
+            await parentTransport.MessageWriter.WriteAsync(message, cancellationToken).ConfigureAwait(false);
+        }
+
+        // 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);
+
         return true;
     }
 
@@ -72,31 +90,136 @@ internal sealed class StreamableHttpPostTransport(StreamableHttpServerTransport
             throw new InvalidOperationException("Server to client requests are not supported in stateless mode.");
         }
 
-        bool isAccepted = await _sseWriter.SendMessageAsync(message, cancellationToken).ConfigureAwait(false);
-        if (!isAccepted)
+        using var _ = await _messageLock.LockAsync().ConfigureAwait(false);
+
+        try
         {
-            // The underlying writer didn't accept the message because the underlying request has completed.
-            // Rather than drop the message, fall back to sending it via the parent transport.
-            await parentTransport.SendMessageAsync(message, cancellationToken).ConfigureAwait(false);
+
+            if (_finalResponseMessageSent)
+            {
+                // The final response message has already been sent.
+                // Rather than drop the message, fall back to sending it via the parent transport.
+                await parentTransport.SendMessageAsync(message, cancellationToken).ConfigureAwait(false);
+                return;
+            }
+
+            var item = new SseItem<JsonRpcMessage?>(message, SseParser.EventTypeDefault);
+
+            if (_storeSseWriter is not null)
+            {
+                item = await _storeSseWriter.WriteEventAsync(item, cancellationToken).ConfigureAwait(false);
+            }
+
+            if (!_httpResponseCompleted)
+            {
+                // Only write the message to the response if the response has not completed.
+
+                try
+                {
+                    await _httpSseWriter.WriteAsync(item, cancellationToken).ConfigureAwait(false);
+                }
+                catch (Exception ex) when (!cancellationToken.IsCancellationRequested)
+                {
+                    _httpResponseTcs.TrySetException(ex);
+                }
+            }
+        }
+        finally
+        {
+            // Complete the response if this is the final message.
+            if ((message is JsonRpcResponse or JsonRpcError) && ((JsonRpcMessageWithId)message).Id == _pendingRequest)
+            {
+                _finalResponseMessageSent = true;
+                _httpResponseTcs.TrySetResult(true);
+                _storeStreamTcs?.TrySetResult(true);
+            }
+        }
+    }
+
+    public async ValueTask EnablePollingAsync(TimeSpan retryInterval, CancellationToken cancellationToken)
+    {
+        if (parentTransport.Stateless)
+        {
+            throw new InvalidOperationException("Polling is not supported in stateless mode.");
+        }
+
+        using var _ = await _messageLock.LockAsync(cancellationToken).ConfigureAwait(false);
+
+        if (_storeSseWriter is null)
+        {
+            throw new InvalidOperationException($"Polling requires an event stream store to be configured.");
+        }
+
+        // Send the priming event with the new retry interval.
+        var primingItem = await _storeSseWriter.WriteEventAsync(
+            sseItem: new SseItem<JsonRpcMessage?>() { ReconnectionInterval = retryInterval },
+            cancellationToken)
+            .ConfigureAwait(false);
+
+        // Write to the response stream if it still exists.
+        if (!_httpResponseCompleted)
+        {
+            await _httpSseWriter.WriteAsync(primingItem, cancellationToken).ConfigureAwait(false);
+        }
+
+        // Set the mode to 'Polling' so that the replay stream ends as soon as all available messages have been sent.
+        // This prevents the client from immediately establishing another long-lived connection.
+        await _storeSseWriter.SetModeAsync(SseEventStreamMode.Polling, cancellationToken).ConfigureAwait(false);
+
+        // Signal completion so HandlePostAsync can return.
+        _httpResponseTcs.TrySetResult(true);
+    }
+
+    private async ValueTask<SseItem<JsonRpcMessage?>?> TryStartSseEventStreamAsync(RequestId requestId)
+    {
+        Debug.Assert(_storeSseWriter is null);
+
+        _storeSseWriter = await parentTransport.TryCreateEventStreamAsync(
+            streamId: requestId.Id!.ToString()!,
+            cancellationToken: sessionCancellationToken)
+            .ConfigureAwait(false);
+
+        if (_storeSseWriter is null)
+        {
+            return null;
+        }
+
+        _storeStreamTcs = new TaskCompletionSource<bool>(TaskCreationOptions.RunContinuationsAsynchronously);
+        _ = HandleStoreStreamDisposalAsync(_storeStreamTcs.Task);
+
+        return await _storeSseWriter.WriteEventAsync(SseItem.Prime<JsonRpcMessage>(), sessionCancellationToken).ConfigureAwait(false);
+
+        async Task HandleStoreStreamDisposalAsync(Task streamTask)
+        {
+            try
+            {
+                await streamTask.WaitAsync(sessionCancellationToken).ConfigureAwait(false);
+            }
+            finally
+            {
+                using var _ = await _messageLock.LockAsync().ConfigureAwait(false);
+
+                await _storeSseWriter!.DisposeAsync().ConfigureAwait(false);
+            }
         }
     }
 
     public async ValueTask DisposeAsync()
     {
-        await _sseWriter.DisposeAsync().ConfigureAwait(false);
-    }
+        using var _ = await _messageLock.LockAsync().ConfigureAwait(false);
 
-    private async IAsyncEnumerable<SseItem<JsonRpcMessage?>> StopOnFinalResponseFilter(IAsyncEnumerable<SseItem<JsonRpcMessage?>> messages, [EnumeratorCancellation] CancellationToken cancellationToken)
-    {
-        await foreach (var message in messages.WithCancellation(cancellationToken))
+        if (_httpResponseCompleted)
         {
-            yield return message;
-
-            if (message.Data is JsonRpcResponse or JsonRpcError && ((JsonRpcMessageWithId)message.Data).Id == _pendingRequest)
-            {
-                // Complete the SSE response stream now that all pending requests have been processed.
-                break;
-            }
+            return;
         }
+
+        _httpResponseCompleted = true;
+
+        _httpResponseTcs.TrySetResult(true);
+
+        _httpSseWriter.Dispose();
+
+        // Don't dispose the event stream writer here, as we may continue to write to the event store
+        // after disposal if there are pending messages.
     }
 }

src/ModelContextProtocol.Core/Server/StreamableHttpServerTransport.cs

@@ -1,5 +1,7 @@
 using ModelContextProtocol.Protocol;
-using System.IO.Pipelines;
+using System.Diagnostics;
+using System.Diagnostics.CodeAnalysis;
+using System.Net.ServerSentEvents;
 using System.Security.Claims;
 using System.Threading.Channels;
 
@@ -21,24 +23,27 @@ namespace ModelContextProtocol.Server;
 /// </remarks>
 public sealed class StreamableHttpServerTransport : ITransport
 {
-    // For JsonRpcMessages without a RelatedTransport, we don't want to block just because the client didn't make a GET request to handle unsolicited messages.
-    private readonly SseWriter _sseWriter = new(channelOptions: new BoundedChannelOptions(1)
-    {
-        SingleReader = true,
-        SingleWriter = false,
-        FullMode = BoundedChannelFullMode.DropOldest,
-    });
+    /// <summary>
+    /// The stream ID used for unsolicited messages sent via the standalone GET SSE stream.
+    /// </summary>
+    public static readonly string UnsolicitedMessageStreamId = "__get__";
+
     private readonly Channel<JsonRpcMessage> _incomingChannel = Channel.CreateBounded<JsonRpcMessage>(new BoundedChannelOptions(1)
     {
         SingleReader = true,
         SingleWriter = false,
     });
-    private readonly CancellationTokenSource _disposeCts = new();
+    private readonly CancellationTokenSource _transportDisposedCts = new();
+    private readonly SemaphoreSlim _unsolicitedMessageLock = new(1, 1);
 
-    private int _getRequestStarted;
+    private SseEventWriter? _httpSseWriter;
+    private ISseEventStreamWriter? _storeSseWriter;
+    private TaskCompletionSource<bool>? _httpResponseTcs;
+    private bool _getHttpRequestStarted;
+    private bool _getHttpResponseCompleted;
 
     /// <inheritdoc/>
-    public string? SessionId { get; set; }
+    public string? SessionId { get; init; }
 
     /// <summary>
     /// Gets or initializes a value that indicates whether the transport should be in stateless mode that does not require all requests for a given session
@@ -59,15 +64,30 @@ public sealed class StreamableHttpServerTransport : ITransport
     public bool FlowExecutionContextFromRequests { get; init; }
 
     /// <summary>
-    /// Gets or sets a callback to be invoked before handling the initialize request.
+    /// Gets or sets the event store for resumability support.
+    /// When set, events are stored and can be replayed when clients reconnect with a Last-Event-ID header.
     /// </summary>
-    public Func<InitializeRequestParams?, ValueTask>? OnInitRequestReceived { get; set; }
+    public ISseEventStreamStore? EventStreamStore { get; init; }
+
+    /// <summary>
+    /// Gets or sets the negotiated protocol version for this session.
+    /// </summary>
+    internal string? NegotiatedProtocolVersion { get; private set; }
 
     /// <inheritdoc/>
     public ChannelReader<JsonRpcMessage> MessageReader => _incomingChannel.Reader;
 
     internal ChannelWriter<JsonRpcMessage> MessageWriter => _incomingChannel.Writer;
 
+    /// <summary>
+    /// Handles the initialize request by capturing the protocol version and invoking the user callback.
+    /// </summary>
+    internal async ValueTask HandleInitRequestAsync(InitializeRequestParams? initParams)
+    {
+        // Capture the negotiated protocol version for resumability checks
+        NegotiatedProtocolVersion = initParams?.ProtocolVersion;
+    }
+
     /// <summary>
     /// Handles an optional SSE GET request a client using the Streamable HTTP transport might make by
     /// writing any unsolicited JSON-RPC messages sent via <see cref="SendMessageAsync"/>
@@ -78,8 +98,7 @@ public sealed class StreamableHttpServerTransport : ITransport
     /// <returns>A task representing the send loop that writes JSON-RPC messages to the SSE response stream.</returns>
     /// <exception cref="ArgumentNullException"><paramref name="sseResponseStream"/> is <see langword="null"/>.</exception>
     /// <exception cref="InvalidOperationException">
-    /// <see cref="Stateless"/> is <see langword="true"/> and GET requests are not supported in stateless mode,
-    /// or a GET request has already been started for this session.
+    /// <see cref="Stateless"/> is <see langword="true"/> and GET requests are not supported in stateless mode.
     /// </exception>
     public async Task HandleGetRequestAsync(Stream sseResponseStream, CancellationToken cancellationToken = default)
     {
@@ -90,13 +109,27 @@ public sealed class StreamableHttpServerTransport : ITransport
             throw new InvalidOperationException("GET requests are not supported in stateless mode.");
         }
 
-        if (Interlocked.Exchange(ref _getRequestStarted, 1) == 1)
+        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);
+            }
         }
 
-        // We do not need to reference _disposeCts like in HandlePostRequest, because the session ending completes the _sseWriter gracefully.
-        await _sseWriter.WriteAllAsync(sseResponseStream, cancellationToken).ConfigureAwait(false);
+        // 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>
@@ -122,9 +155,15 @@ public sealed class StreamableHttpServerTransport : ITransport
         Throw.IfNull(message);
         Throw.IfNull(responseStream);
 
-        using var postCts = CancellationTokenSource.CreateLinkedTokenSource(_disposeCts.Token, cancellationToken);
-        await using var postTransport = new StreamableHttpPostTransport(this, responseStream);
-        return await postTransport.HandlePostAsync(message, postCts.Token).ConfigureAwait(false);
+        var postTransport = new StreamableHttpPostTransport(this, responseStream, _transportDisposedCts.Token);
+        using var postCts = CancellationTokenSource.CreateLinkedTokenSource(_transportDisposedCts.Token, cancellationToken);
+        await using (postTransport.ConfigureAwait(false))
+        {
+            return await postTransport.HandlePostAsync(
+                message,
+                cancellationToken: postCts.Token)
+                .ConfigureAwait(false);
+        }
     }
 
     /// <inheritdoc/>
@@ -137,28 +176,94 @@ public sealed class StreamableHttpServerTransport : ITransport
             throw new InvalidOperationException("Unsolicited server to client messages are not supported in stateless mode.");
         }
 
-        // If the underlying writer has been disposed, just drop the message.
-        await _sseWriter.SendMessageAsync(message, cancellationToken).ConfigureAwait(false);
+        using var _ = await _unsolicitedMessageLock.LockAsync(cancellationToken).ConfigureAwait(false);
+
+        if (!_getHttpRequestStarted)
+        {
+            // Clients are not required to make a GET request for unsolicited messages.
+            // If no GET request has been made, drop the message.
+            return;
+        }
+
+        Debug.Assert(_httpSseWriter is not null);
+        Debug.Assert(_httpResponseTcs is not null);
+
+        var item = SseItem.Message(message);
+
+        if (_storeSseWriter is not null)
+        {
+            item = await _storeSseWriter.WriteEventAsync(item, cancellationToken).ConfigureAwait(false);
+        }
+
+        if (!_getHttpResponseCompleted)
+        {
+            // Only write the message to the response if the response has not completed.
+
+            try
+            {
+                await _httpSseWriter!.WriteAsync(item, cancellationToken).ConfigureAwait(false);
+            }
+            catch (Exception ex) when (!cancellationToken.IsCancellationRequested)
+            {
+                _httpResponseTcs!.TrySetException(ex);
+            }
+        }
     }
 
     /// <inheritdoc/>
     public async ValueTask DisposeAsync()
     {
+        using var _ = await _unsolicitedMessageLock.LockAsync().ConfigureAwait(false);
+
+        if (_getHttpResponseCompleted)
+        {
+            return;
+        }
+
+        _getHttpResponseCompleted = true;
+
         try
         {
             _incomingChannel.Writer.TryComplete();
-            await _disposeCts.CancelAsync();
+            await _transportDisposedCts.CancelAsync().ConfigureAwait(false);
         }
         finally
         {
             try
             {
-                await _sseWriter.DisposeAsync().ConfigureAwait(false);
+                _httpResponseTcs?.TrySetResult(true);
+                _httpSseWriter?.Dispose();
+
+                if (_storeSseWriter is not null)
+                {
+                    await _storeSseWriter.DisposeAsync().ConfigureAwait(false);
+                }
             }
             finally
             {
-                _disposeCts.Dispose();
+                _transportDisposedCts.Dispose();
             }
         }
     }
+
+    internal async ValueTask<ISseEventStreamWriter?> TryCreateEventStreamAsync(string streamId, CancellationToken cancellationToken)
+    {
+        if (EventStreamStore is null || !McpSessionHandler.SupportsPrimingEvent(NegotiatedProtocolVersion))
+        {
+            return null;
+        }
+
+        // We use the 'Streaming' stream mode so that in the case of an unexpected network disconnection,
+        // the client can continue reading the remaining messages in a single, streamed response.
+        const SseEventStreamMode Mode = SseEventStreamMode.Streaming;
+
+        var sseEventStreamWriter = await EventStreamStore.CreateStreamAsync(new SseEventStreamOptions
+        {
+            SessionId = SessionId ?? Guid.NewGuid().ToString("N"),
+            StreamId = streamId,
+            Mode = Mode,
+        }, cancellationToken).ConfigureAwait(false);
+
+        return sseEventStreamWriter;
+    }
 }