From 4c23b45756d07ea8bdda8c1a968caed454cf00ab Mon Sep 17 00:00:00 2001
From: Stephen Toub <stoub@microsoft.com>
Date: Sun, 6 Apr 2025 21:49:22 -0400
Subject: [PATCH 1/2] Stop the host when the single session server service
 finishes

---
 .../Hosting/SingleSessionMcpServerHostedService.cs | 14 ++++++++++++--
 1 file changed, 12 insertions(+), 2 deletions(-)

diff --git a/src/ModelContextProtocol/Hosting/SingleSessionMcpServerHostedService.cs b/src/ModelContextProtocol/Hosting/SingleSessionMcpServerHostedService.cs
index 42937759..677247e9 100644
--- a/src/ModelContextProtocol/Hosting/SingleSessionMcpServerHostedService.cs
+++ b/src/ModelContextProtocol/Hosting/SingleSessionMcpServerHostedService.cs
@@ -6,8 +6,18 @@ namespace ModelContextProtocol.Hosting;
 /// <summary>
 /// Hosted service for a single-session (e.g. stdio) MCP server.
 /// </summary>
-internal sealed class SingleSessionMcpServerHostedService(IMcpServer session) : BackgroundService
+internal sealed class SingleSessionMcpServerHostedService(IMcpServer session, IHostApplicationLifetime lifetime) : 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();
+        }
+    }
 }

From 4b1869bfa2ffa3ed2b79f50b4cc1fcc6ed9f63cd Mon Sep 17 00:00:00 2001
From: Stephen Toub <stoub@microsoft.com>
Date: Mon, 7 Apr 2025 12:25:56 -0400
Subject: [PATCH 2/2] Add test

---
 .../SingleSessionMcpServerHostedService.cs    |  8 ++++--
 ...pServerBuilderExtensionsTransportsTests.cs | 26 +++++++++++++++++--
 2 files changed, 30 insertions(+), 4 deletions(-)

diff --git a/src/ModelContextProtocol/Hosting/SingleSessionMcpServerHostedService.cs b/src/ModelContextProtocol/Hosting/SingleSessionMcpServerHostedService.cs
index 677247e9..60316bcf 100644
--- a/src/ModelContextProtocol/Hosting/SingleSessionMcpServerHostedService.cs
+++ b/src/ModelContextProtocol/Hosting/SingleSessionMcpServerHostedService.cs
@@ -6,7 +6,11 @@ namespace ModelContextProtocol.Hosting;
 /// <summary>
 /// Hosted service for a single-session (e.g. stdio) MCP server.
 /// </summary>
-internal sealed class SingleSessionMcpServerHostedService(IMcpServer session, IHostApplicationLifetime lifetime) : 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 async Task ExecuteAsync(CancellationToken stoppingToken)
@@ -17,7 +21,7 @@ protected override async Task ExecuteAsync(CancellationToken stoppingToken)
         }
         finally
         {
-            lifetime.StopApplication();
+            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;
+    }
 }