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

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)