Fix slow shutdown when a Streamable HTTP client is connected (#843)

* Add "Async" suffix to async methods
This commit is contained in:
Stephen Halter 2025-10-10 17:50:57 -07:00 • committed by GitHub
commit e01edc143a
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 103 additions and 43 deletions

View file

@ -1,6 +1,7 @@
using Microsoft.AspNetCore.Http;
using Microsoft.AspNetCore.Http.Features;
using Microsoft.AspNetCore.WebUtilities;
using Microsoft.Extensions.Hosting;
using Microsoft.Extensions.Logging;
using Microsoft.Extensions.Options;
using Microsoft.Net.Http.Headers;
@ -17,8 +18,9 @@ internal sealed class StreamableHttpHandler(
IOptionsFactory<McpServerOptions> mcpServerOptionsFactory,
IOptions<HttpServerTransportOptions> httpServerTransportOptions,
StatefulSessionManager sessionManager,
ILoggerFactory loggerFactory,
IServiceProvider applicationServices)
IHostApplicationLifetime hostApplicationLifetime,
IServiceProvider applicationServices,
ILoggerFactory loggerFactory)
{
private const string McpSessionIdHeaderName = "Mcp-Session-Id";
@ -60,7 +62,7 @@ internal sealed class StreamableHttpHandler(
}
InitializeSseResponse(context);
var wroteResponse = await session.Transport.HandlePostRequest(message, context.Response.Body, context.RequestAborted);
var wroteResponse = await session.Transport.HandlePostRequestAsync(message, context.Response.Body, context.RequestAborted);
if (!wroteResponse)
{
// We wound up writing nothing, so there should be no Content-Type response header.
@ -94,14 +96,28 @@ internal sealed class StreamableHttpHandler(
return;
}
await using var _ = await session.AcquireReferenceAsync(context.RequestAborted);
InitializeSseResponse(context);
// Link the GET request to both RequestAborted and ApplicationStopping.
// The GET request should complete immediately during graceful shutdown without waiting for
// in-flight POST requests to complete. This prevents slow shutdown when clients are still connected.
using var sseCts = CancellationTokenSource.CreateLinkedTokenSource(context.RequestAborted, hostApplicationLifetime.ApplicationStopping);
var cancellationToken = sseCts.Token;
// We should flush headers to indicate a 200 success quickly, because the initialization response
// will be sent in response to a different POST request. It might be a while before we send a message
// over this response body.
await context.Response.Body.FlushAsync(context.RequestAborted);
await session.Transport.HandleGetRequest(context.Response.Body, context.RequestAborted);
try
{
await using var _ = await session.AcquireReferenceAsync(cancellationToken);
InitializeSseResponse(context);
// We should flush headers to indicate a 200 success quickly, because the initialization response
// will be sent in response to a different POST request. It might be a while before we send a message
// over this response body.
await context.Response.Body.FlushAsync(cancellationToken);
await session.Transport.HandleGetRequestAsync(context.Response.Body, cancellationToken);
}
catch (OperationCanceledException) when (cancellationToken.IsCancellationRequested)
{
// RequestAborted always triggers when the client disconnects before a complete response body is written,
// but this is how SSE connections are typically closed.
}
}
public async Task HandleDeleteRequestAsync(HttpContext context)

View file

@ -175,16 +175,27 @@ internal sealed partial class StreamableHttpClientSessionTransport : TransportBa
request.Headers.Accept.Add(s_textEventStreamMediaType);
CopyAdditionalHeaders(request.Headers, _options.AdditionalHeaders, SessionId, _negotiatedProtocolVersion);
using var response = await _httpClient.SendAsync(request, message: null, _connectionCts.Token).ConfigureAwait(false);
if (!response.IsSuccessStatusCode)
// Server support for the GET request is optional. If it fails, we don't care. It just means we won't receive unsolicited messages.
HttpResponseMessage response;
try
{
response = await _httpClient.SendAsync(request, message: null, _connectionCts.Token).ConfigureAwait(false);
}
catch (HttpRequestException)
{
// Server support for the GET request is optional. If it fails, we don't care. It just means we won't receive unsolicited messages.
return;
}
using var responseStream = await response.Content.ReadAsStreamAsync(_connectionCts.Token).ConfigureAwait(false);
await ProcessSseResponseAsync(responseStream, relatedRpcRequest: null, _connectionCts.Token).ConfigureAwait(false);
using (response)
{
if (!response.IsSuccessStatusCode)
{
return;
}
using var responseStream = await response.Content.ReadAsStreamAsync(_connectionCts.Token).ConfigureAwait(false);
await ProcessSseResponseAsync(responseStream, relatedRpcRequest: null, _connectionCts.Token).ConfigureAwait(false);
}
}
private async Task<JsonRpcMessageWithId?> ProcessSseResponseAsync(Stream responseStream, JsonRpcRequest? relatedRpcRequest, CancellationToken cancellationToken)

