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,19 +1,19 @@
using System.ComponentModel;
using System.Text.Json;
using ModelContextProtocol.Server;
using Microsoft.Extensions.AI;
using Microsoft.Extensions.DependencyInjection;
using ModelContextProtocol.Protocol.Transport;
using System.IO.Pipelines;
using Microsoft.Extensions.Logging;
using Microsoft.Extensions.Options;
using ModelContextProtocol.Client;
using ModelContextProtocol.Configuration;
using ModelContextProtocol.Tests.Transport;
using System.Text.RegularExpressions;
using Microsoft.Extensions.AI;
using System.Threading.Channels;
using ModelContextProtocol.Protocol.Messages;
using Microsoft.Extensions.Options;
using ModelContextProtocol.Protocol.Transport;
using ModelContextProtocol.Server;
using ModelContextProtocol.Tests.Transport;
using ModelContextProtocol.Tests.Utils;
using Microsoft.Extensions.Logging;
using System.ComponentModel;
using System.IO.Pipelines;
using System.Text.Json;
using System.Text.RegularExpressions;
using System.Threading.Channels;
namespace ModelContextProtocol.Tests.Configuration;
@ -23,31 +23,41 @@ public class McpServerBuilderExtensionsToolsTests : LoggedTest, IAsyncDisposable
private readonly Pipe _serverToClientPipe = new();
private readonly ServiceProvider _serviceProvider;
private readonly IMcpServerBuilder _builder;
private readonly IMcpServer _server;
private readonly CancellationTokenSource _cts;
private readonly Task _serverTask;
public McpServerBuilderExtensionsToolsTests(ITestOutputHelper testOutputHelper)
: base(testOutputHelper)
{
ServiceCollection sc = new();
sc.AddSingleton(LoggerFactory);
sc.AddSingleton<IServerTransport>(new StdioServerTransport("TestServer", _clientToServerPipe.Reader.AsStream(), _serverToClientPipe.Writer.AsStream()));
_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(new ObjectWithId());
_builder = sc.AddMcpServer().WithTools<EchoTool>();
_serviceProvider = sc.BuildServiceProvider();
_server = _serviceProvider.GetRequiredService<IMcpServer>();
var server = _serviceProvider.GetRequiredService<IMcpServer>();
_cts = CancellationTokenSource.CreateLinkedTokenSource(TestContext.Current.CancellationToken);
_serverTask = server.RunAsync(cancellationToken: _cts.Token);
}
public ValueTask DisposeAsync()
public async ValueTask DisposeAsync()
{
await _cts.CancelAsync();
_clientToServerPipe.Writer.Complete();
_serverToClientPipe.Writer.Complete();
return _serviceProvider.DisposeAsync();
await _serverTask;
await _serviceProvider.DisposeAsync();
_cts.Dispose();
Dispose();
}
private async Task<IMcpClient> CreateMcpClientForServer()
{
await _server.StartAsync(TestContext.Current.CancellationToken);
var serverStdinWriter = new StreamWriter(_clientToServerPipe.Writer.AsStream());
var serverStdoutReader = new StreamReader(_serverToClientPipe.Reader.AsStream());
@ -60,14 +70,16 @@ public class McpServerBuilderExtensionsToolsTests : LoggedTest, IAsyncDisposable
return await McpClientFactory.CreateAsync(
serverConfig,
createTransportFunc: (_, _) => new StreamClientTransport(serverStdinWriter, serverStdoutReader),
createTransportFunc: (_, _) => new StreamClientTransport(serverStdinWriter, serverStdoutReader, LoggerFactory),
loggerFactory: LoggerFactory,
cancellationToken: TestContext.Current.CancellationToken);
}
[Fact]
public void Adds_Tools_To_Server()
{
var tools = _server.ServerOptions?.Capabilities?.Tools?.ToolCollection;
var serverOptions = _serviceProvider.GetRequiredService<IOptions<McpServerOptions>>().Value;
var tools = serverOptions.Capabilities?.Tools?.ToolCollection;
Assert.NotNull(tools);
Assert.NotEmpty(tools);
}
@ -93,7 +105,6 @@ public class McpServerBuilderExtensionsToolsTests : LoggedTest, IAsyncDisposable
Assert.Equal("Echoes the input back to the client.", doubleEchoTool.Description);
}
[Fact]
public async Task Can_Create_Multiple_Servers_From_Options_And_List_Registered_Tools()
{
@ -105,28 +116,26 @@ public class McpServerBuilderExtensionsToolsTests : LoggedTest, IAsyncDisposable
var stdinPipe = new Pipe();
var stdoutPipe = new Pipe();
try
await using var transport = new StdioServerTransport($"TestServer_{i}", 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()
{
var transport = new StdioServerTransport($"TestServer_{i}", stdinPipe.Reader.AsStream(), stdoutPipe.Writer.AsStream());
var server = McpServerFactory.Create(transport, options, loggerFactory, _serviceProvider);
await server.StartAsync(TestContext.Current.CancellationToken);
var serverStdinWriter = new StreamWriter(stdinPipe.Writer.AsStream());
var serverStdoutReader = new StreamReader(stdoutPipe.Reader.AsStream());
var serverConfig = new McpServerConfig()
{
Id = $"TestServer_{i}",
Name = $"TestServer_{i}",
TransportType = "ignored",
};
var client = await McpClientFactory.CreateAsync(
serverConfig,
createTransportFunc: (_, _) => new StreamClientTransport(serverStdinWriter, serverStdoutReader),
cancellationToken: TestContext.Current.CancellationToken);
Id = $"TestServer_{i}",
Name = $"TestServer_{i}",
TransportType = "ignored",
};
await using (var client = await McpClientFactory.CreateAsync(
serverConfig,
createTransportFunc: (_, _) => new StreamClientTransport(serverStdinWriter, serverStdoutReader, LoggerFactory),
loggerFactory: LoggerFactory,
cancellationToken: TestContext.Current.CancellationToken))
{
var tools = await client.ListToolsAsync(TestContext.Current.CancellationToken);
Assert.Equal(11, tools.Count);
@ -142,11 +151,10 @@ public class McpServerBuilderExtensionsToolsTests : LoggedTest, IAsyncDisposable
Assert.Equal("double_echo", doubleEchoTool.Name);
Assert.Equal("Echoes the input back to the client.", doubleEchoTool.Description);
}
finally
{
stdinPipe.Writer.Complete();
stdoutPipe.Writer.Complete();
}
stdinPipe.Writer.Complete();
await serverRunTask;
stdoutPipe.Writer.Complete();
}
}
@ -168,7 +176,8 @@ public class McpServerBuilderExtensionsToolsTests : LoggedTest, IAsyncDisposable
var notificationRead = listChanged.Reader.ReadAsync(TestContext.Current.CancellationToken);
Assert.False(notificationRead.IsCompleted);
var serverTools = _server.ServerOptions.Capabilities?.Tools?.ToolCollection;
var serverOptions = _serviceProvider.GetRequiredService<IOptions<McpServerOptions>>().Value;
var serverTools = serverOptions.Capabilities?.Tools?.ToolCollection;
Assert.NotNull(serverTools);
var newTool = McpServerTool.Create([McpServerTool(Name = "NewTool")] () => "42");