Skip to content

Commit cedcddc

Browse files
committed
Add client completion notification and details
Introduce a new mechanism for reporting MCP client session completion details, including both graceful and error terminations. Add a Completion property to McpClient, exposing a `Task<ClientCompletionDetails>` that completes with details about session closure. Implement extensible completion details (including StdioClientCompletionDetails for stdio transports) and propagate them via a new internal TransportClosedException.
1 parent ce2a5e4 commit cedcddc

17 files changed

+366
-14
lines changed
Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,21 @@
1+
namespace ModelContextProtocol.Client;
2+
3+
/// <summary>
4+
/// Provides details about why an MCP client session completed.
5+
/// </summary>
6+
/// <remarks>
7+
/// <para>
8+
/// Transport implementations may return derived types with additional strongly-typed
9+
/// information, such as <see cref="StdioClientCompletionDetails"/>.
10+
/// </para>
11+
/// </remarks>
12+
public class ClientCompletionDetails
13+
{
14+
/// <summary>
15+
/// Gets the exception that caused the session to close, if any.
16+
/// </summary>
17+
/// <remarks>
18+
/// This is <see langword="null"/> for graceful closure.
19+
/// </remarks>
20+
public Exception? Exception { get; set; }
21+
}

src/ModelContextProtocol.Core/Client/McpClient.cs

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -44,4 +44,21 @@ public abstract partial class McpClient : McpSession
4444
/// </para>
4545
/// </remarks>
4646
public abstract string? ServerInstructions { get; }
47+
48+
/// <summary>
49+
/// Gets a <see cref="Task{TResult}"/> that completes when the client session has completed.
50+
/// </summary>
51+
/// <remarks>
52+
/// <para>
53+
/// The task always completes successfully. The result provides details about why the session
54+
/// completed. Transport implementations may return derived types with additional strongly-typed
55+
/// information, such as <see cref="StdioClientCompletionDetails"/>.
56+
/// </para>
57+
/// <para>
58+
/// For graceful closure (e.g., explicit disposal), <see cref="ClientCompletionDetails.Exception"/>
59+
/// will be <see langword="null"/>. For unexpected closure (e.g., process crash, network failure),
60+
/// it may contain an exception that caused or that represents the failure.
61+
/// </para>
62+
/// </remarks>
63+
public abstract Task<ClientCompletionDetails> Completion { get; }
4764
}

src/ModelContextProtocol.Core/Client/McpClientImpl.cs

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -519,6 +519,9 @@ private void RegisterTaskHandlers(RequestHandlers requestHandlers, IMcpTaskStore
519519
/// <inheritdoc/>
520520
public override string? ServerInstructions => _serverInstructions;
521521

522+
/// <inheritdoc/>
523+
public override Task<ClientCompletionDetails> Completion => _sessionHandler.CompletionTask;
524+
522525
/// <summary>
523526
/// Asynchronously connects to an MCP server, establishes the transport connection, and completes the initialization handshake.
524527
/// </summary>
@@ -653,6 +656,14 @@ public override async ValueTask DisposeAsync()
653656
_taskCancellationTokenProvider?.Dispose();
654657
await _sessionHandler.DisposeAsync().ConfigureAwait(false);
655658
await _transport.DisposeAsync().ConfigureAwait(false);
659+
660+
// After disposal, the channel writer is complete but ProcessMessagesCoreAsync
661+
// may have been cancelled with unread items still buffered. ChannelReader.Completion
662+
// only resolves once all items are consumed, so drain remaining items.
663+
while (_transport.MessageReader.TryRead(out var _));
664+
665+
// Then ensure all work has quiesced.
666+
await Completion.ConfigureAwait(false);
656667
}
657668

658669
[LoggerMessage(Level = LogLevel.Information, Message = "{EndpointName} client received server '{ServerInfo}' capabilities: '{Capabilities}'.")]

