Refactor transports to help enable graceful shutdown (#142)

- This also paves the way for better multi-session support
- We should definitely rethink names for the transport API
- For now, I kept the names similar as possible, so we can
  focus on the API shape

Co-authored-by: Stephen Toub <stoub@microsoft.com>
This commit is contained in:
Stephen Halter 2025-03-30 17:10:14 -07:00 • committed by GitHub
commit 2ab9db0878
No known key found for this signature in database
GPG key ID: B5690EEEBB952194

View file

@ -1,4 +1,4 @@
using Microsoft.Extensions.Logging.Abstractions;
using Microsoft.Extensions.Logging;
using ModelContextProtocol.Protocol.Messages;
using ModelContextProtocol.Protocol.Transport;
using ModelContextProtocol.Utils.Json;
@ -14,8 +14,8 @@ internal sealed class StreamClientTransport : TransportBase, IClientTransport
private readonly TextReader _serverStdoutReader;
private readonly TextWriter _serverStdinWriter;
public StreamClientTransport(TextWriter serverStdinWriter, TextReader serverStdoutReader)
: base(NullLoggerFactory.Instance)
public StreamClientTransport(TextWriter serverStdinWriter, TextReader serverStdoutReader, ILoggerFactory loggerFactory)
: base(loggerFactory)
{
_serverStdoutReader = serverStdoutReader;
_serverStdinWriter = serverStdinWriter;
@ -23,7 +23,7 @@ internal sealed class StreamClientTransport : TransportBase, IClientTransport
SetConnected(true);
}
public Task ConnectAsync(CancellationToken cancellationToken = default) => Task.CompletedTask;
public Task<ITransport> ConnectAsync(CancellationToken cancellationToken = default) => Task.FromResult<ITransport>(this);
public override async Task SendMessageAsync(IJsonRpcMessage message, CancellationToken cancellationToken = default)
{
@ -37,22 +37,28 @@ internal sealed class StreamClientTransport : TransportBase, IClientTransport
private async Task ReadMessagesAsync(CancellationToken cancellationToken)
{
while (await _serverStdoutReader.ReadLineAsync(cancellationToken).ConfigureAwait(false) is string line)
try
{
if (!string.IsNullOrWhiteSpace(line))
while (await _serverStdoutReader.ReadLineAsync(cancellationToken).ConfigureAwait(false) is string line)
{
try
if (!string.IsNullOrWhiteSpace(line))
{
if (JsonSerializer.Deserialize<IJsonRpcMessage>(line.Trim(), _jsonOptions) is { } message)
try
{
if (JsonSerializer.Deserialize<IJsonRpcMessage>(line.Trim(), _jsonOptions) is { } message)
{
await WriteMessageAsync(message, cancellationToken).ConfigureAwait(false);
}
}
catch (JsonException)
{
await WriteMessageAsync(message, cancellationToken).ConfigureAwait(false);
}
}
catch (JsonException)
{
}
}
}
catch (OperationCanceledException)
{
}
}
public override async ValueTask DisposeAsync()