From bab74d58957d8358d96ac1ec706bc34260acda70 Mon Sep 17 00:00:00 2001
From: Stephen Halter <halter73@gmail.com>
Date: Tue, 8 Apr 2025 10:20:16 -0700
Subject: [PATCH] Fix StdioServerTransport.DisposeAsync hang (#235)

- It's not just Unix, it's all platforms that can hang
- Windows doesn't hang on Ctrl-C because that triggers
  the end of the Stream, but that doesn't help in IHostApplicationLifetime.StopAsync() scenarios

https://github.com/dotnet/runtime/blob/4e9627e311806330be72b7f7d3660be699878ebd/src/libraries/System.Console/src/System/ConsolePal.Unix.ConsoleStream.cs#L13
https://github.com/dotnet/runtime/blob/4e9627e311806330be72b7f7d3660be699878ebd/src/libraries/System.Console/src/System/ConsolePal.Windows.cs#L1149
---
 .../TestServerWithHosting.csproj              |  2 +-
 .../Transport/StdioServerTransport.cs         | 35 +++++++++-
 .../ModelContextProtocol.Tests.csproj         |  7 ++
 .../StdioServerIntegrationTests.cs            | 66 +++++++++++++++++++
 .../Utils/KestrelInMemoryConnection.cs        | 12 ++--
 5 files changed, 114 insertions(+), 8 deletions(-)
 create mode 100644 tests/ModelContextProtocol.Tests/StdioServerIntegrationTests.cs

diff --git a/samples/TestServerWithHosting/TestServerWithHosting.csproj b/samples/TestServerWithHosting/TestServerWithHosting.csproj
index 0f3918d..d12277e 100644
--- a/samples/TestServerWithHosting/TestServerWithHosting.csproj
+++ b/samples/TestServerWithHosting/TestServerWithHosting.csproj
@@ -2,7 +2,7 @@
 
   <PropertyGroup>
     <OutputType>Exe</OutputType>
-    <TargetFramework>net9.0</TargetFramework>
+    <TargetFrameworks>net9.0;net8.0</TargetFrameworks>
     <ImplicitUsings>enable</ImplicitUsings>
     <Nullable>enable</Nullable>
     <!--
diff --git a/src/ModelContextProtocol/Protocol/Transport/StdioServerTransport.cs b/src/ModelContextProtocol/Protocol/Transport/StdioServerTransport.cs
index b3ae2e4..4fa7912 100644
--- a/src/ModelContextProtocol/Protocol/Transport/StdioServerTransport.cs
+++ b/src/ModelContextProtocol/Protocol/Transport/StdioServerTransport.cs
@@ -60,7 +60,7 @@ public sealed class StdioServerTransport : StreamServerTransport
     /// </para>
     /// </remarks>
     public StdioServerTransport(string serverName, ILoggerFactory? loggerFactory = null)
-        : base(Console.OpenStandardInput(),
+        : base(new CancellableStdinStream(Console.OpenStandardInput()),
                new BufferedStream(Console.OpenStandardOutput()),
                serverName ?? throw new ArgumentNullException(nameof(serverName)),
                loggerFactory)
@@ -73,4 +73,37 @@ public sealed class StdioServerTransport : StreamServerTransport
 
         return serverOptions.ServerInfo?.Name ?? McpServer.DefaultImplementation.Name;
     }
+
+    // Neither WindowsConsoleStream nor UnixConsoleStream respect CancellationTokens or cancel any I/O on Dispose.
+    // WindowsConsoleStream will return an EOS on Ctrl-C, but that is not the only reason the shutdownToken may fire.
+    private sealed class CancellableStdinStream(Stream stdinStream) : Stream
+    {
+        public override bool CanRead => true;
+        public override bool CanSeek => false;
+        public override bool CanWrite => false;
+
+        public override Task<int> ReadAsync(byte[] buffer, int offset, int count, CancellationToken cancellationToken)
+            => stdinStream.ReadAsync(buffer, offset, count, cancellationToken).WaitAsync(cancellationToken);
+
+#if NET
+        public override ValueTask<int> ReadAsync(Memory<byte> buffer, CancellationToken cancellationToken = default)
+        {
+            ValueTask<int> vt = stdinStream.ReadAsync(buffer, cancellationToken);
+            return vt.IsCompletedSuccessfully ? vt : new(vt.AsTask().WaitAsync(cancellationToken));
+        }
+#endif
+
+        // The McpServer shouldn't call flush on the stdin Stream, but it doesn't need to throw just in case.
+        public override void Flush() { }
+
+        public override long Length => throw new NotSupportedException();
+
+        public override long Position { get => throw new NotSupportedException(); set => throw new NotSupportedException(); }
+
+        public override int Read(byte[] buffer, int offset, int count) => throw new NotSupportedException();
+        public override long Seek(long offset, SeekOrigin origin) => throw new NotSupportedException();
+
+        public override void SetLength(long value) => throw new NotSupportedException();
+        public override void Write(byte[] buffer, int offset, int count) => throw new NotSupportedException();
+    }
 }
