Fix HttpClient timeout for long-running tools without event stream store (#1268)

Copilot · halter73 · web-flow · commit fb2daf2f0dce · 2026-02-14T10:19:53.000-05:00
Co-authored-by: copilot-swe-agent[bot] &lt;198982749+Copilot@users.noreply.github.com&gt;
Co-authored-by: halter73 &lt;54385+halter73@users.noreply.github.com&gt;
diff --git a/src/ModelContextProtocol.Core/Server/StreamableHttpPostTransport.cs b/src/ModelContextProtocol.Core/Server/StreamableHttpPostTransport.cs
@@ -74,6 +74,13 @@ public async ValueTask<bool> HandlePostAsync(JsonRpcMessage message, Cancellatio
             {
                 await _httpSseWriter.WriteAsync(primingItem.Value, cancellationToken).ConfigureAwait(false);
             }
+            else
+            {
+                // 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.
+                await responseStream.FlushAsync(cancellationToken).ConfigureAwait(false);
+            }
 
             // Ensure that we've sent the priming event before processing the incoming request.
             await parentTransport.MessageWriter.WriteAsync(message, cancellationToken).ConfigureAwait(false);
diff --git a/tests/ModelContextProtocol.AspNetCore.Tests/MapMcpTests.cs b/tests/ModelContextProtocol.AspNetCore.Tests/MapMcpTests.cs
@@ -231,6 +231,50 @@ public async Task Server_ShutsDownQuickly_WhenClientIsConnected()
             "This suggests the GET request is not respecting ApplicationStopping token.");
     }
 
+    [Fact]
+    public async Task LongRunningToolCall_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.
+
+        Builder.Services.AddMcpServer().WithHttpTransport(ConfigureStateless).WithTools<LongRunningTools>();
+
+        await using var app = Builder.Build();
+        app.MapMcp();
+        await app.StartAsync(TestContext.Current.CancellationToken);
+
+        // Create a custom HttpClient with a very short timeout (1 second)
+        // The tool will take 2 seconds to complete
+        using var shortTimeoutClient = new HttpClient(SocketsHttpHandler)
+        {
+            BaseAddress = new Uri("http://localhost:5000/"),
+            Timeout = TimeSpan.FromSeconds(1)
+        };
+
+        var path = UseStreamableHttp ? "/" : "/sse";
+        var transportMode = UseStreamableHttp ? HttpTransportMode.StreamableHttp : HttpTransportMode.Sse;
+
+        await using var transport = new HttpClientTransport(new()
+        {
+            Endpoint = new($"http://localhost:5000{path}"),
+            TransportMode = transportMode,
+        }, shortTimeoutClient, LoggerFactory);
+
+        await using var mcpClient = await McpClient.CreateAsync(transport, loggerFactory: LoggerFactory, cancellationToken: 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
+        var response = await mcpClient.CallToolAsync(
+            "long_running_operation",
+            new Dictionary<string, object?>() { ["durationMs"] = 2000 },
+            cancellationToken: TestContext.Current.CancellationToken);
+
+        var content = Assert.Single(response.Content.OfType<TextContentBlock>());
+        Assert.Equal("Operation completed after 2000ms", content.Text);
+    }
+
     private ClaimsPrincipal CreateUser(string name)
         => new(new ClaimsIdentity(
             [new Claim("name", name), new Claim(ClaimTypes.NameIdentifier, name)],
@@ -288,4 +332,17 @@ public static async Task<string> SamplingToolAsync(McpServer server, string prom
             return $"Sampling completed successfully. Client responded: {Assert.IsType<TextContentBlock>(Assert.Single(samplingResult.Content)).Text}";
         }
     }
+
+    [McpServerToolType]
+    protected class LongRunningTools
+    {
+        [McpServerTool, Description("Simulates a long-running operation")]
+        public static async Task<string> LongRunningOperation(
+            [Description("Duration of the operation in milliseconds")] int durationMs,
+            CancellationToken cancellationToken)
+        {
+            await Task.Delay(durationMs, cancellationToken);
+            return $"Operation completed after {durationMs}ms";
+        }
+    }
 }
diff --git a/tests/ModelContextProtocol.AspNetCore.Tests/OAuth/AuthTests.cs b/tests/ModelContextProtocol.AspNetCore.Tests/OAuth/AuthTests.cs
@@ -1,5 +1,6 @@
 using Microsoft.AspNetCore.Authentication;
 using Microsoft.AspNetCore.Authentication.JwtBearer;
