From a88c02d36c139421183b5c8c51f0701514d5ab98 Mon Sep 17 00:00:00 2001 From: Stephen Halter Date: Thu, 6 Aug 2026 16:42:31 -0700 Subject: [PATCH] Make Streamable HTTP status mapping deterministic Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- .../StreamableHttpResponseStartFeature.cs | 26 ++++ .../ModelContextProtocol.AspNetCore.csproj | 1 + .../StreamableHttpHandler.cs | 16 +-- .../ModelContextProtocol.Core.csproj | 1 + .../Server/StreamableHttpPostTransport.cs | 132 +++++++----------- .../StreamableHttpResponseStartOptions.cs | 8 ++ .../Server/StreamableHttpServerTransport.cs | 38 ++--- .../MapMcpTests.cs | 15 +- .../RawHttpConformanceTests.cs | 74 +++++++++- 9 files changed, 182 insertions(+), 129 deletions(-) create mode 100644 src/Common/StreamableHttpResponseStartFeature.cs create mode 100644 src/ModelContextProtocol.Core/Server/StreamableHttpResponseStartOptions.cs diff --git a/src/Common/StreamableHttpResponseStartFeature.cs b/src/Common/StreamableHttpResponseStartFeature.cs new file mode 100644 index 000000000..b28b8b80d --- /dev/null +++ b/src/Common/StreamableHttpResponseStartFeature.cs @@ -0,0 +1,26 @@ +using ModelContextProtocol.Protocol; + +namespace ModelContextProtocol; + +internal static class StreamableHttpResponseStartFeature +{ + private const string ResponseStartingCallbackKey = "ModelContextProtocol.StreamableHttp.ResponseStartingCallback"; + + public static void Set(JsonRpcMessage message, Action callback) + { + message.Context ??= new(); + (message.Context.Items ??= new Dictionary())[ResponseStartingCallbackKey] = callback; + } + + public static Action? Take(JsonRpcMessage message) + { + if (message.Context?.Items is not { } items || + !items.TryGetValue(ResponseStartingCallbackKey, out var value)) + { + return null; + } + + items.Remove(ResponseStartingCallbackKey); + return (Action)value!; + } +} diff --git a/src/ModelContextProtocol.AspNetCore/ModelContextProtocol.AspNetCore.csproj b/src/ModelContextProtocol.AspNetCore/ModelContextProtocol.AspNetCore.csproj index 4c46acdd1..fe0f0d4f9 100644 --- a/src/ModelContextProtocol.AspNetCore/ModelContextProtocol.AspNetCore.csproj +++ b/src/ModelContextProtocol.AspNetCore/ModelContextProtocol.AspNetCore.csproj @@ -26,6 +26,7 @@ + diff --git a/src/ModelContextProtocol.AspNetCore/StreamableHttpHandler.cs b/src/ModelContextProtocol.AspNetCore/StreamableHttpHandler.cs index f0b0b1a12..056a1afd9 100644 --- a/src/ModelContextProtocol.AspNetCore/StreamableHttpHandler.cs +++ b/src/ModelContextProtocol.AspNetCore/StreamableHttpHandler.cs @@ -153,14 +153,13 @@ await WriteJsonRpcErrorAsync(context, await using var _ = await session.AcquireReferenceAsync(context.RequestAborted); - Func? onResponseStarting = null; if (RequiresPerRequestMetadataProtocol(context)) { // SEP-2575 maps some JSON-RPC error codes onto HTTP statuses (404 for a method the server - // does not implement, 400 for missing-capability and unsupported-version rejections). The - // status line can only be chosen before the first response byte, so the transport defers - // its eager header flush and reports the first response message here. - onResponseStarting = firstMessage => + // does not implement, 400 for missing-capability and unsupported-version rejections). + // Wait for the actual first JSON-RPC message so this mapping is deterministic regardless + // of how long the request handler takes. + StreamableHttpResponseStartFeature.Set(message, firstMessage => { if (firstMessage is JsonRpcError { Error: { } errorDetail } && !context.Response.HasStarted) { @@ -173,13 +172,12 @@ await WriteJsonRpcErrorAsync(context, _ => context.Response.StatusCode, }; } - - return default; - }; + }); } InitializeSseResponse(context); - var wroteResponse = await session.Transport.HandlePostRequestAsync(message, context.Response.Body, onResponseStarting, 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. diff --git a/src/ModelContextProtocol.Core/ModelContextProtocol.Core.csproj b/src/ModelContextProtocol.Core/ModelContextProtocol.Core.csproj index 3fbef0377..c37352f46 100644 --- a/src/ModelContextProtocol.Core/ModelContextProtocol.Core.csproj +++ b/src/ModelContextProtocol.Core/ModelContextProtocol.Core.csproj @@ -34,6 +34,7 @@ + diff --git a/src/ModelContextProtocol.Core/Server/StreamableHttpPostTransport.cs b/src/ModelContextProtocol.Core/Server/StreamableHttpPostTransport.cs index 95411f7e2..7e9e5e297 100644 --- a/src/ModelContextProtocol.Core/Server/StreamableHttpPostTransport.cs +++ b/src/ModelContextProtocol.Core/Server/StreamableHttpPostTransport.cs @@ -16,13 +16,14 @@ internal sealed partial class StreamableHttpPostTransport( Stream responseStream, CancellationToken sessionCancellationToken, ILogger logger, - Func? onResponseStarting = null) : ITransport + StreamableHttpResponseStartOptions? responseStartOptions = null) : ITransport { private readonly SemaphoreSlim _messageLock = new(1, 1); private readonly TaskCompletionSource _httpResponseTcs = new(TaskCreationOptions.RunContinuationsAsynchronously); private readonly SseEventWriter _httpSseWriter = new(responseStream); private TaskCompletionSource? _storeStreamTcs; + private List>? _pendingHttpSseItems; #pragma warning disable MCP9006 // Stateful Streamable HTTP resumability types are obsolete but still wired up internally. private ISseEventStreamWriter? _storeSseWriter; #pragma warning restore MCP9006 @@ -76,114 +77,74 @@ public async ValueTask HandlePostAsync(JsonRpcMessage message, Cancellatio return false; } - CancellationTokenSource? deferredFlushCts = null; - Task? deferredFlushTask = null; - bool deferHeaderFlush = false; using (await _messageLock.LockAsync(cancellationToken).ConfigureAwait(false)) { var primingItem = await TryStartSseEventStreamAsync(_pendingRequest).ConfigureAwait(false); if (primingItem.HasValue) { - await NotifyResponseStartingAsync(firstMessage: null).ConfigureAwait(false); - await _httpSseWriter.WriteAsync(primingItem.Value, cancellationToken).ConfigureAwait(false); + if (responseStartOptions is null) + { + StartHttpResponse(firstMessage: null); + await _httpSseWriter.WriteAsync(primingItem.Value, cancellationToken).ConfigureAwait(false); + } + else + { + QueuePendingHttpSseItem(primingItem.Value); + } } - else if (onResponseStarting is null) + else if (responseStartOptions is null) { // If there's no priming write, flush the stream to ensure HTTP response headers are // sent to the client now that the server is ready to process the request. // This prevents HttpClient timeout for long-running requests. + StartHttpResponse(firstMessage: null); await responseStream.FlushAsync(cancellationToken).ConfigureAwait(false); } - else - { - deferHeaderFlush = true; - } - // Ensure that we've sent the priming event before processing the incoming request. + // Legacy responses send the priming event before processing the request. A response + // waiting for its first JSON-RPC message keeps the event queued until that message arrives. await parentTransport.MessageWriter.WriteAsync(message, cancellationToken).ConfigureAwait(false); } - if (deferHeaderFlush) - { - // Defer the flush (and the header commit it implies) so the callback can still choose - // the HTTP status line for an immediate JSON-RPC error. Start the bounded grace period - // only after the request has been queued for dispatch. - deferredFlushCts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken); - deferredFlushTask = DeferredHeaderFlushAsync(deferredFlushCts.Token); - } - - try - { - // 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); - } - finally - { - if (deferredFlushCts is not null) - { - deferredFlushCts.Cancel(); - await deferredFlushTask!.ConfigureAwait(false); - deferredFlushCts.Dispose(); - } - } + // 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; } /// - /// Bounds the deferred header flush: after a short grace window, flushes the response headers - /// if no response message has been written yet. Immediate rejections land well inside the - /// window, so the response-starting callback can still map their JSON-RPC error codes onto the - /// HTTP status line; a handler that runs longer commits the headers here so clients see them - /// promptly (long-running tool calls must not trip HttpClient's response timeout). + /// Notifies the HTTP application exactly once, immediately before the first response write. /// - private async Task DeferredHeaderFlushAsync(CancellationToken cancellationToken) + private void StartHttpResponse(JsonRpcMessage? firstMessage) { - try - { - await Task.Delay(DeferredHeaderFlushGrace, cancellationToken).ConfigureAwait(false); - using var _ = await _messageLock.LockAsync(cancellationToken).ConfigureAwait(false); - if (!_httpResponseStarted && !_httpResponseCompleted) - { - await NotifyResponseStartingAsync(firstMessage: null).ConfigureAwait(false); - await responseStream.FlushAsync(cancellationToken).ConfigureAwait(false); - } - } - catch (OperationCanceledException) when (cancellationToken.IsCancellationRequested) + if (_httpResponseStarted) { - // The response was written or the request ended before the grace window elapsed. + return; } - catch (Exception ex) + + _httpResponseStarted = true; + if (responseStartOptions is not null) { - // Surface the failure to the awaiting HandlePostAsync when possible. If the response - // future has already been resolved (the response started or completed on another path), - // TrySetException is a no-op, so log here to keep the deferred-flush failure diagnosable. - if (!_httpResponseTcs.TrySetException(ex)) - { - LogDeferredHeaderFlushFailed(ex); - } + responseStartOptions.OnResponseStarting( + firstMessage ?? throw new InvalidOperationException("A JSON-RPC message is required to start this response.")); } } - /// How long the response-header flush may be deferred waiting for the first response message. - internal static readonly TimeSpan DeferredHeaderFlushGrace = TimeSpan.FromMilliseconds(250); + private void QueuePendingHttpSseItem(SseItem item) + => (_pendingHttpSseItems ??= []).Add(item); - /// - /// Invokes the response-starting callback exactly once, immediately before the first write to - /// the HTTP response stream, so the HTTP application can still set the response status line. - /// - private async ValueTask NotifyResponseStartingAsync(JsonRpcMessage? firstMessage) + private async ValueTask WritePendingHttpSseItemsAsync(CancellationToken cancellationToken) { - if (_httpResponseStarted) + if (_pendingHttpSseItems is not { } pendingItems) { return; } - _httpResponseStarted = true; - if (onResponseStarting is not null) + _pendingHttpSseItems = null; + foreach (var item in pendingItems) { - await onResponseStarting(firstMessage).ConfigureAwait(false); + await _httpSseWriter.WriteAsync(item, cancellationToken).ConfigureAwait(false); } } @@ -222,7 +183,8 @@ public async Task SendMessageAsync(JsonRpcMessage message, CancellationToken can try { - await NotifyResponseStartingAsync(message).ConfigureAwait(false); + StartHttpResponse(message); + await WritePendingHttpSseItemsAsync(cancellationToken).ConfigureAwait(false); await _httpSseWriter.WriteAsync(item, cancellationToken).ConfigureAwait(false); } catch (Exception ex) when (!cancellationToken.IsCancellationRequested) @@ -266,8 +228,15 @@ public async ValueTask EnablePollingAsync(TimeSpan retryInterval, CancellationTo // Write to the response stream if it still exists. if (!_httpResponseCompleted) { - await NotifyResponseStartingAsync(firstMessage: null).ConfigureAwait(false); - await _httpSseWriter.WriteAsync(primingItem, cancellationToken).ConfigureAwait(false); + if (responseStartOptions is not null && !_httpResponseStarted) + { + QueuePendingHttpSseItem(primingItem); + } + else + { + StartHttpResponse(firstMessage: null); + 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. @@ -276,8 +245,12 @@ public async ValueTask EnablePollingAsync(TimeSpan retryInterval, CancellationTo await _storeSseWriter.SetModeAsync(SseEventStreamMode.Polling, cancellationToken).ConfigureAwait(false); #pragma warning restore MCP9006 - // Signal completion so HandlePostAsync can return. - _httpResponseTcs.TrySetResult(true); + // A response waiting for its first JSON-RPC message cannot complete until that message has + // started the response and flushed any queued priming events. + if (responseStartOptions is null || _httpResponseStarted) + { + _httpResponseTcs.TrySetResult(true); + } } private async ValueTask?> TryStartSseEventStreamAsync(RequestId requestId) @@ -342,7 +315,4 @@ public async ValueTask DisposeAsync() [LoggerMessage(Level = LogLevel.Warning, Message = "Failed to dispose SSE event stream writer.")] private partial void LogStoreStreamDisposalFailed(Exception exception); - - [LoggerMessage(Level = LogLevel.Warning, Message = "Failed to flush deferred Streamable HTTP response headers.")] - private partial void LogDeferredHeaderFlushFailed(Exception exception); } diff --git a/src/ModelContextProtocol.Core/Server/StreamableHttpResponseStartOptions.cs b/src/ModelContextProtocol.Core/Server/StreamableHttpResponseStartOptions.cs new file mode 100644 index 000000000..571c04066 --- /dev/null +++ b/src/ModelContextProtocol.Core/Server/StreamableHttpResponseStartOptions.cs @@ -0,0 +1,8 @@ +using ModelContextProtocol.Protocol; + +namespace ModelContextProtocol.Server; + +internal sealed class StreamableHttpResponseStartOptions(Action onResponseStarting) +{ + public Action OnResponseStarting { get; } = onResponseStarting; +} diff --git a/src/ModelContextProtocol.Core/Server/StreamableHttpServerTransport.cs b/src/ModelContextProtocol.Core/Server/StreamableHttpServerTransport.cs index f143eaaa7..045b0a8ed 100644 --- a/src/ModelContextProtocol.Core/Server/StreamableHttpServerTransport.cs +++ b/src/ModelContextProtocol.Core/Server/StreamableHttpServerTransport.cs @@ -207,39 +207,19 @@ public async Task HandleGetRequestAsync(Stream sseResponseStream, CancellationTo /// If an authenticated sent the message, that can be included in the . /// No other part of the context should be set. /// - public Task HandlePostRequestAsync(JsonRpcMessage message, Stream responseStream, CancellationToken cancellationToken = default) - => HandlePostRequestAsync(message, responseStream, onResponseStarting: null, cancellationToken); - - /// - /// Handles a Streamable HTTP POST request, processing the JSON-RPC message and writing any - /// JSON-RPC responses to the response stream. - /// This overload additionally reports the first JSON-RPC message written to the response via - /// , before any response bytes are written, so the HTTP - /// application can still choose the response status line (SEP-2575 maps some JSON-RPC error - /// codes to HTTP statuses). When is provided, the eager - /// response-header flush that normally precedes request processing is deferred until that first - /// message; the callback receives when the first write is not a JSON-RPC - /// message (e.g. a resumability priming event). - /// The status line can only be influenced by the FIRST write: when a handler streams a - /// notification (e.g. progress) before failing, or runs past the transport's bounded - /// header-flush grace window, the status is already committed and a later JSON-RPC error - /// rides the committed status. - /// - /// The JSON-RPC message to process. - /// The response stream to write any JSON-RPC responses to. - /// Callback invoked once, immediately before the first write to . - /// The to monitor for cancellation requests. The default is . - /// - /// if data was written to the response body. - /// if nothing was written because the request body did not contain any messages to respond to. - /// - /// or is . - public async Task HandlePostRequestAsync(JsonRpcMessage message, Stream responseStream, Func? onResponseStarting, CancellationToken cancellationToken) + public async Task HandlePostRequestAsync( + JsonRpcMessage message, + Stream responseStream, + CancellationToken cancellationToken = default) { Throw.IfNull(message); Throw.IfNull(responseStream); - var postTransport = new StreamableHttpPostTransport(this, responseStream, _transportDisposedCts.Token, _logger, onResponseStarting); + var onResponseStarting = StreamableHttpResponseStartFeature.Take(message); + var responseStartOptions = onResponseStarting is null + ? null + : new StreamableHttpResponseStartOptions(onResponseStarting); + var postTransport = new StreamableHttpPostTransport(this, responseStream, _transportDisposedCts.Token, _logger, responseStartOptions); using var postCts = CancellationTokenSource.CreateLinkedTokenSource(_transportDisposedCts.Token, cancellationToken); await using (postTransport.ConfigureAwait(false)) { diff --git a/tests/ModelContextProtocol.AspNetCore.Tests/MapMcpTests.cs b/tests/ModelContextProtocol.AspNetCore.Tests/MapMcpTests.cs index ef6832101..9c316e7b1 100644 --- a/tests/ModelContextProtocol.AspNetCore.Tests/MapMcpTests.cs +++ b/tests/ModelContextProtocol.AspNetCore.Tests/MapMcpTests.cs @@ -237,12 +237,11 @@ public async Task Server_ShutsDownQuickly_WhenClientIsConnected() } [Fact] - public async Task LongRunningToolCall_DoesNotTimeout_WhenNoEventStreamStore() + public async Task LegacyLongRunningToolCall_DoesNotTimeout_WhenNoEventStreamStore() { - // Regression test for: Tool calls that last over HttpClient timeout without producing - // intermediate notifications will timeout because HttpClient doesn't see the 200 response - // until the first message is written. When primingItem is null (no ISseEventStreamStore), - // we should flush the response stream so HttpClient sees the 200 immediately. + // Legacy protocol revisions do not map JSON-RPC errors onto the HTTP status line, so the + // response headers should be flushed before a long-running handler emits its first message. + // The 2026-07-28 revision intentionally waits for that message to map SEP-2575 statuses. Builder.Services.AddMcpServer().WithHttpTransport(ConfigureStateless).WithTools(); @@ -274,7 +273,11 @@ public async Task LongRunningToolCall_DoesNotTimeout_WhenNoEventStreamStore() TransportMode = transportMode, }, shortTimeoutClient, LoggerFactory); - await using var mcpClient = await McpClient.CreateAsync(transport, loggerFactory: LoggerFactory, cancellationToken: TestContext.Current.CancellationToken); + await using var mcpClient = await McpClient.CreateAsync( + transport, + new McpClientOptions { ProtocolVersion = McpProtocolVersions.November2025ProtocolVersion }, + LoggerFactory, + TestContext.Current.CancellationToken); // Call a tool that takes 2 seconds - this should succeed despite the 1 second HttpClient timeout // because the response stream is flushed immediately after receiving the request diff --git a/tests/ModelContextProtocol.AspNetCore.Tests/RawHttpConformanceTests.cs b/tests/ModelContextProtocol.AspNetCore.Tests/RawHttpConformanceTests.cs index 8520f929c..a7526f482 100644 --- a/tests/ModelContextProtocol.AspNetCore.Tests/RawHttpConformanceTests.cs +++ b/tests/ModelContextProtocol.AspNetCore.Tests/RawHttpConformanceTests.cs @@ -23,8 +23,18 @@ public class RawHttpConformanceTests(ITestOutputHelper outputHelper) : KestrelIn private WebApplication? _app; - private async Task StartAsync(string? protocolVersion = null) + private async Task StartAsync(string? protocolVersion = null, McpServerTool? additionalTool = null) { + List tools = + [ + McpServerTool.Create((string text) => $"echo:{text}", new() { Name = "echo" }), + ]; + + if (additionalTool is not null) + { + tools.Add(additionalTool); + } + Builder.Services .AddMcpServer(options => { @@ -32,7 +42,7 @@ private async Task StartAsync(string? protocolVersion = null) options.ProtocolVersion = protocolVersion; }) .WithHttpTransport() - .WithTools([McpServerTool.Create((string text) => $"echo:{text}", new() { Name = "echo" })]) + .WithTools(tools) .WithTools(); _app = Builder.Build(); @@ -187,7 +197,7 @@ public async Task July2026Post_RemovedMethod_Returns404_WithMethodNotFound(strin } [Fact] - public async Task July2026Post_MissingRequiredCapability_Returns400() + public async Task July2026Post_SlowMissingRequiredCapability_Returns400() { await StartAsync(); @@ -207,6 +217,58 @@ public async Task July2026Post_MissingRequiredCapability_Returns400() Assert.Equal((int)McpErrorCode.MissingRequiredClientCapability, json["error"]!["code"]!.GetValue()); } + [Fact] + public async Task LegacyPost_LongRunningHandler_FlushesResponseHeadersPromptly() + { + var handlerStarted = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + var releaseHandler = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + var longRunningTool = McpServerTool.Create( + async (CancellationToken cancellationToken) => + { + handlerStarted.TrySetResult(true); + await releaseHandler.Task.WaitAsync(cancellationToken); + return "released"; + }, + new() { Name = "wait_for_release" }); + await StartAsync(additionalTool: longRunningTool); + + var initializeBody = @"{""jsonrpc"":""2.0"",""id"":1,""method"":""initialize"",""params"":{""protocolVersion"":""2025-11-25"",""capabilities"":{},""clientInfo"":{""name"":""initialize-handshake"",""version"":""1.0""}}}"; + using (var initializeRequest = new HttpRequestMessage(HttpMethod.Post, "") { Content = JsonContent(initializeBody) }) + using (var initializeResponse = await HttpClient.SendAsync(initializeRequest, TestContext.Current.CancellationToken)) + { + Assert.Equal(HttpStatusCode.OK, initializeResponse.StatusCode); + } + + var body = @"{""jsonrpc"":""2.0"",""id"":2,""method"":""tools/call"",""params"":{""name"":""wait_for_release"",""arguments"":{}}}"; + using var request = new HttpRequestMessage(HttpMethod.Post, "") { Content = JsonContent(body) }; + request.Headers.Add(ProtocolVersionHeader, McpProtocolVersions.November2025ProtocolVersion); + var responseTask = HttpClient.SendAsync( + request, + HttpCompletionOption.ResponseHeadersRead, + TestContext.Current.CancellationToken); + + await handlerStarted.Task.WaitAsync(TestConstants.DefaultTimeout, TestContext.Current.CancellationToken); + + HttpResponseMessage response; + try + { + response = await responseTask.WaitAsync( + TestConstants.HttpClientPollingTimeout, + TestContext.Current.CancellationToken); + } + finally + { + releaseHandler.TrySetResult(true); + } + + using (response) + { + Assert.Equal(HttpStatusCode.OK, response.StatusCode); + var json = await ReadJsonResponseAsync(response, TestContext.Current.CancellationToken); + Assert.Equal("released", json["result"]!["content"]![0]!["text"]!.GetValue()); + } + } + [Fact] public async Task ServerDiscover_WithConfiguredPerRequestMetadataProtocol_ReturnsOnlyConfiguredVersion() { @@ -501,9 +563,13 @@ public async Task July2026Post_MalformedClientCapabilities_Returns400_WithInvali private sealed class CapabilityTools { [McpServerTool(Name = "requires_sampling")] - public static string RequiresSampling() => + public static async Task RequiresSampling(CancellationToken cancellationToken) + { + // Reproduce the old 250 ms header-grace race before returning the SEP-2575 error. + await Task.Delay(TimeSpan.FromMilliseconds(500), cancellationToken); throw new MissingRequiredClientCapabilityException( new ClientCapabilities { Sampling = new() }, "sampling capability required but not declared by client"); + } } }