Fix process crash when testing Stderr

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 { }
     }
 
@@ -45,16 +47,24 @@ 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);
 
+        // 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
         // is still valid so we can read the exit code.
@@ -89,13 +99,17 @@ protected override async ValueTask CleanupAsync(Exception? error = null, Cancell
         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.
+            // Use a bounded wait: the process is already dead, we're just draining pipe
+            // buffers. If the caller's token is never canceled (e.g. _shutdownCts hasn't
+            // been canceled yet), an unbounded wait here can hang indefinitely when the
+            // threadpool is slow to deliver the stderr EOF callback.
 #if NET
-            await _process.WaitForExitAsync(cancellationToken).ConfigureAwait(false);
+            using var timeoutCts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
+            timeoutCts.CancelAfter(_options.ShutdownTimeout);
+            await _process.WaitForExitAsync(timeoutCts.Token).ConfigureAwait(false);
+#else
+            _process.WaitForExit((int)_options.ShutdownTimeout.TotalMilliseconds);
 #endif
-            _process.WaitForExit();
         }
         catch { }
 

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

@@ -164,6 +164,21 @@ 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()
+    {
+        try
+        {
+            _shutdownCts?.Cancel();
+        }
+        catch (ObjectDisposedException)
+        {
+        }
+    }
+
     protected virtual async ValueTask CleanupAsync(Exception? error = null, CancellationToken cancellationToken = default)
     {
         LogTransportShuttingDown(Name);

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(),