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

* Stop the host when the single session server service finishes

* Add test

Author: stephentoub

src/ModelContextProtocol/Hosting/SingleSessionMcpServerHostedService.cs

@@ -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();
+        }
+    }
 }

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;
+    }
 }