Improve exception diagnostics for stdio client (#376)

* Improve exception diagnostics for stdio client

Track the last several lines of stderr and include those in an exception created as part of CleanupAsync if the server process has already exited. This is based on an assumption that the server should never exit prior to CleanupAsync being called.

* Address feedback and fix test failure
This commit is contained in:
Stephen Toub 2025-05-14 10:46:49 -04:00 • committed by GitHub
commit c11c634945
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
11 changed files with 222 additions and 131 deletions

View file

@ -1,37 +0,0 @@
// Licensed to the .NET Foundation under one or more agreements.
// The .NET Foundation licenses this file to you under the MIT license.
namespace System.Diagnostics;
internal static class ProcessExtensions
{
public static void Kill(this Process process, bool entireProcessTree)
{
_ = entireProcessTree;
process.Kill();
}
public static async Task WaitForExitAsync(this Process process, CancellationToken cancellationToken = default)
{
if (process.HasExited)
{
return;
}
var tcs = new TaskCompletionSource<bool>();
void ProcessExitedHandler(object? sender, EventArgs e) => tcs.TrySetResult(true);
try
{
process.EnableRaisingEvents = true;
process.Exited += ProcessExitedHandler;
using var _ = cancellationToken.Register(() => tcs.TrySetCanceled(cancellationToken));
await tcs.Task.ConfigureAwait(false);
}
finally
{
process.Exited -= ProcessExitedHandler;
}
}
}

View file

@ -137,7 +137,7 @@ internal sealed partial class SseClientSessionTransport : TransportBase
} }
finally finally
{ {
SetConnected(false); SetDisconnected();
} }
} }
@ -203,7 +203,7 @@ internal sealed partial class SseClientSessionTransport : TransportBase
} }
finally finally
{ {
SetConnected(false); SetDisconnected();
} }
} }
@ -251,7 +251,7 @@ internal sealed partial class SseClientSessionTransport : TransportBase
_messageEndpoint = new Uri(_sseEndpoint, data); _messageEndpoint = new Uri(_sseEndpoint, data);
// Set connected state // Set connected state
SetConnected(true); SetConnected();
_connectionEstablished.TrySetResult(true); _connectionEstablished.TrySetResult(true);
} }

View file

