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,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;
|
||||
}
|
||||
}
|
||||
|
|
|
|||
Loading…
Reference in a new issue