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,6 +1,5 @@
using Microsoft.Extensions.AI;
using Microsoft.Extensions.DependencyInjection;
using Microsoft.Extensions.Logging;
using ModelContextProtocol.Client;
using ModelContextProtocol.Protocol.Messages;
using ModelContextProtocol.Protocol.Transport;
@ -14,16 +13,14 @@ namespace ModelContextProtocol.Tests.Server;
public class McpServerTests : LoggedTest
{
private readonly Mock<IServerTransport> _serverTransport;
private readonly Mock<ILogger> _logger;
private readonly Mock<ITransport> _serverTransport;
private readonly McpServerOptions _options;
private readonly IServiceProvider _serviceProvider;
public McpServerTests(ITestOutputHelper testOutputHelper)
: base(testOutputHelper)
{
_serverTransport = new Mock<IServerTransport>();
_logger = new Mock<ILogger>();
_serverTransport = new Mock<ITransport>();
_options = CreateOptions();
_serviceProvider = new ServiceCollection().BuildServiceProvider();
}
@ -84,61 +81,23 @@ public class McpServerTests : LoggedTest
}
[Fact]
public async Task StartAsync_Should_Throw_InvalidOperationException_If_Already_Initializing()
public async Task RunAsync_Should_Throw_InvalidOperationException_If_Already_Running()
{
// Arrange
await using var server = McpServerFactory.Create(_serverTransport.Object, _options, LoggerFactory, _serviceProvider);
var task = server.StartAsync(TestContext.Current.CancellationToken);
var runTask = server.RunAsync(TestContext.Current.CancellationToken);
// Act & Assert
await Assert.ThrowsAsync<InvalidOperationException>(() => server.StartAsync(TestContext.Current.CancellationToken));
await Assert.ThrowsAsync<InvalidOperationException>(() => server.RunAsync(TestContext.Current.CancellationToken));
await task;
}
[Fact]
public async Task StartAsync_Should_Do_Nothing_If_Already_Initialized()
{
// Arrange
await using var server = McpServerFactory.Create(_serverTransport.Object, _options, LoggerFactory, _serviceProvider);
SetInitialized(server, true);
await server.StartAsync(TestContext.Current.CancellationToken);
// Assert
_serverTransport.Verify(t => t.StartListeningAsync(It.IsAny<CancellationToken>()), Times.Never);
}
[Fact]
public async Task StartAsync_ShouldStartListening()
{
// Arrange
await using var server = McpServerFactory.Create(_serverTransport.Object, _options, LoggerFactory, _serviceProvider);
// Act
await server.StartAsync(TestContext.Current.CancellationToken);
// Assert
_serverTransport.Verify(t => t.StartListeningAsync(It.IsAny<CancellationToken>()), Times.Once);
}
[Fact]
public async Task StartAsync_Sets_Initialized_After_Transport_Responses_Initialized_Notification()
{
await using var transport = new TestServerTransport();
await using var server = McpServerFactory.Create(transport, _options, LoggerFactory, _serviceProvider);
await server.StartAsync(TestContext.Current.CancellationToken);
// Send initialized notification
await transport.SendMessageAsync(new JsonRpcNotification
{
Method = "notifications/initialized"
}, TestContext.Current.CancellationToken);
await Task.Delay(50, TestContext.Current.CancellationToken);
Assert.True(server.IsInitialized);
try
{
await runTask;
}
catch (NullReferenceException)
{
// _serverTransport.Object returns a null MessageReader
}
}
[Fact]
@ -162,7 +121,7 @@ public class McpServerTests : LoggedTest
await using var server = McpServerFactory.Create(transport, _options, LoggerFactory, _serviceProvider);
SetClientCapabilities(server, new ClientCapabilities { Sampling = new SamplingCapability() });
await server.StartAsync(TestContext.Current.CancellationToken);
var runTask = server.RunAsync(TestContext.Current.CancellationToken);
// Act
var result = await server.RequestSamplingAsync(new CreateMessageRequestParams { Messages = [] }, CancellationToken.None);
@ -171,6 +130,9 @@ public class McpServerTests : LoggedTest
Assert.NotEmpty(transport.SentMessages);
Assert.IsType<JsonRpcRequest>(transport.SentMessages[0]);
Assert.Equal("sampling/createMessage", ((JsonRpcRequest)transport.SentMessages[0]).Method);
await transport.DisposeAsync();
await runTask;
}
[Fact]
@ -191,7 +153,7 @@ public class McpServerTests : LoggedTest
await using var transport = new TestServerTransport();
await using var server = McpServerFactory.Create(transport, _options, LoggerFactory, _serviceProvider);
SetClientCapabilities(server, new ClientCapabilities { Roots = new RootsCapability() });
await server.StartAsync(TestContext.Current.CancellationToken);
var runTask = server.RunAsync(TestContext.Current.CancellationToken);
// Act
var result = await server.RequestRootsAsync(new ListRootsRequestParams(), CancellationToken.None);
@ -201,6 +163,9 @@ public class McpServerTests : LoggedTest
Assert.NotEmpty(transport.SentMessages);
Assert.IsType<JsonRpcRequest>(transport.SentMessages[0]);
Assert.Equal("roots/list", ((JsonRpcRequest)transport.SentMessages[0]).Method);
await transport.DisposeAsync();
await runTask;
}
[Fact]
@ -555,7 +520,7 @@ public class McpServerTests : LoggedTest
await using var server = McpServerFactory.Create(transport, options, LoggerFactory, _serviceProvider);
await server.StartAsync();
var runTask = server.RunAsync(TestContext.Current.CancellationToken);
var receivedMessage = new TaskCompletionSource<JsonRpcResponse>();
@ -573,11 +538,14 @@ public class McpServerTests : LoggedTest
}
);
var response = await receivedMessage.Task.WaitAsync(TimeSpan.FromSeconds(1));
var response = await receivedMessage.Task.WaitAsync(TimeSpan.FromSeconds(5));
Assert.NotNull(response);
Assert.NotNull(response.Result);
assertResult(response.Result);
await transport.DisposeAsync();
await runTask;
}
private async Task Throws_Exception_If_No_Handler_Assigned(ServerCapabilities serverCapabilities, string method, string expectedError)
@ -633,13 +601,6 @@ public class McpServerTests : LoggedTest
property.SetValue(server, capabilities);
}
private static void SetInitialized(IMcpServer server, bool isInitialized)
{
PropertyInfo? property = server.GetType().GetProperty("IsInitialized", BindingFlags.Public | BindingFlags.Instance);
Assert.NotNull(property);
property.SetValue(server, isInitialized);
}
private sealed class TestServerForIChatClient(bool supportsSampling) : IMcpServer
{
public ClientCapabilities? ClientCapabilities =>
@ -675,8 +636,6 @@ public class McpServerTests : LoggedTest
public ValueTask DisposeAsync() => default;
public bool IsInitialized => throw new NotImplementedException();
public Implementation? ClientInfo => throw new NotImplementedException();
public McpServerOptions ServerOptions => throw new NotImplementedException();
public IServiceProvider? Services => throw new NotImplementedException();
@ -684,9 +643,7 @@ public class McpServerTests : LoggedTest
throw new NotImplementedException();
public Task SendMessageAsync(IJsonRpcMessage message, CancellationToken cancellationToken = default) =>
throw new NotImplementedException();
public Task StartAsync(CancellationToken cancellationToken = default) =>
public Task RunAsync(CancellationToken cancellationToken = default) =>
throw new NotImplementedException();
public object? GetService(Type serviceType, object? serviceKey = null) => null;
}
}