@ -9,54 +9,51 @@ internal sealed class StdioClientSessionTransport : StreamClientSessionTransport
{ {
private readonly StdioClientTransportOptions _options; private readonly StdioClientTransportOptions _options;
private readonly Process _process; private readonly Process _process;
private readonly Queue<string> _stderrRollingLog;
private int _cleanedUp = 0;
public StdioClientSessionTransport(StdioClientTransportOptions options, Process process, string endpointName, ILoggerFactory? loggerFactory) public StdioClientSessionTransport(StdioClientTransportOptions options, Process process, string endpointName, Queue<string> stderrRollingLog, ILoggerFactory? loggerFactory)
: base(process.StandardInput, process.StandardOutput, endpointName, loggerFactory) : base(process.StandardInput, process.StandardOutput, endpointName, loggerFactory)
{ {
_process = process; _process = process;
_options = options; _options = options;
_stderrRollingLog = stderrRollingLog;
} }
/// <inheritdoc/> /// <inheritdoc/>
/// <remarks>
/// <para>
/// For stdio-based transports, this implementation first verifies that the underlying process
/// is still running before attempting to send the message. If the process has exited or cannot
/// be accessed, a <see cref="InvalidOperationException"/> is thrown with details about the failure.
/// </para>
/// <para>
/// After verifying the process state, this method delegates to the base class implementation
/// to handle the actual message serialization and transmission to the process's standard input stream.
/// </para>
/// </remarks>
/// <exception cref="InvalidOperationException">
/// Thrown when the underlying process has exited or cannot be accessed.
/// </exception>
public override async Task SendMessageAsync(JsonRpcMessage message, CancellationToken cancellationToken = default) public override async Task SendMessageAsync(JsonRpcMessage message, CancellationToken cancellationToken = default)
{ {
Exception? processException = null;
bool hasExited = false;
try try
{ {
hasExited = _process.HasExited; await base.SendMessageAsync(message, cancellationToken);
} }
catch (Exception e) catch (IOException)
{ {
processException = e; // We failed to send due to an I/O error. If the server process has exited, which is then very likely the cause
hasExited = true; // for the I/O error, we should throw an exception for that instead.
} if (await GetUnexpectedExitExceptionAsync(cancellationToken).ConfigureAwait(false) is Exception processExitException)
{
throw processExitException;
}
if (hasExited) throw;
{
throw new InvalidOperationException("Transport is not connected", processException);
} }
await base.SendMessageAsync(message, cancellationToken).ConfigureAwait(false);
} }
/// <inheritdoc/> /// <inheritdoc/>
protected override ValueTask CleanupAsync(CancellationToken cancellationToken) protected override async ValueTask CleanupAsync(Exception? error = null, CancellationToken cancellationToken = default)
{ {
// Only clean up once.
if (Interlocked.Exchange(ref _cleanedUp, 1) != 0)
{
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);
// Now terminate the server process.
try try
{ {
StdioClientTransport.DisposeProcess(_process, processRunning: true, _options.ShutdownTimeout, Name); StdioClientTransport.DisposeProcess(_process, processRunning: true, _options.ShutdownTimeout, Name);
@ -66,6 +63,49 @@ internal sealed class StdioClientSessionTransport : StreamClientSessionTransport
LogTransportShutdownFailed(Name, ex); LogTransportShutdownFailed(Name, ex);
} }
return base.CleanupAsync(cancellationToken); // And handle cleanup in the base type.
await base.CleanupAsync(error, cancellationToken);
}
private async ValueTask<Exception?> GetUnexpectedExitExceptionAsync(CancellationToken cancellationToken)
{
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.
#if NET
await _process.WaitForExitAsync(cancellationToken).ConfigureAwait(false);
#else
_process.WaitForExit();
#endif
}
catch { }
string errorMessage = "MCP server process exited unexpectedly";
string? exitCode = null;
try
{
exitCode = $" (exit code: {(uint)_process.ExitCode})";
}
catch { }
lock (_stderrRollingLog)
{
if (_stderrRollingLog.Count > 0)
{
errorMessage =
$"{errorMessage}{exitCode}{Environment.NewLine}" +
$"Server's stderr tail:{Environment.NewLine}" +
$"{string.Join(Environment.NewLine, _stderrRollingLog)}";
}
}
return new IOException(errorMessage);
} }
} }

View file