src/ModelContextProtocol.Core/Client/SseClientSessionTransport.cs

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -179,6 +179,8 @@ private async Task ReceiveMessagesAsync(CancellationToken cancellationToken)
179179
}
180180
else
181181
{
182+
SetDisconnected(ex);
183+
182184
LogTransportReadMessagesFailed(Name, ex);
183185
_connectionEstablished.TrySetException(ex);
184186
throw;
Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,22 @@
1+
namespace ModelContextProtocol.Client;
2+
3+
/// <summary>
4+
/// Provides details about the completion of a stdio-based MCP client session.
5+
/// </summary>
6+
public sealed class StdioClientCompletionDetails : ClientCompletionDetails
7+
{
8+
/// <summary>
9+
/// Gets the process ID of the server process, or <see langword="null"/> if unavailable.
10+
/// </summary>
11+
public int? ProcessId { get; set; }
12+
13+
/// <summary>
14+
/// Gets the exit code of the server process, or <see langword="null"/> if unavailable.
15+
/// </summary>
16+
public int? ExitCode { get; set; }
17+
18+
/// <summary>
19+
/// Gets the last lines of the server process's standard error output, or <see langword="null"/> if unavailable.
20+
/// </summary>
21+
public IEnumerable<string>? StandardErrorTail { get; set; }
22+
}

src/ModelContextProtocol.Core/Client/StdioClientSessionTransport.cs

Lines changed: 54 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -5,14 +5,22 @@
55
namespace ModelContextProtocol.Client;
66

77
/// <summary>Provides the client side of a stdio-based session transport.</summary>
8-
internal sealed class StdioClientSessionTransport(
9-
StdioClientTransportOptions options, Process process, string endpointName, Queue<string> stderrRollingLog, ILoggerFactory? loggerFactory) :
10-
StreamClientSessionTransport(process.StandardInput.BaseStream, process.StandardOutput.BaseStream, encoding: null, endpointName, loggerFactory)
8+
internal sealed class StdioClientSessionTransport : StreamClientSessionTransport
119
{
12-
private readonly StdioClientTransportOptions _options = options;
13-
private readonly Process _process = process;
14-
private readonly Queue<string> _stderrRollingLog = stderrRollingLog;
10+
private readonly StdioClientTransportOptions _options;
11+
private readonly Process _process;
12+
private readonly Queue<string> _stderrRollingLog;
1513
private int _cleanedUp = 0;
14+
private readonly int? _processId;
15+
16+
public StdioClientSessionTransport(StdioClientTransportOptions options, Process process, string endpointName, Queue<string> stderrRollingLog, ILoggerFactory? loggerFactory) :
17+
base(process.StandardInput.BaseStream, process.StandardOutput.BaseStream, encoding: null, endpointName, loggerFactory)
18+
{
19+
_options = options;
20+
_process = process;
21+
_stderrRollingLog = stderrRollingLog;
22+
try { _processId = process.Id; } catch { }
23+
}
1624

1725
/// <inheritdoc/>
1826
public override async Task SendMessageAsync(JsonRpcMessage message, CancellationToken cancellationToken = default)
@@ -47,17 +55,26 @@ protected override async ValueTask CleanupAsync(Exception? error = null, Cancell
4755
// so create an exception with details about that.
4856
error ??= await GetUnexpectedExitExceptionAsync(cancellationToken).ConfigureAwait(false);
4957

50-
// Now terminate the server process.
58+
// Terminate the server process (or confirm it already exited), then build
59+
// and publish strongly-typed completion details while the process handle
60+
// is still valid so we can read the exit code.
5161
try
5262
{
53-
StdioClientTransport.DisposeProcess(_process, processRunning: true, shutdownTimeout: _options.ShutdownTimeout);
63+
StdioClientTransport.DisposeProcess(
64+
_process,
65+
processRunning: true,
66+
_options.ShutdownTimeout,
67+
beforeDispose: () => SetDisconnected(new TransportClosedException(BuildCompletionDetails(error))));
5468
}
5569
catch (Exception ex)
5670
{
5771
LogTransportShutdownFailed(Name, ex);
72+
SetDisconnected(new TransportClosedException(BuildCompletionDetails(error)));
5873
}
5974

60-
// And handle cleanup in the base type.
75+
// And handle cleanup in the base type. SetDisconnected has already been
76+
// called above, so the base call is a no-op for disconnect state but
77+
// still performs other cleanup (cancelling the read task, etc.).
6178
await base.CleanupAsync(error, cancellationToken).ConfigureAwait(false);
6279
}
6380

@@ -104,4 +121,32 @@ protected override async ValueTask CleanupAsync(Exception? error = null, Cancell
104121

105122
return new IOException(errorMessage);
106123
}
124+
125+
private StdioClientCompletionDetails BuildCompletionDetails(Exception? error)
126+
{
127+
StdioClientCompletionDetails details = new()
128+
{
129+
Exception = error,
130+
ProcessId = _processId,
131+
};
132+
133+
try
134+
{
135+
if (StdioClientTransport.HasExited(_process))
136+
{
137+
details.ExitCode = _process.ExitCode;
138+
}
139+
}
140+
catch { }
141+
142+
lock (_stderrRollingLog)
143+
{
144+
if (_stderrRollingLog.Count > 0)
145+
{
146+
details.StandardErrorTail = _stderrRollingLog.ToArray();
147+
}
148+
}
149+
150+
return details;
151+
}
107152
}

src/ModelContextProtocol.Core/Client/StdioClientTransport.cs

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -213,7 +213,7 @@ public async Task<ITransport> ConnectAsync(CancellationToken cancellationToken =
213213
}
214214

215215
internal static void DisposeProcess(
216-
Process? process, bool processRunning, TimeSpan shutdownTimeout)
216+
Process? process, bool processRunning, TimeSpan shutdownTimeout, Action? beforeDispose = null)
217217
{
218218
if (process is not null)
219219
{
@@ -239,6 +239,10 @@ internal static void DisposeProcess(
239239
{
240240
process.WaitForExit();
241241
}
242+
243+
// Invoke the callback while the process handle is still valid,
244+
// e.g. to read ExitCode before Dispose() invalidates it.
245+
beforeDispose?.Invoke();
242246
}
243247
finally
244248
{
Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,19 @@
1+
using ModelContextProtocol.Protocol;
2+
using System.Threading.Channels;
3+
4+
namespace ModelContextProtocol.Client;
5+
6+
/// <summary>
7+
/// <see cref="IOException"/> used to smuggle <see cref="ClientCompletionDetails"/> through
8+
/// the <see cref="ChannelWriter{T}.TryComplete(Exception?)"/> mechanism.
9+
/// </summary>
10+
/// <remarks>
11+
/// This could be made public in the future to allow custom <see cref="ITransport"/>
12+
/// implementations to provide their own <see cref="ClientCompletionDetails"/>-derived types
13+
/// by completing their channel with this exception.
14+
/// </remarks>
15+
internal sealed class TransportClosedException(ClientCompletionDetails details) :
16+
IOException(details.Exception?.Message, details.Exception)
17+
{
18+
public ClientCompletionDetails Details { get; } = details;
19+
}

src/ModelContextProtocol.Core/McpSessionHandler.cs

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -132,6 +132,7 @@ public McpSessionHandler(
132132
_incomingMessageFilter = incomingMessageFilter ?? (next => next);
133133
_outgoingMessageFilter = outgoingMessageFilter ?? (next => next);
134134
_logger = logger;
135+
135136
LogSessionCreated(EndpointName, _sessionId, _transportKind);
136137
}
137138

@@ -145,6 +146,15 @@ public McpSessionHandler(
145146
/// </summary>
146147
public string? NegotiatedProtocolVersion { get; set; }
147148

149+
/// <summary>
150+
/// Gets a task that completes when the client session has completed, providing details about the closure.
151+
/// Completion details are resolved from the transport's channel completion exception: if a transport
152+
/// completes its channel with a <see cref="TransportClosedException"/>, the wrapped
153+
/// <see cref="ClientCompletionDetails"/> is unwrapped. Otherwise, a default instance is returned.
154+
/// </summary>
155+
internal Task<ClientCompletionDetails> CompletionTask =>
156+
field ??= GetCompletionDetailsAsync(_transport.MessageReader.Completion);
157+
148158
/// <summary>
149159
/// Starts processing messages from the transport. This method will block until the transport is disconnected.
150160
/// This is generally started in a background task or thread from the initialization logic of the derived class.
@@ -297,6 +307,28 @@ ex is OperationCanceledException &&
297307
}
298308
}
299309

310+
/// <summary>
311+
/// Resolves <see cref="ClientCompletionDetails"/> from the transport's channel completion.
312+
/// If the channel was completed with a <see cref="TransportClosedException"/>, the wrapped
313+
/// details are returned. Otherwise a default instance is created from the completion state.
314+
/// </summary>
315+
private static async Task<ClientCompletionDetails> GetCompletionDetailsAsync(Task channelCompletion)
316+
{
317+
try
318+
{
319+
await channelCompletion.ConfigureAwait(false);
320+
return new ClientCompletionDetails();
321+
}
322+
catch (TransportClosedException tce)
323+
{
324+
return tce.Details;
325+
}
326+
catch (Exception ex)
327+
{
328+
return new ClientCompletionDetails { Exception = ex };
329+
}
330+
}
331+
300332
private async Task HandleMessageAsync(JsonRpcMessage message, CancellationToken cancellationToken)
301333
{
302334
Histogram<double> durationMetric = _isServer ? s_serverOperationDuration : s_clientOperationDuration;

tests/ModelContextProtocol.AspNetCore.Tests/HttpServerIntegrationTests.cs

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -303,4 +303,17 @@ public async Task CallTool_Sse_EchoServer_Concurrently()
303303
Assert.Equal($"Echo: Hello MCP! {i}", textContent.Text);
304304
}
305305
}
306+
307+
[Fact]
308+
public async Task Completion_GracefulDisposal_ReturnsCompletionDetails()
309+
{
310+
var client = await GetClientAsync();
311+
Assert.False(client.Completion.IsCompleted);
312+
313+
await client.DisposeAsync();
314+
Assert.True(client.Completion.IsCompleted);
315+
316+
var details = await client.Completion;
317+
Assert.Null(details.Exception);
318+
}
306319
}

0 commit comments

Comments
 (0)