Fix slow shutdown when a Streamable HTTP client is connected

Author: halter73

src/ModelContextProtocol.AspNetCore/StreamableHttpHandler.cs

@@ -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";
 
@@ -94,14 +96,28 @@ await WriteJsonRpcErrorAsync(context,
             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.HandleGetRequestAsync(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(sseCts.Token);
+            await session.Transport.HandleGetRequestAsync(context.Response.Body, sseCts.Token);
+        }
+        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)

tests/ModelContextProtocol.AspNetCore.Tests/MapMcpTests.cs

@@ -8,6 +8,7 @@
 using ModelContextProtocol.Server;
 using ModelContextProtocol.Tests.Utils;
 using System.ComponentModel;
+using System.Diagnostics;
 using System.Net;
 using System.Security.Claims;
 
@@ -114,10 +115,9 @@ public async Task Messages_FromNewUser_AreRejected()
     }
 
     [Fact]
-    public async Task ClaimsPrincipal_CanBeInjectedIntoToolMethod()
+    public async Task ClaimsPrincipal_CanBeInjected_IntoToolMethod()
     {
         Builder.Services.AddMcpServer().WithHttpTransport(ConfigureStateless).WithTools<ClaimsPrincipalTools>();
-        Builder.Services.AddHttpContextAccessor();
 
         await using var app = Builder.Build();
 
@@ -211,6 +211,35 @@ public async Task Sampling_DoesNotCloseStream_Prematurely()
             m.Message.Contains("request '2' for method 'sampling/createMessage'"));
     }
 
+    [Fact]
+    public async Task Server_ShutsDownQuickly_WhenClientIsConnected()
+    {
+        Builder.Services.AddMcpServer().WithHttpTransport().WithTools<ClaimsPrincipalTools>();
+
+        await using var app = Builder.Build();
+        app.MapMcp();
+
+        await app.StartAsync(TestContext.Current.CancellationToken);
+
+        // Connect a client which will open a long-running GET request (SSE or Streamable HTTP)
+        await using var mcpClient = await ConnectAsync();
+
+        // Verify the client is connected
+        var tools = await mcpClient.ListToolsAsync(cancellationToken: TestContext.Current.CancellationToken);
+        Assert.NotEmpty(tools);
+
+        // Now measure how long it takes to stop the server
+        var stopwatch = Stopwatch.StartNew();
+        await app.StopAsync(TestContext.Current.CancellationToken);
+        stopwatch.Stop();
+
+        // The server should shut down quickly (within a few seconds). We use 5 seconds as a generous threshold.
+        // This is much less than the default HostOptions.ShutdownTimeout of 30 seconds.
+        Assert.True(stopwatch.Elapsed < TimeSpan.FromSeconds(5),
+            $"Server took {stopwatch.Elapsed.TotalSeconds:F2} seconds to shut down with a connected client. " +
+            "This suggests the GET request is not respecting ApplicationStopping token.");
+    }
+
     private ClaimsPrincipal CreateUser(string name)
         => new ClaimsPrincipal(new ClaimsIdentity(
             [new Claim("name", name), new Claim(ClaimTypes.NameIdentifier, name)],

tests/ModelContextProtocol.AspNetCore.Tests/Utils/KestrelInMemoryTransport.cs

@@ -13,7 +13,11 @@ public sealed class KestrelInMemoryTransport : IConnectionListenerFactory
     public KestrelInMemoryConnection CreateConnection(EndPoint endpoint)
     {
         var connection = new KestrelInMemoryConnection();
-        GetAcceptQueue(endpoint).Writer.TryWrite(connection);
+        if (!GetAcceptQueue(endpoint).Writer.TryWrite(connection))
+        {
+            throw new IOException("The KestrelInMemoryTransport has been shut down.");
+        };
+
         return connection;
     }
 
@@ -37,7 +41,7 @@ private sealed class KestrelInMemoryListener(EndPoint endpoint, Channel<Connecti
 
         public async ValueTask<ConnectionContext?> AcceptAsync(CancellationToken cancellationToken = default)
         {
-            if (await acceptQueue.Reader.WaitToReadAsync(cancellationToken))
+            while (await acceptQueue.Reader.WaitToReadAsync(cancellationToken))
             {
                 while (acceptQueue.Reader.TryRead(out var item))
                 {

tests/ModelContextProtocol.TestSseServer/Program.cs

@@ -370,6 +370,7 @@ private static void HandleStatelessMcp(IApplicationBuilder app)
         var serviceCollection = new ServiceCollection();
         serviceCollection.AddLogging();
         serviceCollection.AddSingleton(app.ApplicationServices.GetRequiredService<ILoggerFactory>());
+        serviceCollection.AddSingleton(app.ApplicationServices.GetRequiredService<IHostApplicationLifetime>());
         serviceCollection.AddSingleton(app.ApplicationServices.GetRequiredService<DiagnosticListener>());
         serviceCollection.AddRoutingCore();