@ -127,8 +127,28 @@ public sealed partial class StdioClientTransport : IClientTransport
process = new() { StartInfo = startInfo }; process = new() { StartInfo = startInfo };
// Set up error logging // Set up stderr handling. Log all stderr output, and keep the last
process.ErrorDataReceived += (sender, args) => LogReadStderr(logger, endpointName, args.Data ?? "(no data)"); // 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) =>
{
string? data = args.Data;
if (data is not null)
{
lock (stderrRollingLog)
{
if (stderrRollingLog.Count >= MaxStderrLength)
{
stderrRollingLog.Dequeue();
}
stderrRollingLog.Enqueue(data);
}
LogReadStderr(logger, endpointName, data);
}
};
// We need both stdin and stdout to use a no-BOM UTF-8 encoding. On .NET Core, // We need both stdin and stdout to use a no-BOM UTF-8 encoding. On .NET Core,
// we can use ProcessStartInfo.StandardOutputEncoding/StandardInputEncoding, but // we can use ProcessStartInfo.StandardOutputEncoding/StandardInputEncoding, but
@ -154,14 +174,14 @@ public sealed partial class StdioClientTransport : IClientTransport
if (!processStarted) if (!processStarted)
{ {
LogTransportProcessStartFailed(logger, endpointName); LogTransportProcessStartFailed(logger, endpointName);
throw new InvalidOperationException("Failed to start MCP server process"); throw new IOException("Failed to start MCP server process.");
} }
LogTransportProcessStarted(logger, endpointName, process.Id); LogTransportProcessStarted(logger, endpointName, process.Id);
process.BeginErrorReadLine(); process.BeginErrorReadLine();
return new StdioClientSessionTransport(_options, process, endpointName, _loggerFactory); return new StdioClientSessionTransport(_options, process, endpointName, stderrRollingLog, _loggerFactory);
} }
catch (Exception ex) catch (Exception ex)
{ {
@ -176,7 +196,7 @@ public sealed partial class StdioClientTransport : IClientTransport
LogTransportShutdownFailed(logger, endpointName, ex2); LogTransportShutdownFailed(logger, endpointName, ex2);
} }
throw new InvalidOperationException("Failed to connect transport", ex); throw new IOException("Failed to connect transport.", ex);
} }
} }
@ -185,20 +205,9 @@ public sealed partial class StdioClientTransport : IClientTransport
{ {
if (process is not null) if (process is not null)
{ {
if (processRunning)
{
try
{
processRunning = !process.HasExited;
}
catch
{
processRunning = false;
}
}
try try
{ {
processRunning = processRunning && !HasExited(process);
if (processRunning) if (processRunning)
{ {
// Wait for the process to exit. // Wait for the process to exit.
@ -214,6 +223,19 @@ public sealed partial class StdioClientTransport : IClientTransport
} }
} }
/// <summary>Gets whether <paramref name="process"/> has exited.</summary>
internal static bool HasExited(Process process)
{
try
{
return process.HasExited;
}
catch
{
return true;
}
}
[LoggerMessage(Level = LogLevel.Information, Message = "{EndpointName} connecting.")] [LoggerMessage(Level = LogLevel.Information, Message = "{EndpointName} connecting.")]
private static partial void LogTransportConnecting(ILogger logger, string endpointName); private static partial void LogTransportConnecting(ILogger logger, string endpointName);

View file

@ -53,17 +53,12 @@ internal class StreamClientSessionTransport : TransportBase
_readTask = readTask.Unwrap(); _readTask = readTask.Unwrap();
readTask.Start(); readTask.Start();
SetConnected(true); SetConnected();
} }
/// <inheritdoc/> /// <inheritdoc/>
public override async Task SendMessageAsync(JsonRpcMessage message, CancellationToken cancellationToken = default) public override async Task SendMessageAsync(JsonRpcMessage message, CancellationToken cancellationToken = default)
{ {
if (!IsConnected)
{
throw new InvalidOperationException("Transport is not connected");
}
string id = "(no id)"; string id = "(no id)";
if (message is JsonRpcMessageWithId messageWithId) if (message is JsonRpcMessageWithId messageWithId)
{ {
@ -82,31 +77,22 @@ internal class StreamClientSessionTransport : TransportBase
catch (Exception ex) catch (Exception ex)
{ {
LogTransportSendFailed(Name, id, ex); LogTransportSendFailed(Name, id, ex);
throw new InvalidOperationException("Failed to send message", ex); throw new IOException("Failed to send message.", ex);
} }
} }
/// <inheritdoc/> /// <inheritdoc/>
/// <summary>
/// Asynchronously releases all resources used by the stream client session transport.
/// </summary>
/// <returns>A task that represents the asynchronous dispose operation.</returns>
/// <remarks>
/// This method cancels ongoing operations and waits for the read task to complete
/// before marking the transport as disconnected. It calls <see cref="CleanupAsync"/>
/// to perform the actual cleanup work.
/// After disposal, the transport can no longer be used to send or receive messages.
/// </remarks>
public override ValueTask DisposeAsync() => public override ValueTask DisposeAsync() =>
CleanupAsync(CancellationToken.None); CleanupAsync(cancellationToken: CancellationToken.None);
private async Task ReadMessagesAsync(CancellationToken cancellationToken) private async Task ReadMessagesAsync(CancellationToken cancellationToken)
{ {
Exception? error = null;
try try
{ {
LogTransportEnteringReadMessagesLoop(Name); LogTransportEnteringReadMessagesLoop(Name);
while (!cancellationToken.IsCancellationRequested) while (true)
{ {
if (await _serverOutput.ReadLineAsync(cancellationToken).ConfigureAwait(false) is not string line) if (await _serverOutput.ReadLineAsync(cancellationToken).ConfigureAwait(false) is not string line)
{ {
@ -130,12 +116,13 @@ internal class StreamClientSessionTransport : TransportBase
} }
catch (Exception ex) catch (Exception ex)
{ {
error = ex;
LogTransportReadMessagesFailed(Name, ex); LogTransportReadMessagesFailed(Name, ex);
} }
finally finally
{ {
_readTask = null; _readTask = null;
await CleanupAsync(cancellationToken).ConfigureAwait(false); await CleanupAsync(error, cancellationToken).ConfigureAwait(false);
} }
} }
@ -166,7 +153,7 @@ internal class StreamClientSessionTransport : TransportBase
} }
} }
protected virtual async ValueTask CleanupAsync(CancellationToken cancellationToken) protected virtual async ValueTask CleanupAsync(Exception? error = null, CancellationToken cancellationToken = default)
{ {
LogTransportShuttingDown(Name); LogTransportShuttingDown(Name);
@ -191,7 +178,7 @@ internal class StreamClientSessionTransport : TransportBase
} }
} }
SetConnected(false); SetDisconnected(error);
LogTransportShutDown(Name); LogTransportShutDown(Name);
} }
} }