diff --git a/tests/ModelContextProtocol.Tests/ModelContextProtocol.Tests.csproj b/tests/ModelContextProtocol.Tests/ModelContextProtocol.Tests.csproj
index 16e6a89..a72caf5 100644
--- a/tests/ModelContextProtocol.Tests/ModelContextProtocol.Tests.csproj
+++ b/tests/ModelContextProtocol.Tests/ModelContextProtocol.Tests.csproj
@@ -47,6 +47,7 @@
     <ProjectReference Include="..\..\src\ModelContextProtocol\ModelContextProtocol.csproj" />
     <ProjectReference Include="..\ModelContextProtocol.TestServer\ModelContextProtocol.TestServer.csproj" />
     <ProjectReference Include="..\ModelContextProtocol.TestSseServer\ModelContextProtocol.TestSseServer.csproj" />
+    <ProjectReference Include="..\..\samples\TestServerWithHosting\TestServerWithHosting.csproj" />
   </ItemGroup>
 
   <ItemGroup>
@@ -56,6 +57,12 @@
     <Content Condition="$([MSBuild]::IsOSPlatform('Linux'))" Include="$([System.IO.Path]::GetFullPath('$(ArtifactsBinDir)ModelContextProtocol.TestServer\$(Configuration)'))\$(TargetFramework)\TestServer.dll">
       <CopyToOutputDirectory>PreserveNewest</CopyToOutputDirectory>
     </Content>
+    <Content Condition="$([MSBuild]::IsOSPlatform('Windows'))" Include="$([System.IO.Path]::GetFullPath('$(ArtifactsBinDir)TestServerWithHosting\$(Configuration)'))\$(TargetFramework)\TestServerWithHosting.exe">
+      <CopyToOutputDirectory>PreserveNewest</CopyToOutputDirectory>
+    </Content>
+    <Content Condition="$([MSBuild]::IsOSPlatform('Linux'))" Include="$([System.IO.Path]::GetFullPath('$(ArtifactsBinDir)TestServerWithHosting\$(Configuration)'))\$(TargetFramework)\TestServerWithHosting.dll">
+      <CopyToOutputDirectory>PreserveNewest</CopyToOutputDirectory>
+    </Content>
   </ItemGroup>
 
 </Project>
