diff --git a/src/ModelContextProtocol/Hosting/SingleSessionMcpServerHostedService.cs b/src/ModelContextProtocol/Hosting/SingleSessionMcpServerHostedService.cs index 42937759..60316bcf 100644 --- a/src/ModelContextProtocol/Hosting/SingleSessionMcpServerHostedService.cs +++ b/src/ModelContextProtocol/Hosting/SingleSessionMcpServerHostedService.cs @@ -6,8 +6,22 @@ namespace ModelContextProtocol.Hosting; /// /// Hosted service for a single-session (e.g. stdio) MCP server. /// -internal sealed class SingleSessionMcpServerHostedService(IMcpServer session) : BackgroundService +/// The server representing the session being hosted. +/// +/// The host's application lifetime. If available, it will have termination requested when the session's run completes. +/// +internal sealed class SingleSessionMcpServerHostedService(IMcpServer session, IHostApplicationLifetime? lifetime = null) : BackgroundService { /// - protected override Task ExecuteAsync(CancellationToken stoppingToken) => session.RunAsync(stoppingToken); + protected override async Task ExecuteAsync(CancellationToken stoppingToken) + { + try + { + await session.RunAsync(stoppingToken); + } + finally + { + lifetime?.StopApplication(); + } + } } diff --git a/tests/ModelContextProtocol.Tests/Configuration/McpServerBuilderExtensionsTransportsTests.cs b/tests/ModelContextProtocol.Tests/Configuration/McpServerBuilderExtensionsTransportsTests.cs index 93d546b0..51663d78 100644 --- a/tests/ModelContextProtocol.Tests/Configuration/McpServerBuilderExtensionsTransportsTests.cs +++ b/tests/ModelContextProtocol.Tests/Configuration/McpServerBuilderExtensionsTransportsTests.cs @@ -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 void WithStdioServerTransport_Sets_Transport() 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; + } }