View file

@ -51,7 +51,7 @@ public class StreamServerTransport : TransportBase
_inputReader = new StreamReader(inputStream, Encoding.UTF8); _inputReader = new StreamReader(inputStream, Encoding.UTF8);
_outputStream = outputStream; _outputStream = outputStream;
SetConnected(true); SetConnected();
_readLoopCompleted = Task.Run(ReadMessagesAsync, _shutdownCts.Token); _readLoopCompleted = Task.Run(ReadMessagesAsync, _shutdownCts.Token);
} }
@ -60,7 +60,7 @@ public class StreamServerTransport : TransportBase
{ {
if (!IsConnected) if (!IsConnected)
{ {
throw new InvalidOperationException("Transport is not connected"); return;
} }
using var _ = await _sendLock.LockAsync(cancellationToken).ConfigureAwait(false); using var _ = await _sendLock.LockAsync(cancellationToken).ConfigureAwait(false);
@ -80,13 +80,14 @@ public class StreamServerTransport : TransportBase
catch (Exception ex) catch (Exception ex)
{ {
LogTransportSendFailed(Name, id, ex); LogTransportSendFailed(Name, id, ex);
throw new InvalidOperationException("Failed to send message", ex); throw new IOException("Failed to send message.", ex);
} }
} }
private async Task ReadMessagesAsync() private async Task ReadMessagesAsync()
{ {
CancellationToken shutdownToken = _shutdownCts.Token; CancellationToken shutdownToken = _shutdownCts.Token;
Exception? error = null;
try try
{ {
LogTransportEnteringReadMessagesLoop(Name); LogTransportEnteringReadMessagesLoop(Name);
@ -140,10 +141,11 @@ public class StreamServerTransport : TransportBase
catch (Exception ex) catch (Exception ex)
{ {
LogTransportReadMessagesFailed(Name, ex); LogTransportReadMessagesFailed(Name, ex);
error = ex;
} }
finally finally
{ {
SetConnected(false); SetDisconnected(error);
} }
} }
@ -183,7 +185,7 @@ public class StreamServerTransport : TransportBase
} }
finally finally
{ {
SetConnected(false); SetDisconnected();
LogTransportShutDown(Name); LogTransportShutDown(Name);
} }

