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:
parent
d47c834fa9
commit
2ab9db0878
1 changed files with 2015 additions and 1847 deletions
|
|
@ -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");
|
||||
|
|
|
|||
Loading…
Reference in a new issue