+using Microsoft.AspNetCore.Authorization;
 using Microsoft.AspNetCore.Builder;
 using Microsoft.AspNetCore.Http;
 using Microsoft.AspNetCore.WebUtilities;
@@ -8,10 +9,12 @@
 using ModelContextProtocol.AspNetCore.Authentication;
 using ModelContextProtocol.Authentication;
 using ModelContextProtocol.Client;
+using ModelContextProtocol.Protocol;
 using ModelContextProtocol.Server;
 using System.Net;
 using System.Net.Http.Json;
 using System.Security.Claims;
+using System.Text.Json;
 using Xunit.Sdk;
 
 namespace ModelContextProtocol.AspNetCore.Tests.OAuth;
@@ -211,32 +214,46 @@ public async Task CanAuthenticate_WithTokenRefresh()
     {
         var hasForcedRefresh = false;
 
-        Builder.Services.AddHttpContextAccessor();
         Builder.Services.AddMcpServer(options =>
+        {
+            options.ToolCollection = new();
+        });
+
+        await using var app = await StartMcpServerAsync(configureMiddleware: app =>
+        {
+            // Add middleware to intercept list tools requests and force a token refresh on the first call
+            app.Use(async (context, next) =>
             {
-                options.ToolCollection = new();
-            })
-            .AddListToolsFilter(next =>
-            {
-                return async (mcpContext, cancellationToken) =>
+                if (context.Request.Method == HttpMethods.Post && context.Request.Path == "/" && !hasForcedRefresh)
                 {
-                    if (!hasForcedRefresh)
+                    // Enable buffering so we can read the request body multiple times
+                    context.Request.EnableBuffering();
+
+                    // Read the request body to check if it's calling tools/list
+                    var message = await JsonSerializer.DeserializeAsync(
+                        context.Request.Body,
+                        McpJsonUtilities.DefaultOptions.GetTypeInfo(typeof(JsonRpcMessage)),
+                        context.RequestAborted) as JsonRpcMessage;
+
+                    // Reset the request body position so MapMcp can read it
+                    context.Request.Body.Position = 0;
+
+                    // Check if this is a tools/list request
+                    if (message is JsonRpcRequest request && request.Method == "tools/list")
                     {
                         hasForcedRefresh = true;
 
-                        var httpContext = mcpContext.Services!.GetRequiredService<IHttpContextAccessor>().HttpContext!;
-                        await httpContext.ChallengeAsync(JwtBearerDefaults.AuthenticationScheme);
-                        await httpContext.Response.CompleteAsync();
-                        throw new Exception("This exception will not impact the client because the response has already been completed.");
-                    }
-                    else
-                    {
-                        return await next(mcpContext, cancellationToken);
+                        // Return 401 to force token refresh
+                        await context.ChallengeAsync(JwtBearerDefaults.AuthenticationScheme);
+                        await context.Response.StartAsync(context.RequestAborted);
+                        await context.Response.Body.FlushAsync(context.RequestAborted);
+                        return; // Short-circuit, don't call next()
                     }
-                };
-            });
+                }
 
-        await using var app = await StartMcpServerAsync();
+                await next(context);
+            });
+        });
 
         await using var transport = new HttpClientTransport(new()
         {
@@ -451,29 +468,66 @@ public async Task AuthorizationFlow_UsesScopeFromForbiddenHeader()
     {
         var adminScopes = "admin:read admin:write";
 
-        Builder.Services.AddHttpContextAccessor();
         Builder.Services.AddMcpServer()
             .WithTools([
                 McpServerTool.Create([McpServerTool(Name = "admin-tool")]
-                async (IServiceProvider serviceProvider, ClaimsPrincipal user) =>
+                (ClaimsPrincipal user) =>
                 {
-                    if (!user.HasClaim("scope", adminScopes))
-                    {
-                        var httpContext = serviceProvider.GetRequiredService<IHttpContextAccessor>().HttpContext!;
-                        httpContext.Response.StatusCode = StatusCodes.Status403Forbidden;
-                        httpContext.Response.Headers.WWWAuthenticate = $"Bearer error=\"insufficient_scope\", resource_metadata=\"{McpServerUrl}/.well-known/oauth-protected-resource\", scope=\"{adminScopes}\"";
-                        await httpContext.Response.CompleteAsync();
-
-                        throw new Exception("This exception will not impact the client because the response has already been completed.");
-                    }
-
+                    // Tool now just checks if user has the required scopes
+                    // If they don't, it shouldn't get here due to middleware
+                    Assert.True(user.HasClaim("scope", adminScopes), "User should have admin scopes when tool executes");
                     return "Admin tool executed.";
                 }),
             ]);
 
         string? requestedScope = null;
 
-        await using var app = await StartMcpServerAsync();
+        await using var app = await StartMcpServerAsync(configureMiddleware: app =>
+        {
+            // Add middleware to intercept requests and check for admin-tool calls
+            app.Use(async (context, next) =>
+            {
+                if (context.Request.Method == HttpMethods.Post && context.Request.Path == "/")
+                {
+                    // Enable buffering so we can read the request body multiple times
+                    context.Request.EnableBuffering();
+
+                    // Read the request body to check if it's calling admin-tool
+                    var message = await JsonSerializer.DeserializeAsync(
+                        context.Request.Body,
+                        McpJsonUtilities.DefaultOptions.GetTypeInfo(typeof(JsonRpcMessage)),
+                        context.RequestAborted) as JsonRpcMessage;
+
+                    // Reset the request body position so MapMcp can read it
+                    context.Request.Body.Position = 0;
+
+                    // Check if this is a tools/call request for admin-tool
+                    if (message is JsonRpcRequest request && request.Method == "tools/call")
+                    {
+                        var toolCallParams = JsonSerializer.Deserialize(
+                            request.Params,
+                            McpJsonUtilities.DefaultOptions.GetTypeInfo(typeof(CallToolRequestParams))) as CallToolRequestParams;
+
+                        if (toolCallParams?.Name == "admin-tool")
+                        {
+                            // Check if user has required scopes
+                            var user = context.User;
+                            if (!user.HasClaim("scope", adminScopes))
+                            {
+                                // User lacks required scopes, return 403 before MapMcp processes the request
+                                context.Response.StatusCode = StatusCodes.Status403Forbidden;
+                                context.Response.Headers.WWWAuthenticate = $"Bearer error=\"insufficient_scope\", resource_metadata=\"{McpServerUrl}/.well-known/oauth-protected-resource\", scope=\"{adminScopes}\"";
+                                await context.Response.StartAsync(context.RequestAborted);
+                                await context.Response.Body.FlushAsync(context.RequestAborted);
+                                return; // Short-circuit, don't call next()
+                            }
+                        }
+                    }
+                }
+
+                await next(context);
+            });
+        });
 
         await using var transport = new HttpClientTransport(new()
         {
diff --git a/tests/ModelContextProtocol.AspNetCore.Tests/OAuth/OAuthTestBase.cs b/tests/ModelContextProtocol.AspNetCore.Tests/OAuth/OAuthTestBase.cs
@@ -81,7 +81,7 @@ public async ValueTask DisposeAsync()
         }
     }
 
-    protected async Task<WebApplication> StartMcpServerAsync(string path = "", string? authScheme = null)
+    protected async Task<WebApplication> StartMcpServerAsync(string path = "", string? authScheme = null, Action<WebApplication>? configureMiddleware = null)
     {
         // Wait for the OAuth server to be ready before starting the MCP server.
         // This prevents race conditions in CI where the OAuth server may not be
@@ -94,6 +94,10 @@ protected async Task<WebApplication> StartMcpServerAsync(string path = "", strin
         });
 
         var app = Builder.Build();
+        
+        // Allow tests to add custom middleware before MapMcp
+        configureMiddleware?.Invoke(app);
+        
         app.MapMcp(path).RequireAuthorization(new AuthorizeAttribute
         {
             AuthenticationSchemes = authScheme
diff --git a/tests/ModelContextProtocol.AspNetCore.Tests/StreamableHttpServerConformanceTests.cs b/tests/ModelContextProtocol.AspNetCore.Tests/StreamableHttpServerConformanceTests.cs
@@ -350,7 +350,13 @@ public async Task DeleteRequest_CompletesSession_WhichCancelsLongRunningToolCall
         await CallInitializeAndValidateAsync();
 
         Task<HttpResponseMessage> CallLongRunningToolAsync() =>
-            HttpClient.PostAsync("", JsonContent(CallTool("long-running")), TestContext.Current.CancellationToken);
+            HttpClient.SendAsync(
+                new HttpRequestMessage(HttpMethod.Post, "")
+                {
+                    Content = JsonContent(CallTool("long-running"))
+                },
+                HttpCompletionOption.ResponseHeadersRead,
+                TestContext.Current.CancellationToken);
 
         var longRunningToolTasks = new Task<HttpResponseMessage>[10];
         for (int i = 0; i < longRunningToolTasks.Length; i++)
@@ -360,25 +366,28 @@ Task<HttpResponseMessage> CallLongRunningToolAsync() =>
 
         var getResponse = await HttpClient.GetAsync("", HttpCompletionOption.ResponseHeadersRead, TestContext.Current.CancellationToken);
 
-        for (int i = 0; i < longRunningToolTasks.Length; i++)
+        // Wait for all long-running tool calls to receive 200 response headers before sending DELETE
+        var responseHeaders = await Task.WhenAll(longRunningToolTasks);
+        foreach (var response in responseHeaders)
         {
-            Assert.False(longRunningToolTasks[i].IsCompleted);
+            Assert.Equal(HttpStatusCode.OK, response.StatusCode);
         }
 
+        // Now send DELETE to cancel the session
         await HttpClient.DeleteAsync("", TestContext.Current.CancellationToken);
 
         // Get request should complete gracefully.
         var sseResponseBody = await getResponse.Content.ReadAsStringAsync(TestContext.Current.CancellationToken);
         Assert.Empty(sseResponseBody);
 
-        // Currently, the OCE thrown by the canceled session is unhandled and turned into a 500 error by Kestrel.
+        // Currently, responses are flushed immediately to prevent HttpClient timeouts for long-running requests.
+        // This means the response starts with a 200 status code. When the session is canceled, Kestrel closes
+        // the connection without writing the chunk terminator, causing an HttpRequestException when reading the response body.
         // The spec suggests sending CancelledNotifications. That would be good, but we can do that later.
-        // For now, the important thing is that request completes without indicating success.
-        await Task.WhenAll(longRunningToolTasks);
-        foreach (var task in longRunningToolTasks)
+        // For now, the important thing is that reading the response body fails.
+        foreach (var response in responseHeaders)
         {
-            var response = await task;
-            Assert.False(response.IsSuccessStatusCode);
+            await Assert.ThrowsAsync<HttpRequestException>(async () => await response.Content.ReadAsByteArrayAsync(TestContext.Current.CancellationToken));
         }
     }
 

Original file line number	Diff line number	Diff line change
`@@ -74,6 +74,13 @@ public async ValueTask<bool> HandlePostAsync(JsonRpcMessage message, Cancellatio`
`74`	`74`	`{`
`75`	`75`	`await _httpSseWriter.WriteAsync(primingItem.Value, cancellationToken).ConfigureAwait(false);`
`76`	`76`	`}`
	`77`	`+ else`
	`78`	`+ {`
	`79`	`+ // If there's no priming write, flush the stream to ensure HTTP response headers are`
	`80`	`+ // sent to the client now that the server is ready to process the request.`
	`81`	`+ // This prevents HttpClient timeout for long-running requests.`
	`82`	`+ await responseStream.FlushAsync(cancellationToken).ConfigureAwait(false);`
	`83`	`+ }`
`77`	`84`
`78`	`85`	`// Ensure that we've sent the priming event before processing the incoming request.`
`79`	`86`	`await parentTransport.MessageWriter.WriteAsync(message, cancellationToken).ConfigureAwait(false);`
Original file line number	Diff line number	Diff line change
`@@ -81,7 +81,7 @@ public async ValueTask DisposeAsync()`
`81`	`81`	`}`
`82`	`82`	`}`
`83`	`83`
`84`		`- protected async Task<WebApplication> StartMcpServerAsync(string path = "", string? authScheme = null)`
	`84`	`+ protected async Task<WebApplication> StartMcpServerAsync(string path = "", string? authScheme = null, Action<WebApplication>? configureMiddleware = null)`
`85`	`85`	`{`
`86`	`86`	`// Wait for the OAuth server to be ready before starting the MCP server.`
`87`	`87`	`// This prevents race conditions in CI where the OAuth server may not be`
`@@ -94,6 +94,10 @@ protected async Task<WebApplication> StartMcpServerAsync(string path = "", strin`
`94`	`94`	`});`
`95`	`95`
`96`	`96`	`var app = Builder.Build();`
	`97`	`+`
	`98`	`+ // Allow tests to add custom middleware before MapMcp`
	`99`	`+ configureMiddleware?.Invoke(app);`
	`100`	`+`
`97`	`101`	`app.MapMcp(path).RequireAuthorization(new AuthorizeAttribute`
`98`	`102`	`{`
`99`	`103`	`AuthenticationSchemes = authScheme`