View file

@ -45,7 +45,7 @@ internal sealed partial class StreamableHttpClientSessionTransport : TransportBa
// We connect with the initialization request with the MCP transport. This means that any errors won't be observed // We connect with the initialization request with the MCP transport. This means that any errors won't be observed
// until the first call to SendMessageAsync. Fortunately, that happens internally in McpClientFactory.ConnectAsync // until the first call to SendMessageAsync. Fortunately, that happens internally in McpClientFactory.ConnectAsync
// so we still throw any connection-related Exceptions from there and never expose a pre-connected client to the user. // so we still throw any connection-related Exceptions from there and never expose a pre-connected client to the user.
SetConnected(true); SetConnected();
} }
/// <inheritdoc/> /// <inheritdoc/>
@ -139,7 +139,7 @@ internal sealed partial class StreamableHttpClientSessionTransport : TransportBa
} }
finally finally
{ {
SetConnected(false); SetDisconnected();
} }
} }

View file

@ -1,7 +1,8 @@
using System.Diagnostics;
using System.Threading.Channels; using System.Threading.Channels;
using Microsoft.Extensions.Logging; using Microsoft.Extensions.Logging;
using ModelContextProtocol.Protocol.Messages;
using Microsoft.Extensions.Logging.Abstractions; using Microsoft.Extensions.Logging.Abstractions;
using ModelContextProtocol.Protocol.Messages;
namespace ModelContextProtocol.Protocol.Transport; namespace ModelContextProtocol.Protocol.Transport;
@ -23,7 +24,14 @@ public abstract partial class TransportBase : ITransport
{ {
private readonly Channel<JsonRpcMessage> _messageChannel; private readonly Channel<JsonRpcMessage> _messageChannel;
private readonly ILogger _logger; private readonly ILogger _logger;
private int _isConnected; private volatile int _state = StateInitial;
/// <summary>The transport has not yet been connected.</summary>
private const int StateInitial = 0;
/// <summary>The transport is connected.</summary>
private const int StateConnected = 1;
/// <summary>The transport was previously connected and is now disconnected.</summary>
private const int StateDisconnected = 2;
/// <summary> /// <summary>
/// Initializes a new instance of the <see cref="TransportBase"/> class. /// Initializes a new instance of the <see cref="TransportBase"/> class.
@ -53,7 +61,7 @@ public abstract partial class TransportBase : ITransport
protected string Name { get; } protected string Name { get; }
/// <inheritdoc/> /// <inheritdoc/>
public bool IsConnected => _isConnected == 1; public bool IsConnected => _state == StateConnected;
/// <inheritdoc/> /// <inheritdoc/>
public ChannelReader<JsonRpcMessage> MessageReader => _messageChannel.Reader; public ChannelReader<JsonRpcMessage> MessageReader => _messageChannel.Reader;
@ -73,7 +81,7 @@ public abstract partial class TransportBase : ITransport
{ {
if (!IsConnected) if (!IsConnected)
{ {
throw new InvalidOperationException("Transport is not connected"); throw new InvalidOperationException("Transport is not connected.");
} }
if (_logger.IsEnabled(LogLevel.Debug)) if (_logger.IsEnabled(LogLevel.Debug))
@ -82,24 +90,61 @@ public abstract partial class TransportBase : ITransport
LogTransportReceivedMessage(Name, messageId); LogTransportReceivedMessage(Name, messageId);
} }
await _messageChannel.Writer.WriteAsync(message, cancellationToken).ConfigureAwait(false); bool wrote = _messageChannel.Writer.TryWrite(message);
Debug.Assert(wrote || !IsConnected, "_messageChannel is unbounded; this should only ever return false if the channel has been closed.");
} }
/// <summary> /// <summary>
/// Sets the connected state of the transport. /// Sets the transport to a connected state.
/// </summary> /// </summary>
/// <param name="isConnected">Whether the transport is connected.</param> protected void SetConnected()
protected void SetConnected(bool isConnected)
{ {
var newIsConnected = isConnected ? 1 : 0; while (true)
if (Interlocked.Exchange(ref _isConnected, newIsConnected) == newIsConnected)
{ {
return; int state = _state;
} switch (state)
{
case StateInitial:
if (Interlocked.CompareExchange(ref _state, StateConnected, StateInitial) == StateInitial)
{
return;
}
break;
if (!isConnected) case StateConnected:
return;
case StateDisconnected:
throw new IOException("Transport is already disconnected and can't be reconnected.");
default:
Debug.Fail($"Unexpected state: {state}");
return;
}
}
}
/// <summary>
/// Sets the transport to a disconnected state.
/// </summary>
/// <param name="error">Optional error information associated with the transport disconnecting. Should be <see langwor="null"/> if the disconnect was graceful and expected.</param>
protected void SetDisconnected(Exception? error = null)
{
int state = _state;
switch (state)
{ {
_messageChannel.Writer.Complete(); case StateInitial:
case StateConnected:
_state = StateDisconnected;
_messageChannel.Writer.TryComplete(error);
break;
case StateDisconnected:
return;
default:
Debug.Fail($"Unexpected state: {state}");
break;
} }
} }

