Split StdioClient/ServerTransports into Stdio : Stream (#172)

This commit is contained in:
Stephen Toub 2025-03-31 18:44:54 -04:00 • committed by GitHub
commit 96fa96c1b2
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 667 additions and 687 deletions

View file

@ -23,7 +23,7 @@ public class McpClientExtensionsTests : LoggedTest
sc.AddSingleton(LoggerFactory);
sc.AddMcpServer().WithStdioServerTransport();
// Call WithStdioServerTransport to get the IMcpServer registration, then overwrite default transport with a pipe transport.
sc.AddSingleton<ITransport>(new StdioServerTransport("TestServer", _clientToServerPipe.Reader.AsStream(), _serverToClientPipe.Writer.AsStream()));
sc.AddSingleton<ITransport>(new StreamServerTransport(_clientToServerPipe.Reader.AsStream(), _serverToClientPipe.Writer.AsStream()));
for (int f = 0; f < 10; f++)
{
string name = $"Method{f}";
@ -53,19 +53,17 @@ public class McpClientExtensionsTests : LoggedTest
private async Task<IMcpClient> CreateMcpClientForServer()
{
var serverStdinWriter = new StreamWriter(_clientToServerPipe.Writer.AsStream());
var serverStdoutReader = new StreamReader(_serverToClientPipe.Reader.AsStream());
var serverConfig = new McpServerConfig()
{
Id = "TestServer",
Name = "TestServer",
TransportType = "ignored",
};
return await McpClientFactory.CreateAsync(
serverConfig,
createTransportFunc: (_, _) => new StreamClientTransport(serverStdinWriter, serverStdoutReader, LoggerFactory),
new McpServerConfig()
{
Id = "TestServer",
Name = "TestServer",
TransportType = "ignored",
},
createTransportFunc: (_, _) => new StreamClientTransport(
serverInput: _clientToServerPipe.Writer.AsStream(),
serverOutput: _serverToClientPipe.Reader.AsStream(),
LoggerFactory),
loggerFactory: LoggerFactory,
cancellationToken: TestContext.Current.CancellationToken);
}

View file

@ -30,7 +30,7 @@ public class McpServerBuilderExtensionsPromptsTests : LoggedTest, IAsyncDisposab
sc.AddSingleton(LoggerFactory);
_builder = sc.AddMcpServer().WithStdioServerTransport().WithPrompts<SimplePrompts>();
// Call WithStdioServerTransport to get the IMcpServer registration, then overwrite default transport with a pipe transport.
sc.AddSingleton<ITransport>(new StdioServerTransport("TestServer", _clientToServerPipe.Reader.AsStream(), _serverToClientPipe.Writer.AsStream(), LoggerFactory));
sc.AddSingleton<ITransport>(new StreamServerTransport(_clientToServerPipe.Reader.AsStream(), _serverToClientPipe.Writer.AsStream(), loggerFactory: LoggerFactory));
sc.AddSingleton(new ObjectWithId());
_serviceProvider = sc.BuildServiceProvider();
@ -55,19 +55,17 @@ public class McpServerBuilderExtensionsPromptsTests : LoggedTest, IAsyncDisposab
private async Task<IMcpClient> CreateMcpClientForServer()
{
var serverStdinWriter = new StreamWriter(_clientToServerPipe.Writer.AsStream());
var serverStdoutReader = new StreamReader(_serverToClientPipe.Reader.AsStream());
var serverConfig = new McpServerConfig()
{
Id = "TestServer",
Name = "TestServer",
TransportType = "ignored",
};
return await McpClientFactory.CreateAsync(
serverConfig,
createTransportFunc: (_, _) => new StreamClientTransport(serverStdinWriter, serverStdoutReader, LoggerFactory),
new McpServerConfig()
{
Id = "TestServer",
Name = "TestServer",
TransportType = "ignored",
},
createTransportFunc: (_, _) => new StreamClientTransport(
serverInput: _clientToServerPipe.Writer.AsStream(),
serverOutput: _serverToClientPipe.Reader.AsStream(),
LoggerFactory),
loggerFactory: LoggerFactory,
cancellationToken: TestContext.Current.CancellationToken);
}

View file

@ -34,7 +34,7 @@ public class McpServerBuilderExtensionsToolsTests : LoggedTest, IAsyncDisposable
sc.AddSingleton(LoggerFactory);
_builder = sc.AddMcpServer().WithStdioServerTransport().WithTools<EchoTool>();
// Call WithStdioServerTransport to get the IMcpServer registration, then overwrite default transport with a pipe transport.
sc.AddSingleton<ITransport>(new StdioServerTransport("TestServer", _clientToServerPipe.Reader.AsStream(), _serverToClientPipe.Writer.AsStream(), LoggerFactory));
sc.AddSingleton<ITransport>(new StreamServerTransport(_clientToServerPipe.Reader.AsStream(), _serverToClientPipe.Writer.AsStream(), loggerFactory: LoggerFactory));
sc.AddSingleton(new ObjectWithId());
_serviceProvider = sc.BuildServiceProvider();
@ -59,19 +59,17 @@ public class McpServerBuilderExtensionsToolsTests : LoggedTest, IAsyncDisposable
private async Task<IMcpClient> CreateMcpClientForServer()
{
var serverStdinWriter = new StreamWriter(_clientToServerPipe.Writer.AsStream());
var serverStdoutReader = new StreamReader(_serverToClientPipe.Reader.AsStream());
var serverConfig = new McpServerConfig()
{
Id = "TestServer",
Name = "TestServer",
TransportType = "ignored",
};
return await McpClientFactory.CreateAsync(
serverConfig,
createTransportFunc: (_, _) => new StreamClientTransport(serverStdinWriter, serverStdoutReader, LoggerFactory),
new McpServerConfig()
{
Id = "TestServer",
Name = "TestServer",
TransportType = "ignored",
},
createTransportFunc: (_, _) => new StreamClientTransport(
serverInput: _clientToServerPipe.Writer.AsStream(),
_serverToClientPipe.Reader.AsStream(),
LoggerFactory),
loggerFactory: LoggerFactory,
cancellationToken: TestContext.Current.CancellationToken);
}
@ -117,23 +115,21 @@ public class McpServerBuilderExtensionsToolsTests : LoggedTest, IAsyncDisposable
var stdinPipe = new Pipe();
var stdoutPipe = new Pipe();
await using var transport = new StdioServerTransport($"TestServer_{i}", stdinPipe.Reader.AsStream(), stdoutPipe.Writer.AsStream());
await using var transport = new StreamServerTransport(stdinPipe.Reader.AsStream(), stdoutPipe.Writer.AsStream());
await using var server = McpServerFactory.Create(transport, options, loggerFactory, _serviceProvider);
var serverRunTask = server.RunAsync(TestContext.Current.CancellationToken);
using var serverStdinWriter = new StreamWriter(stdinPipe.Writer.AsStream());
using var serverStdoutReader = new StreamReader(stdoutPipe.Reader.AsStream());
var serverConfig = new McpServerConfig()
{
Id = $"TestServer_{i}",
Name = $"TestServer_{i}",
TransportType = "ignored",
};
await using (var client = await McpClientFactory.CreateAsync(
serverConfig,
createTransportFunc: (_, _) => new StreamClientTransport(serverStdinWriter, serverStdoutReader, LoggerFactory),
new McpServerConfig()
{
Id = $"TestServer_{i}",
Name = $"TestServer_{i}",
TransportType = "ignored",
},
createTransportFunc: (_, _) => new StreamClientTransport(
serverInput: stdinPipe.Writer.AsStream(),
serverOutput: stdoutPipe.Reader.AsStream(),
LoggerFactory),
loggerFactory: LoggerFactory,
cancellationToken: TestContext.Current.CancellationToken))
{

View file

@ -55,7 +55,7 @@ public class StdioServerTransportTests : LoggedTest
[Fact]
public async Task Should_Start_In_Connected_State()
{
await using var transport = new StdioServerTransport(_serverOptions.ServerInfo.Name, new Pipe().Reader.AsStream(), Stream.Null, LoggerFactory);
await using var transport = new StreamServerTransport(new Pipe().Reader.AsStream(), Stream.Null, loggerFactory: LoggerFactory);
Assert.True(transport.IsConnected);
}
@ -65,11 +65,10 @@ public class StdioServerTransportTests : LoggedTest
{
using var output = new MemoryStream();
await using var transport = new StdioServerTransport(
_serverOptions.ServerInfo.Name,
await using var transport = new StreamServerTransport(
new Pipe().Reader.AsStream(),
output,
LoggerFactory);
loggerFactory: LoggerFactory);
// Verify transport is connected
Assert.True(transport.IsConnected, "Transport should be connected after StartListeningAsync");
@ -87,7 +86,7 @@ public class StdioServerTransportTests : LoggedTest
[Fact]
public async Task DisposeAsync_Should_Dispose_Resources()
{
await using var transport = new StdioServerTransport(_serverOptions.ServerInfo.Name, Stream.Null, Stream.Null, LoggerFactory);
await using var transport = new StreamServerTransport(Stream.Null, Stream.Null, loggerFactory: LoggerFactory);
await transport.DisposeAsync();
@ -104,11 +103,10 @@ public class StdioServerTransportTests : LoggedTest
Pipe pipe = new();
using var input = pipe.Reader.AsStream();
await using var transport = new StdioServerTransport(
_serverOptions.ServerInfo.Name,
await using var transport = new StreamServerTransport(
input,
Stream.Null,
LoggerFactory);
loggerFactory: LoggerFactory);
// Verify transport is connected
Assert.True(transport.IsConnected, "Transport should be connected after StartListeningAsync");
@ -128,7 +126,7 @@ public class StdioServerTransportTests : LoggedTest
[Fact]
public async Task CleanupAsync_Should_Cleanup_Resources()
{
var transport = new StdioServerTransport(_serverOptions.ServerInfo.Name, Stream.Null, Stream.Null, LoggerFactory);
var transport = new StreamServerTransport(Stream.Null, Stream.Null, loggerFactory: LoggerFactory);
await transport.DisposeAsync();
@ -141,11 +139,10 @@ public class StdioServerTransportTests : LoggedTest
// Use a reader that won't terminate
using var output = new MemoryStream();
await using var transport = new StdioServerTransport(
_serverOptions.ServerInfo.Name,
await using var transport = new StreamServerTransport(
new Pipe().Reader.AsStream(),
output,
LoggerFactory);
loggerFactory: LoggerFactory);
// Verify transport is connected
Assert.True(transport.IsConnected, "Transport should be connected after StartListeningAsync");

View file

@ -1,79 +0,0 @@
using Microsoft.Extensions.Logging;
using ModelContextProtocol.Protocol.Messages;
using ModelContextProtocol.Protocol.Transport;
using ModelContextProtocol.Utils.Json;
using System.Text.Json;
namespace ModelContextProtocol.Tests.Transport;
internal sealed class StreamClientTransport : TransportBase, IClientTransport
{
private readonly JsonSerializerOptions _jsonOptions = McpJsonUtilities.DefaultOptions;
private readonly Task? _readTask;
private readonly CancellationTokenSource _shutdownCts = new CancellationTokenSource();
private readonly TextReader _serverStdoutReader;
private readonly TextWriter _serverStdinWriter;
public StreamClientTransport(TextWriter serverStdinWriter, TextReader serverStdoutReader, ILoggerFactory loggerFactory)
: base(loggerFactory)
{
_serverStdoutReader = serverStdoutReader;
_serverStdinWriter = serverStdinWriter;
_readTask = Task.Run(() => ReadMessagesAsync(_shutdownCts.Token), CancellationToken.None);
SetConnected(true);
}
public Task<ITransport> ConnectAsync(CancellationToken cancellationToken = default) => Task.FromResult<ITransport>(this);
public override async Task SendMessageAsync(IJsonRpcMessage message, CancellationToken cancellationToken = default)
{
string id = message is IJsonRpcMessageWithId messageWithId ?
messageWithId.Id.ToString() :
"(no id)";
await _serverStdinWriter.WriteLineAsync(JsonSerializer.Serialize(message)).ConfigureAwait(false);
await _serverStdinWriter.FlushAsync(cancellationToken).ConfigureAwait(false);
}
private async Task ReadMessagesAsync(CancellationToken cancellationToken)
{
try
{
while (await _serverStdoutReader.ReadLineAsync(cancellationToken).ConfigureAwait(false) is string line)
{
if (!string.IsNullOrWhiteSpace(line))
{
try
{
if (JsonSerializer.Deserialize<IJsonRpcMessage>(line.Trim(), _jsonOptions) is { } message)
{
await WriteMessageAsync(message, cancellationToken).ConfigureAwait(false);
}
}
catch (JsonException)
{
}
}
}
}
catch (OperationCanceledException)
{
}
}
public override async ValueTask DisposeAsync()
{
if (_shutdownCts is { } shutdownCts)
{
await shutdownCts.CancelAsync().ConfigureAwait(false);
shutdownCts.Dispose();
}
if (_readTask is Task readTask)
{
await readTask.WaitAsync(TimeSpan.FromSeconds(5)).ConfigureAwait(false);
}
SetConnected(false);
}
}