Fix process crash when testing Stderr (#1449)

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Author: ericstj

src/ModelContextProtocol.Core/Client/StdioClientSessionTransport.cs

@@ -10,15 +10,17 @@ internal sealed class StdioClientSessionTransport : StreamClientSessionTransport
     private readonly StdioClientTransportOptions _options;
     private readonly Process _process;
     private readonly Queue<string> _stderrRollingLog;
+    private readonly DataReceivedEventHandler _errorHandler;
     private int _cleanedUp = 0;
     private readonly int? _processId;
 
-    public StdioClientSessionTransport(StdioClientTransportOptions options, Process process, string endpointName, Queue<string> stderrRollingLog, ILoggerFactory? loggerFactory) :
+    public StdioClientSessionTransport(StdioClientTransportOptions options, Process process, string endpointName, Queue<string> stderrRollingLog, DataReceivedEventHandler errorHandler, ILoggerFactory? loggerFactory) :
         base(process.StandardInput.BaseStream, process.StandardOutput.BaseStream, encoding: null, endpointName, loggerFactory)
     {
         _options = options;
         _process = process;
         _stderrRollingLog = stderrRollingLog;
+        _errorHandler = errorHandler;
         try { _processId = process.Id; } catch { }
     }
 
@@ -33,7 +35,7 @@ public override async Task SendMessageAsync(JsonRpcMessage message, Cancellation
         {
             // We failed to send due to an I/O error. If the server process has exited, which is then very likely the cause
             // for the I/O error, we should throw an exception for that instead.
-            if (await GetUnexpectedExitExceptionAsync(cancellationToken).ConfigureAwait(false) is Exception processExitException)
+            if (await GetUnexpectedExitExceptionAsync().ConfigureAwait(false) is Exception processExitException)
             {
                 throw processExitException;
             }
@@ -45,15 +47,32 @@ public override async Task SendMessageAsync(JsonRpcMessage message, Cancellation
     /// <inheritdoc/>
     protected override async ValueTask CleanupAsync(Exception? error = null, CancellationToken cancellationToken = default)
     {
-        // Only clean up once.
+        // Only run the full stdio cleanup once (handler detach, process kill, etc.).
+        // If another call is already handling cleanup, cancel the shutdown token
+        // to unblock it (e.g. if it's stuck in WaitForExitAsync) and let it
+        // call SetDisconnected with full StdioClientCompletionDetails.
         if (Interlocked.Exchange(ref _cleanedUp, 1) != 0)
         {
+            CancelShutdown();
             return;
         }
 
         // We've not yet forcefully terminated the server. If it's already shut down, something went wrong,
         // so create an exception with details about that.
-        error ??= await GetUnexpectedExitExceptionAsync(cancellationToken).ConfigureAwait(false);
+        error ??= await GetUnexpectedExitExceptionAsync().ConfigureAwait(false);
+
+        // Ensure all pending ErrorDataReceived events are drained before detaching
+        // the handler. GetUnexpectedExitExceptionAsync does this when HasExited is
+        // true, but there is a narrow window on Linux where the process has closed
+        // stdout (causing EOF in ReadMessagesAsync) yet hasn't been fully reaped,
+        // so HasExited returns false and the drain is skipped. An unconditional
+        // wait here covers that gap. When the drain already happened above, the
+        // call returns immediately.
+        await WaitForProcessExitAsync().ConfigureAwait(false);
+
+        // Detach the stderr handler so no further ErrorDataReceived events
+        // are dispatched during or after process disposal.
+        _process.ErrorDataReceived -= _errorHandler;
 
         // Terminate the server process (or confirm it already exited), then build
         // and publish strongly-typed completion details while the process handle
@@ -78,26 +97,16 @@ protected override async ValueTask CleanupAsync(Exception? error = null, Cancell
         await base.CleanupAsync(error, cancellationToken).ConfigureAwait(false);
     }
 
-    private async ValueTask<Exception?> GetUnexpectedExitExceptionAsync(CancellationToken cancellationToken)
+    private async ValueTask<Exception?> GetUnexpectedExitExceptionAsync()
     {
         if (!StdioClientTransport.HasExited(_process))
         {
             return null;
         }
 
         Debug.Assert(StdioClientTransport.HasExited(_process));
-        try
-        {
-            // The process has exited, but we still need to ensure stderr has been flushed.
-            // WaitForExitAsync only waits for exit; it does not guarantee that all
-            // ErrorDataReceived events have been dispatched. The synchronous WaitForExit()
-            // (no arguments) does ensure that, so call it after WaitForExitAsync completes.
-#if NET
-            await _process.WaitForExitAsync(cancellationToken).ConfigureAwait(false);
-#endif
-            _process.WaitForExit();
-        }
-        catch { }
+
+        await WaitForProcessExitAsync().ConfigureAwait(false);
 
         string errorMessage = "MCP server process exited unexpectedly";
 
@@ -122,6 +131,35 @@ protected override async ValueTask CleanupAsync(Exception? error = null, Cancell
         return new IOException(errorMessage);
     }
 
+    /// <summary>
+    /// Waits for the process to exit within <see cref="StdioClientTransportOptions.ShutdownTimeout"/>
+    /// and flushes pending <see cref="Process.ErrorDataReceived"/> events.
+    /// </summary>
+    /// <remarks>
+    /// On .NET, <c>Process.WaitForExitAsync</c> also waits for asynchronous output readers
+    /// to complete, ensuring all events have been dispatched. On .NET Framework,
+    /// <see cref="Process.WaitForExit(int)"/> does not guarantee this; the parameterless
+    /// <see cref="Process.WaitForExit()"/> overload is needed to flush the event queue.
+    /// This method is idempotent—calling it after the process has already been waited on
+    /// returns immediately.
+    /// </remarks>
+    private async ValueTask WaitForProcessExitAsync()
+    {
+        try
+        {
+#if NET
+            using var timeoutCts = new CancellationTokenSource(_options.ShutdownTimeout);
+            await _process.WaitForExitAsync(timeoutCts.Token).ConfigureAwait(false);
+#else
+            if (_process.WaitForExit((int)_options.ShutdownTimeout.TotalMilliseconds))
+            {
+                _process.WaitForExit();
+            }
+#endif
+        }
+        catch { }
+    }
+
     private StdioClientCompletionDetails BuildCompletionDetails(Exception? error)
     {
         StdioClientCompletionDetails details = new()

src/ModelContextProtocol.Core/Client/StdioClientTransport.cs

@@ -59,6 +59,7 @@ public async Task<ITransport> ConnectAsync(CancellationToken cancellationToken =
 
         Process? process = null;
         bool processStarted = false;
+        DataReceivedEventHandler? errorHandler = null;
 
         string command = _options.Command;
         IList<string>? arguments = _options.Arguments;
@@ -136,7 +137,7 @@ public async Task<ITransport> ConnectAsync(CancellationToken cancellationToken =
             // few lines in a rolling log for use in exceptions.
             const int MaxStderrLength = 10; // keep the last 10 lines of stderr
             Queue<string> stderrRollingLog = new(MaxStderrLength);
-            process.ErrorDataReceived += (sender, args) =>
+            errorHandler = (sender, args) =>
             {
                 string? data = args.Data;
                 if (data is not null)
@@ -151,11 +152,22 @@ public async Task<ITransport> ConnectAsync(CancellationToken cancellationToken =
                         stderrRollingLog.Enqueue(data);
                     }
 
-                    _options.StandardErrorLines?.Invoke(data);
+                    try
+                    {
+                        _options.StandardErrorLines?.Invoke(data);
+                    }
+                    catch (Exception ex)
+                    {
+                        // Prevent exceptions in the user callback from propagating
+                        // to the background thread that dispatches ErrorDataReceived,
+                        // which would crash the process.
+                        LogStderrCallbackFailed(logger, endpointName, ex);
+                    }
 
                     LogReadStderr(logger, endpointName, data);
                 }
             };
+            process.ErrorDataReceived += errorHandler;
 
             // We need both stdin and stdout to use a no-BOM UTF-8 encoding. On .NET Core,
             // we can use ProcessStartInfo.StandardOutputEncoding/StandardInputEncoding, but
@@ -193,14 +205,19 @@ public async Task<ITransport> ConnectAsync(CancellationToken cancellationToken =
 
             process.BeginErrorReadLine();
 
-            return new StdioClientSessionTransport(_options, process, endpointName, stderrRollingLog, _loggerFactory);
+            return new StdioClientSessionTransport(_options, process, endpointName, stderrRollingLog, errorHandler, _loggerFactory);
         }
         catch (Exception ex)
         {
             LogTransportConnectFailed(logger, endpointName, ex);
 
             try
             {
+                if (process is not null && errorHandler is not null)
+                {
+                    process.ErrorDataReceived -= errorHandler;
+                }
+
                 DisposeProcess(process, processStarted, _options.ShutdownTimeout);
             }
             catch (Exception ex2)
@@ -228,18 +245,6 @@ internal static void DisposeProcess(
                     process.KillTree(shutdownTimeout);
                 }
 
-                // Ensure all redirected stderr/stdout events have been dispatched
-                // before disposing. Only the no-arg WaitForExit() guarantees this;
-                // WaitForExit(int) (as used by KillTree) does not.
-                // This should not hang: either the process already exited on its own
-                // (no child processes holding handles), or KillTree killed the entire
-                // process tree. If it does take too long, the test infrastructure's
-                // own timeout will catch it.
-                if (!processRunning && HasExited(process))
-                {
-                    process.WaitForExit();
-                }
-
                 // Invoke the callback while the process handle is still valid,
                 // e.g. to read ExitCode before Dispose() invalidates it.
                 beforeDispose?.Invoke();
@@ -299,6 +304,9 @@ private static string EscapeArgumentString(string argument) =>
     [LoggerMessage(Level = LogLevel.Information, Message = "{EndpointName} received stderr log: '{Data}'.")]
     private static partial void LogReadStderr(ILogger logger, string endpointName, string data);
 
+    [LoggerMessage(Level = LogLevel.Warning, Message = "{EndpointName} StandardErrorLines callback failed.")]
+    private static partial void LogStderrCallbackFailed(ILogger logger, string endpointName, Exception exception);
+
     [LoggerMessage(Level = LogLevel.Information, Message = "{EndpointName} started server process with PID {ProcessId}.")]
     private static partial void LogTransportProcessStarted(ILogger logger, string endpointName, int processId);
 

src/ModelContextProtocol.Core/Client/StreamClientSessionTransport.cs

@@ -15,7 +15,7 @@ internal class StreamClientSessionTransport : TransportBase
     private readonly TextReader _serverOutput;
     private readonly Stream _serverInputStream;
     private readonly SemaphoreSlim _sendLock = new(1, 1);
-    private CancellationTokenSource? _shutdownCts = new();
+    private readonly CancellationTokenSource _shutdownCts = new();
     private Task? _readTask;
 
     /// <summary>
@@ -97,8 +97,16 @@ public override async Task SendMessageAsync(JsonRpcMessage message, Cancellation
     }
 
     /// <inheritdoc/>
-    public override ValueTask DisposeAsync() =>
-        CleanupAsync(cancellationToken: CancellationToken.None);
+    public override async ValueTask DisposeAsync()
+    {
+        await CleanupAsync(cancellationToken: CancellationToken.None).ConfigureAwait(false);
+
+        // Ensure the channel is always completed after disposal, even if CleanupAsync
+        // returned early because another caller (e.g. ReadMessagesAsync) was already
+        // running cleanup. SetDisconnected is idempotent—if the channel was already
+        // completed by the other cleanup path, this is a no-op.
+        SetDisconnected();
+    }
 
     private async Task ReadMessagesAsync(CancellationToken cancellationToken)
     {
@@ -168,15 +176,20 @@ private async Task ProcessMessageAsync(string line, CancellationToken cancellati
         }
     }
 
+    /// <summary>
+    /// Cancels the shutdown token to signal that the transport is shutting down,
+    /// without performing any other cleanup.
+    /// </summary>
+    protected void CancelShutdown()
+    {
+        _shutdownCts.Cancel();
+    }
+
     protected virtual async ValueTask CleanupAsync(Exception? error = null, CancellationToken cancellationToken = default)
     {
         LogTransportShuttingDown(Name);
 
-        if (Interlocked.Exchange(ref _shutdownCts, null) is { } shutdownCts)
-        {
-            await shutdownCts.CancelAsync().ConfigureAwait(false);
-            shutdownCts.Dispose();
-        }
+        await _shutdownCts.CancelAsync().ConfigureAwait(false);
 
         if (Interlocked.Exchange(ref _readTask, null) is Task readTask)
         {

tests/ModelContextProtocol.AspNetCore.Tests/ClientConformanceTests.cs

@@ -82,24 +82,29 @@ public async Task RunConformanceTest(string scenario)
 
         var process = new Process { StartInfo = startInfo };
 
-        process.OutputDataReceived += (sender, e) =>
+        // Protect callbacks with try/catch to prevent ITestOutputHelper from
+        // throwing on a background thread if events arrive after the test completes.
+        DataReceivedEventHandler outputHandler = (sender, e) =>
         {
             if (e.Data != null)
             {
-                _output.WriteLine(e.Data);
+                try { _output.WriteLine(e.Data); } catch { }
                 outputBuilder.AppendLine(e.Data);
             }
         };
 
-        process.ErrorDataReceived += (sender, e) =>
+        DataReceivedEventHandler errorHandler = (sender, e) =>
         {
             if (e.Data != null)
             {
-                _output.WriteLine(e.Data);
+                try { _output.WriteLine(e.Data); } catch { }
                 errorBuilder.AppendLine(e.Data);
             }
         };
 
+        process.OutputDataReceived += outputHandler;
+        process.ErrorDataReceived += errorHandler;
+
         process.Start();
         process.BeginOutputReadLine();
         process.BeginErrorReadLine();
@@ -112,13 +117,18 @@ public async Task RunConformanceTest(string scenario)
         catch (OperationCanceledException)
         {
             process.Kill(entireProcessTree: true);
+            process.OutputDataReceived -= outputHandler;
+            process.ErrorDataReceived -= errorHandler;
             return (
                 Success: false,
                 Output: outputBuilder.ToString(),
                 Error: errorBuilder.ToString() + "\nProcess timed out after 5 minutes and was killed."
             );
         }
 
+        process.OutputDataReceived -= outputHandler;
+        process.ErrorDataReceived -= errorHandler;
+
         var output = outputBuilder.ToString();
         var error = errorBuilder.ToString();
         var success = process.ExitCode == 0 || HasOnlyWarnings(output, error);

tests/ModelContextProtocol.AspNetCore.Tests/ServerConformanceTests.cs

@@ -136,24 +136,29 @@ public async Task RunPendingConformanceTest_ServerSsePolling()
 
         var process = new Process { StartInfo = startInfo };
 
-        process.OutputDataReceived += (sender, e) =>
+        // Protect callbacks with try/catch to prevent ITestOutputHelper from
+        // throwing on a background thread if events arrive after the test completes.
+        DataReceivedEventHandler outputHandler = (sender, e) =>
         {
             if (e.Data != null)
             {
-                output.WriteLine(e.Data);
+                try { output.WriteLine(e.Data); } catch { }
                 outputBuilder.AppendLine(e.Data);
             }
         };
 
-        process.ErrorDataReceived += (sender, e) =>
+        DataReceivedEventHandler errorHandler = (sender, e) =>
         {
             if (e.Data != null)
             {
-                output.WriteLine(e.Data);
+                try { output.WriteLine(e.Data); } catch { }
                 errorBuilder.AppendLine(e.Data);
             }
         };
 
+        process.OutputDataReceived += outputHandler;
+        process.ErrorDataReceived += errorHandler;
+
         process.Start();
         process.BeginOutputReadLine();
         process.BeginErrorReadLine();
@@ -166,13 +171,18 @@ public async Task RunPendingConformanceTest_ServerSsePolling()
         catch (OperationCanceledException)
         {
             process.Kill(entireProcessTree: true);
+            process.OutputDataReceived -= outputHandler;
+            process.ErrorDataReceived -= errorHandler;
             return (
                 Success: false,
                 Output: outputBuilder.ToString(),
                 Error: errorBuilder.ToString() + "\nProcess timed out after 5 minutes and was killed."
             );
         }
 
+        process.OutputDataReceived -= outputHandler;
+        process.ErrorDataReceived -= errorHandler;
+
         return (
             Success: process.ExitCode == 0,
             Output: outputBuilder.ToString(),