View file

@ -125,7 +125,18 @@ internal abstract partial class McpEndpoint : IAsyncDisposable
} }
protected McpSession GetSessionOrThrow() protected McpSession GetSessionOrThrow()
=> _session ?? throw new InvalidOperationException($"This should be unreachable from public API! Call {nameof(InitializeSession)} before sending messages."); {
#if NET
ObjectDisposedException.ThrowIf(_disposed, this);
#else
if (_disposed)
{
throw new ObjectDisposedException(GetType().Name);
}
#endif
return _session ?? throw new InvalidOperationException($"This should be unreachable from public API! Call {nameof(InitializeSession)} before sending messages.");
}
[LoggerMessage(Level = LogLevel.Information, Message = "{EndpointName} shutting down.")] [LoggerMessage(Level = LogLevel.Information, Message = "{EndpointName} shutting down.")]
private partial void LogEndpointShuttingDown(string endpointName); private partial void LogEndpointShuttingDown(string endpointName);

View file

@ -197,7 +197,7 @@ internal sealed partial class McpSession : IDisposable
// Fail any pending requests, as they'll never be satisfied. // Fail any pending requests, as they'll never be satisfied.
foreach (var entry in _pendingRequests) foreach (var entry in _pendingRequests)
{ {
entry.Value.TrySetException(new InvalidOperationException("The server shut down unexpectedly.")); entry.Value.TrySetException(new IOException("The server shut down unexpectedly."));
} }
} }
} }

View file

@ -0,0 +1,21 @@
using ModelContextProtocol.Client;
using ModelContextProtocol.Protocol.Transport;
using System.Runtime.InteropServices;
namespace ModelContextProtocol.Tests.Transport;
public class StdioClientTransportTests
{
[Fact]
public async Task CreateAsync_ValidProcessInvalidServer_Throws()
{
string id = Guid.NewGuid().ToString("N");
StdioClientTransport transport = RuntimeInformation.IsOSPlatform(OSPlatform.Windows) ?
new(new() { Command = "cmd", Arguments = ["/C", $"echo \"{id}\" >&2"] }) :
new(new() { Command = "ls", Arguments = [id] });
IOException e = await Assert.ThrowsAsync<IOException>(() => McpClientFactory.CreateAsync(transport, cancellationToken: TestContext.Current.CancellationToken));
Assert.Contains(id, e.ToString());
}
}