Fix slow shutdown when a Streamable HTTP client is connected (#843)

* Add "Async" suffix to async methods

tests/ModelContextProtocol.AspNetCore.Tests/MapMcpStreamableHttpTests.cs

@@ -181,9 +181,10 @@ public class MapMcpStreamableHttpTests(ITestOutputHelper outputHelper) : MapMcpT
 
         await mcpClient.DisposeAsync();
 
-        // The header should be included in the GET request, the initialized notification, the tools/list call, and the delete request.
-        // The DELETE request won't be sent for Stateless mode due to the lack of an Mcp-Session-Id.
-        Assert.Equal(Stateless ? 3 : 4, protocolVersionHeaderValues.Count);
+        // The GET request might not have started in time, and the DELETE request won't be sent in
+        // Stateless mode due to the lack of an Mcp-Session-Id, but the header should be included in the
+        // initialized notification and the tools/list call at a minimum.
+        Assert.True(protocolVersionHeaderValues.Count > 1);
         Assert.All(protocolVersionHeaderValues, v => Assert.Equal("2025-03-26", v));
     }
 }

tests/ModelContextProtocol.AspNetCore.Tests/MapMcpTests.cs

@@ -8,6 +8,7 @@ using ModelContextProtocol.Protocol;
 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 abstract class MapMcpTests(ITestOutputHelper testOutputHelper) : KestrelI
     }
 
     [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 abstract class MapMcpTests(ITestOutputHelper testOutputHelper) : KestrelI
             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/KestrelInMemoryConnection.cs

@@ -6,28 +6,29 @@ namespace ModelContextProtocol.AspNetCore.Tests.Utils;
 
 public sealed class KestrelInMemoryConnection : ConnectionContext
 {
-    private readonly Pipe _clientToServerPipe = new();
-    private readonly Pipe _serverToClientPipe = new();
     private readonly CancellationTokenSource _connectionClosedCts = new();
     private readonly FeatureCollection _features = new();
 
     public KestrelInMemoryConnection()
     {
+        Pipe clientToServerPipe = new();
+        Pipe serverToClientPipe = new();
+
         ConnectionClosed = _connectionClosedCts.Token;
         Transport = new DuplexPipe
         {
-            Input = _clientToServerPipe.Reader,
-            Output = _serverToClientPipe.Writer,
+            Input = clientToServerPipe.Reader,
+            Output = serverToClientPipe.Writer,
         };
-        Application = new DuplexPipe
+        ClientPipe = new DuplexPipe
         {
-            Input = _serverToClientPipe.Reader,
-            Output = _clientToServerPipe.Writer,
+            Input = serverToClientPipe.Reader,
+            Output = clientToServerPipe.Writer,
         };
-        ClientStream = new DuplexStream(Application, _connectionClosedCts);
+        ClientStream = new DuplexStream(ClientPipe, _connectionClosedCts);
     }
 
-    public IDuplexPipe Application { get; }
+    public IDuplexPipe ClientPipe { get; }
     public Stream ClientStream { get; }
 
     public override IDuplexPipe Transport { get; set; }
@@ -41,8 +42,8 @@ public sealed class KestrelInMemoryConnection : ConnectionContext
     {
         // This is called by Kestrel. The client should dispose the DuplexStream which
         // completes the other half of these pipes.
-        await _serverToClientPipe.Writer.CompleteAsync();
-        await _serverToClientPipe.Reader.CompleteAsync();
+        await Transport.Input.CompleteAsync();
+        await Transport.Output.CompleteAsync();
 
         // Don't bother disposing the _connectionClosedCts, since this is just for testing,
         // and it's annoying to synchronize with DuplexStream.

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,12 +41,9 @@ public sealed class KestrelInMemoryTransport : IConnectionListenerFactory
 
         public async ValueTask<ConnectionContext?> AcceptAsync(CancellationToken cancellationToken = default)
         {
-            if (await acceptQueue.Reader.WaitToReadAsync(cancellationToken))
+            await foreach (var item in acceptQueue.Reader.ReadAllAsync(cancellationToken))
             {
-                while (acceptQueue.Reader.TryRead(out var item))
-                {
-                    return item;
-                }
+                return item;
             }
 
             return null;