Fix race in KestrelInMemoryConnection (#233)

* Also remove some duplication from WithStdioServerTransport and WithStreamServerTransport

Author: halter73

samples/AspNetCoreSseServer/Properties/launchSettings.json

@@ -4,7 +4,6 @@
     "http": {
       "commandName": "Project",
       "dotnetRunMessages": true,
-      "launchBrowser": true,
       "applicationUrl": "http://localhost:3001",
       "environmentVariables": {
         "ASPNETCORE_ENVIRONMENT": "Development"
@@ -13,7 +12,6 @@
     "https": {
       "commandName": "Project",
       "dotnetRunMessages": true,
-      "launchBrowser": true,
       "applicationUrl": "https://localhost:7133;http://localhost:3001",
       "environmentVariables": {
         "ASPNETCORE_ENVIRONMENT": "Development"

src/ModelContextProtocol/Configuration/McpServerBuilderExtensions.cs

@@ -1,4 +1,5 @@
-using Microsoft.Extensions.Logging;
+using Microsoft.Extensions.DependencyInjection.Extensions;
+using Microsoft.Extensions.Logging;
 using Microsoft.Extensions.Options;
 using ModelContextProtocol.Hosting;
 using ModelContextProtocol.Protocol.Transport;
@@ -348,17 +349,8 @@ public static IMcpServerBuilder WithStdioServerTransport(this IMcpServerBuilder
     {
         Throw.IfNull(builder);
 
+        AddSingleSessionServerDependencies(builder.Services);
         builder.Services.AddSingleton<ITransport, StdioServerTransport>();
-        builder.Services.AddHostedService<SingleSessionMcpServerHostedService>();
-
-        builder.Services.AddSingleton(services =>
-        {
-            ITransport serverTransport = services.GetRequiredService<ITransport>();
-            IOptions<McpServerOptions> options = services.GetRequiredService<IOptions<McpServerOptions>>();
-            ILoggerFactory? loggerFactory = services.GetService<ILoggerFactory>();
-
-            return McpServerFactory.Create(serverTransport, options.Value, loggerFactory, services);
-        });
 
         return builder;
     }
@@ -378,19 +370,22 @@ public static IMcpServerBuilder WithStreamServerTransport(
         Throw.IfNull(inputStream);
         Throw.IfNull(outputStream);
 
+        AddSingleSessionServerDependencies(builder.Services);
         builder.Services.AddSingleton<ITransport>(new StreamServerTransport(inputStream, outputStream));
-        builder.Services.AddHostedService<SingleSessionMcpServerHostedService>();
 
-        builder.Services.AddSingleton(services =>
+        return builder;
+    }
+
+    private static void AddSingleSessionServerDependencies(IServiceCollection services)
+    {
+        services.AddHostedService<SingleSessionMcpServerHostedService>();
+        services.TryAddSingleton(services =>
         {
             ITransport serverTransport = services.GetRequiredService<ITransport>();
             IOptions<McpServerOptions> options = services.GetRequiredService<IOptions<McpServerOptions>>();
             ILoggerFactory? loggerFactory = services.GetService<ILoggerFactory>();
-
             return McpServerFactory.Create(serverTransport, options.Value, loggerFactory, services);
         });
-
-        return builder;
     }
     #endregion
 }

tests/ModelContextProtocol.Tests/Transport/SseClientTransportTests.cs

@@ -148,7 +148,7 @@ public async Task SendMessageAsync_Handles_Accepted_Json_RPC_Response()
 
         var eventSourcePipe = new Pipe();
         var eventSourceData = "event: endpoint\r\ndata: /sseendpoint\r\n\r\n"u8;
-        Assert.True(eventSourceData.TryCopyTo(eventSourcePipe.Writer.GetSpan()));
+        eventSourceData.CopyTo(eventSourcePipe.Writer.GetSpan(eventSourceData.Length));
         eventSourcePipe.Writer.Advance(eventSourceData.Length);
         await eventSourcePipe.Writer.FlushAsync(TestContext.Current.CancellationToken);
 

tests/ModelContextProtocol.Tests/Utils/KestrelInMemoryConnection.cs

@@ -43,9 +43,10 @@ public override ValueTask DisposeAsync()
         // completes the other half of these pipes.
         _serverToClientPipe.Writer.Complete();
         _serverToClientPipe.Reader.Complete();
-        // Disposing the CTS without waiting for the client would be problematic
-        // except we always dispose the HttpClient before Kestrel in our tests.
-        _connectionClosedCts.Dispose();
+
+        // Don't bother disposing the _connectionClosedCts, since this is just for testing,
+        // and it's annoying to synchronize with DuplexStream.
+
         return base.DisposeAsync();
     }