Fix race in KestrelInMemoryConnection (#233)

* Also remove some duplication from WithStdioServerTransport and WithStreamServerTransport
This commit is contained in:
Stephen Halter 2025-04-07 14:49:32 -07:00 • committed by GitHub
commit 3a2e9a228d
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 16 additions and 22 deletions

View file

@ -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"

View file

@ -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 partial class McpServerBuilderExtensions
{
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 partial class McpServerBuilderExtensions
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
}

View file

@ -148,7 +148,7 @@ public class SseClientTransportTests : LoggedTest
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);

View file

@ -43,9 +43,10 @@ public sealed class KestrelInMemoryConnection : ConnectionContext
// 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();
}