Stop the host when the single session server service finishes (#226)

* Stop the host when the single session server service finishes

* Add test
This commit is contained in:
Stephen Toub 2025-04-07 12:45:17 -04:00 • committed by GitHub
commit 5d7c2a1a0f
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 40 additions and 4 deletions

View file

@ -6,8 +6,22 @@ namespace ModelContextProtocol.Hosting;
/// <summary>
/// Hosted service for a single-session (e.g. stdio) MCP server.
/// </summary>
internal sealed class SingleSessionMcpServerHostedService(IMcpServer session) : BackgroundService
/// <param name="session">The server representing the session being hosted.</param>
/// <param name="lifetime">
/// The host's application lifetime. If available, it will have termination requested when the session's run completes.
/// </param>
internal sealed class SingleSessionMcpServerHostedService(IMcpServer session, IHostApplicationLifetime? lifetime = null) : BackgroundService
{
/// <inheritdoc />
protected override Task ExecuteAsync(CancellationToken stoppingToken) => session.RunAsync(stoppingToken);
protected override async Task ExecuteAsync(CancellationToken stoppingToken)
{
try
{
await session.RunAsync(stoppingToken);
}
finally
{
lifetime?.StopApplication();
}
}
}

View file

@ -1,6 +1,8 @@
using ModelContextProtocol.Protocol.Transport;
using Microsoft.Extensions.DependencyInjection;
using Microsoft.Extensions.DependencyInjection;
using Microsoft.Extensions.Hosting;
using ModelContextProtocol.Protocol.Transport;
using Moq;
using System.IO.Pipelines;
namespace ModelContextProtocol.Tests.Configuration;
@ -19,4 +21,24 @@ public class McpServerBuilderExtensionsTransportsTests
Assert.NotNull(transportType);
Assert.Equal(typeof(StdioServerTransport), transportType.ImplementationType);
}
[Fact]
public async Task HostExecutionShutsDownWhenSingleSessionServerExits()
{
Pipe clientToServerPipe = new(), serverToClientPipe = new();
var builder = Host.CreateEmptyApplicationBuilder(null);
builder.Services
.AddMcpServer()
.WithStreamServerTransport(clientToServerPipe.Reader.AsStream(), serverToClientPipe.Writer.AsStream());
IHost host = builder.Build();
Task t = host.RunAsync(TestContext.Current.CancellationToken);
await Task.Delay(1, TestContext.Current.CancellationToken);
Assert.False(t.IsCompleted);
clientToServerPipe.Writer.Complete();
await t;
}
}