View file

@ -40,7 +40,7 @@ public abstract class JsonRpcMessage
/// <remarks>
/// This property should only be set when implementing a custom <see cref="ITransport"/>
/// that needs to pass additional per-message context or to pass a <see cref="JsonRpcMessageContext.User"/>
/// to <see cref="StreamableHttpServerTransport.HandlePostRequest(JsonRpcMessage, Stream, CancellationToken)"/>
/// to <see cref="StreamableHttpServerTransport.HandlePostRequestAsync(JsonRpcMessage, Stream, CancellationToken)"/>
/// or <see cref="SseResponseStreamTransport.OnMessageReceivedAsync(JsonRpcMessage, CancellationToken)"/> .
/// </remarks>
[JsonIgnore]

View file

@ -43,14 +43,14 @@ public sealed class StreamableHttpServerTransport : ITransport
/// <summary>
/// Configures whether the transport should be in stateless mode that does not require all requests for a given session
/// to arrive to the same ASP.NET Core application process. Unsolicited server-to-client messages are not supported in this mode,
/// so calling <see cref="HandleGetRequest(Stream, CancellationToken)"/> results in an <see cref="InvalidOperationException"/>.
/// so calling <see cref="HandleGetRequestAsync(Stream, CancellationToken)"/> results in an <see cref="InvalidOperationException"/>.
/// Server-to-client requests are also unsupported, because the responses may arrive at another ASP.NET Core application process.
/// Client sampling and roots capabilities are also disabled in stateless mode, because the server cannot make requests.
/// </summary>
public bool Stateless { get; init; }
/// <summary>
/// Gets a value indicating whether the execution context should flow from the calls to <see cref="HandlePostRequest(JsonRpcMessage, Stream, CancellationToken)"/>
/// Gets a value indicating whether the execution context should flow from the calls to <see cref="HandlePostRequestAsync(JsonRpcMessage, Stream, CancellationToken)"/>
/// to the corresponding <see cref="JsonRpcMessageContext.ExecutionContext"/> property contained in the <see cref="JsonRpcMessage"/> instances returned by the <see cref="MessageReader"/>.
/// </summary>
/// <remarks>
@ -76,7 +76,7 @@ public sealed class StreamableHttpServerTransport : ITransport
/// <param name="sseResponseStream">The response stream to write MCP JSON-RPC messages as SSE events to.</param>
/// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests. The default is <see cref="CancellationToken.None"/>.</param>
/// <returns>A task representing the send loop that writes JSON-RPC messages to the SSE response stream.</returns>
public async Task HandleGetRequest(Stream sseResponseStream, CancellationToken cancellationToken = default)
public async Task HandleGetRequestAsync(Stream sseResponseStream, CancellationToken cancellationToken = default)
{
Throw.IfNull(sseResponseStream);
@ -111,7 +111,7 @@ public sealed class StreamableHttpServerTransport : ITransport
/// If 's an authenticated <see cref="ClaimsPrincipal"/> sent the message, that can be included in the <see cref="JsonRpcMessage.Context"/>.
/// No other part of the context should be set.
/// </para>
public async Task<bool> HandlePostRequest(JsonRpcMessage message, Stream responseStream, CancellationToken cancellationToken = default)
public async Task<bool> HandlePostRequestAsync(JsonRpcMessage message, Stream responseStream, CancellationToken cancellationToken = default)
{
Throw.IfNull(message);
Throw.IfNull(responseStream);