Fix slow shutdown when a Streamable HTTP client is connected (#843)
* Add "Async" suffix to async methods
This commit is contained in:
parent
d914581471
commit
e01edc143a
4 changed files with 103 additions and 43 deletions
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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]
|
||||
|
|
|
|||
|
|
@ -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);
|
||||
|
|
|
|||
Loading…
Reference in a new issue