diff --git a/tests/ModelContextProtocol.Tests/StdioServerIntegrationTests.cs b/tests/ModelContextProtocol.Tests/StdioServerIntegrationTests.cs
new file mode 100644
index 0000000..a28821f
--- /dev/null
+++ b/tests/ModelContextProtocol.Tests/StdioServerIntegrationTests.cs
@@ -0,0 +1,66 @@
+using ModelContextProtocol.Client;
+using ModelContextProtocol.Protocol.Transport;
+using ModelContextProtocol.Tests.Utils;
+using System.Diagnostics;
+using System.Runtime.InteropServices;
+
+namespace ModelContextProtocol.Tests;
+
+public class StdioServerIntegrationTests(ITestOutputHelper testOutputHelper) : LoggedTest(testOutputHelper)
+{
+    public static bool CanSendSigInt { get; } = RuntimeInformation.IsOSPlatform(OSPlatform.Linux) || RuntimeInformation.IsOSPlatform(OSPlatform.OSX);
+    private const int SIGINT = 2;
+
+    [Fact(Skip = "Platform not supported by this test.", SkipUnless = nameof(CanSendSigInt))]
+    public async Task SigInt_DisposesTestServerWithHosting_Gracefully()
+    {
+        using var process = new Process
+        {
+            StartInfo = new ProcessStartInfo
+            {
+                FileName = "dotnet",
+                Arguments = "TestServerWithHosting.dll",
+                RedirectStandardInput = true,
+                RedirectStandardOutput = true,
+                UseShellExecute = false,
+                CreateNoWindow = true,
+            }
+        };
+
+        process.Start();
+
+        await using var streamServerTransport = new StreamServerTransport(
+            process.StandardOutput.BaseStream,
+            process.StandardInput.BaseStream,
+            serverName: "TestServerWithHosting");
+
+        await using var client = await McpClientFactory.CreateAsync(
+            new TestClientTransport(streamServerTransport),
+            loggerFactory: LoggerFactory,
+            cancellationToken: TestContext.Current.CancellationToken);
+
+        // I considered writing a similar test for Windows using Ctrl-C, then saw that dotnet watch doesn't even send a Ctrl-C
+        // signal because it's such a pain without support for CREATE_NEW_PROCESS_GROUP in System.Diagnostics.Process.
+        // https://github.com/dotnet/sdk/blob/43b1c12e3362098a23ca1018503eb56516840b6a/src/BuiltInTools/dotnet-watch/Internal/ProcessRunner.cs#L277-L303
+        // https://github.com/dotnet/runtime/issues/109432, https://github.com/dotnet/runtime/issues/44944
+        Assert.Equal(0, kill(process.Id, SIGINT));
+
+        using var shutdownCts = CancellationTokenSource.CreateLinkedTokenSource(TestContext.Current.CancellationToken);
+        shutdownCts.CancelAfter(TimeSpan.FromSeconds(10));
+        await process.WaitForExitAsync(shutdownCts.Token);
+
+        Assert.True(process.HasExited);
+        Assert.Equal(0, process.ExitCode);
+    }
+
+    [DllImport("libc", SetLastError = true)]
+    private static extern int kill(int pid, int sig);
+
+    private sealed class TestClientTransport(ITransport sessionTransport) : IClientTransport
+    {
+        public string Name => nameof(TestClientTransport);
+
+        public Task<ITransport> ConnectAsync(CancellationToken cancellationToken = default)
+            => Task.FromResult(sessionTransport);
+    }
+}
diff --git a/tests/ModelContextProtocol.Tests/Utils/KestrelInMemoryConnection.cs b/tests/ModelContextProtocol.Tests/Utils/KestrelInMemoryConnection.cs
index d0f1036..c823261 100644
--- a/tests/ModelContextProtocol.Tests/Utils/KestrelInMemoryConnection.cs
+++ b/tests/ModelContextProtocol.Tests/Utils/KestrelInMemoryConnection.cs
@@ -8,8 +8,8 @@ public sealed class KestrelInMemoryConnection : ConnectionContext
 {
     private readonly Pipe _clientToServerPipe = new();
     private readonly Pipe _serverToClientPipe = new();
-    private readonly CancellationTokenSource _connectionClosedCts = new CancellationTokenSource();
-    private readonly IFeatureCollection _features = new FeatureCollection();
+    private readonly CancellationTokenSource _connectionClosedCts = new();
+    private readonly FeatureCollection _features = new();
 
     public KestrelInMemoryConnection()
     {
@@ -37,17 +37,17 @@ public sealed class KestrelInMemoryConnection : ConnectionContext
 
     public override IDictionary<object, object?> Items { get; set; } = new Dictionary<object, object?>();
 
-    public override ValueTask DisposeAsync()
+    public override async ValueTask DisposeAsync()
     {
         // This is called by Kestrel. The client should dispose the DuplexStream which
         // completes the other half of these pipes.
-        _serverToClientPipe.Writer.Complete();
-        _serverToClientPipe.Reader.Complete();
+        await _serverToClientPipe.Writer.CompleteAsync();
+        await _serverToClientPipe.Reader.CompleteAsync();
 
         // Don't bother disposing the _connectionClosedCts, since this is just for testing,
         // and it's annoying to synchronize with DuplexStream.
 
-        return base.DisposeAsync();
+        await base.DisposeAsync();
     }
 
     private class DuplexPipe : IDuplexPipe