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

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

Browse source 
Co-authored-by: Stephen Halter <halter73@gmail.com>
This commit is contained in:
Mackinnon Buck 2026-01-14 22:54:25 -08:00 • committed by GitHub
commit b8311756b4
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 2675 additions and 246 deletions
Download patch file Download diff file Expand all files Collapse all files

185
src/ModelContextProtocol.Core/Server/StreamableHttpPostTransport.cs
View file

@ -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.
}
}