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
3 changed files with 2015 additions and 1847 deletions

View file

@ -2,8 +2,8 @@
using ModelContextProtocol.Protocol.Messages;
using ModelContextProtocol.Protocol.Transport;
using ModelContextProtocol.Tests.Utils;
using System.IO.Pipelines;
using System.Net;
using System.Reflection;
namespace ModelContextProtocol.Tests.Transport;
@ -35,26 +35,6 @@ public class SseClientTransportTests : LoggedTest
};
}
[Fact]
public async Task Constructor_Should_Initialize_With_Valid_Parameters()
{
// Act
await using var transport = new SseClientTransport(_transportOptions, _serverConfig, LoggerFactory);
// Assert
Assert.NotNull(transport);
PropertyInfo? getOptions = transport.GetType().GetProperty("Options", BindingFlags.NonPublic | BindingFlags.Instance);
Assert.NotNull(getOptions);
var options = (SseClientTransportOptions)getOptions.GetValue(transport)!;
Assert.Equal(TimeSpan.FromSeconds(2), options.ConnectionTimeout);
Assert.Equal(3, options.MaxReconnectAttempts);
Assert.Equal(TimeSpan.FromMilliseconds(50), options.ReconnectDelay);
Assert.NotNull(options.AdditionalHeaders);
Assert.Equal("header", options.AdditionalHeaders["test"]);
}
[Fact]
public void Constructor_Throws_For_Null_Options()
{
@ -85,71 +65,19 @@ public class SseClientTransportTests : LoggedTest
bool firstCall = true;
mockHttpHandler.RequestHandler = async (request) =>
mockHttpHandler.RequestHandler = (request) =>
{
if (!firstCall)
{
Assert.True(transport.IsConnected);
await transport.CloseAsync();
}
firstCall = false;
return new HttpResponseMessage
return Task.FromResult(new HttpResponseMessage
{
StatusCode = HttpStatusCode.OK,
Content = new StringContent("event: endpoint\r\ndata: http://localhost\r\n\r\n")
};
});
};
await transport.ConnectAsync(TestContext.Current.CancellationToken);
}
[Fact]
public async Task ConnectAsync_Throws_If_Already_Connected()
{
using var mockHttpHandler = new MockHttpHandler();
using var httpClient = new HttpClient(mockHttpHandler);
await using var transport = new SseClientTransport(_transportOptions, _serverConfig, httpClient, LoggerFactory);
var tcsConnected = new TaskCompletionSource();
var tcsDone = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
var callIndex = 0;
mockHttpHandler.RequestHandler = async (request) =>
{
switch (callIndex++)
{
case 0:
return new HttpResponseMessage
{
StatusCode = HttpStatusCode.OK,
Content = new StringContent("event: endpoint\r\ndata: http://localhost\r\n\r\n")
};
case 1:
tcsConnected.SetResult();
await tcsDone.Task;
return new HttpResponseMessage
{
StatusCode = HttpStatusCode.OK,
Content = new StringContent("")
};
default:
return new HttpResponseMessage
{
StatusCode = HttpStatusCode.OK,
Content = new StringContent("")
};
}
};
var task = transport.ConnectAsync(TestContext.Current.CancellationToken);
await tcsConnected.Task;
Assert.True(transport.IsConnected);
var action = async () => await transport.ConnectAsync();
var exception = await Assert.ThrowsAsync<McpTransportException>(action);
Assert.Equal("Transport is already connected", exception.Message);
tcsDone.SetResult();
await transport.CloseAsync();
await task;
await using var session = await transport.ConnectAsync(TestContext.Current.CancellationToken);
Assert.NotNull(session);
Assert.False(firstCall);
}
[Fact]
@ -174,15 +102,6 @@ public class SseClientTransportTests : LoggedTest
Assert.Equal(_transportOptions.MaxReconnectAttempts, retries);
}
[Fact]
public async Task SendMessageAsync_Throws_Exception_If_MessageEndpoint_Not_Set()
{
await using var transport = new SseClientTransport(_transportOptions, _serverConfig, LoggerFactory);
// Assert
await Assert.ThrowsAsync<InvalidOperationException>(() => transport.SendMessageAsync(new JsonRpcRequest() { Method = "test" }, CancellationToken.None));
}
[Fact]
public async Task SendMessageAsync_Handles_Accepted_Response()
{
@ -216,9 +135,8 @@ public class SseClientTransportTests : LoggedTest
}
};
await transport.ConnectAsync(TestContext.Current.CancellationToken);
await transport.SendMessageAsync(new JsonRpcRequest() { Method = "initialize", Id = RequestId.FromNumber(44) }, CancellationToken.None);
await using var session = await transport.ConnectAsync(TestContext.Current.CancellationToken);
await session.SendMessageAsync(new JsonRpcRequest() { Method = "initialize", Id = RequestId.FromNumber(44) }, CancellationToken.None);
Assert.True(true);
}
@ -229,6 +147,12 @@ public class SseClientTransportTests : LoggedTest
using var httpClient = new HttpClient(mockHttpHandler);
await using var transport = new SseClientTransport(_transportOptions, _serverConfig, httpClient, LoggerFactory);
var eventSourcePipe = new Pipe();
var eventSourceData = "event: endpoint\r\ndata: /sseendpoint\r\n\r\n"u8;
Assert.True(eventSourceData.TryCopyTo(eventSourcePipe.Writer.GetSpan()));
eventSourcePipe.Writer.Advance(eventSourceData.Length);
await eventSourcePipe.Writer.FlushAsync(TestContext.Current.CancellationToken);
var firstCall = true;
mockHttpHandler.RequestHandler = (request) =>
{
@ -250,15 +174,16 @@ public class SseClientTransportTests : LoggedTest
return Task.FromResult(new HttpResponseMessage
{
StatusCode = HttpStatusCode.OK,
Content = new StringContent("event: endpoint\r\ndata: /sseendpoint\r\n\r\n")
Content = new StreamContent(eventSourcePipe.Reader.AsStream()),
});
}
};
await transport.ConnectAsync(TestContext.Current.CancellationToken);
await using var session = await transport.ConnectAsync(TestContext.Current.CancellationToken);
await transport.SendMessageAsync(new JsonRpcRequest() { Method = "initialize", Id = RequestId.FromNumber(44) }, CancellationToken.None);
await session.SendMessageAsync(new JsonRpcRequest() { Method = "initialize", Id = RequestId.FromNumber(44) }, CancellationToken.None);
Assert.True(true);
eventSourcePipe.Writer.Complete();
}
[Fact]
@ -278,45 +203,39 @@ public class SseClientTransportTests : LoggedTest
return Task.FromResult(new HttpResponseMessage
{
StatusCode = HttpStatusCode.OK,
Content = new StringContent("event: endpoint\r\ndata: /sseendpoint\r\n\r\n")
});
}
else if (callIndex == 2)
{
return Task.FromResult(new HttpResponseMessage
{
StatusCode = HttpStatusCode.OK,
Content = new StringContent("event: message\r\ndata: {\"jsonrpc\":\"2.0\", \"id\": \"44\", \"method\": \"test\", \"params\": null}\r\n\r\n")
Content = new StringContent("event: endpoint\r\ndata: /sseendpoint\r\n\r\nevent: message\r\ndata: {\"jsonrpc\":\"2.0\", \"id\": \"44\", \"method\": \"test\", \"params\": null}\r\n\r\n")
});
}
throw new IOException("Abort");
};
await transport.ConnectAsync(TestContext.Current.CancellationToken);
Assert.True(transport.MessageReader.TryRead(out var message));
await using var session = await transport.ConnectAsync(TestContext.Current.CancellationToken);
Assert.True(session.MessageReader.TryRead(out var message));
Assert.NotNull(message);
Assert.IsType<JsonRpcRequest>(message);
Assert.Equal("44", ((JsonRpcRequest)message).Id.AsString);
}
[Fact]
public async Task CloseAsync_Should_Dispose_Resources()
{
await using var transport = new SseClientTransport(_transportOptions, _serverConfig, LoggerFactory);
await transport.CloseAsync();
Assert.False(transport.IsConnected);
}
[Fact]
public async Task DisposeAsync_Should_Dispose_Resources()
{
await using var transport = new SseClientTransport(_transportOptions, _serverConfig, LoggerFactory);
using var mockHttpHandler = new MockHttpHandler();
using var httpClient = new HttpClient(mockHttpHandler);
mockHttpHandler.RequestHandler = request =>
{
return Task.FromResult(new HttpResponseMessage
{
StatusCode = HttpStatusCode.OK,
Content = new StringContent("event: endpoint\r\ndata: http://localhost\r\n\r\n")
});
};
await transport.DisposeAsync();
await using var transport = new SseClientTransport(_transportOptions, _serverConfig, httpClient, LoggerFactory);
await using var session = await transport.ConnectAsync(TestContext.Current.CancellationToken);
Assert.False(transport.IsConnected);
await session.DisposeAsync();
Assert.False(session.IsConnected);
}
}

View file

@ -1,10 +1,12 @@
using ModelContextProtocol.Protocol.Messages;
using Microsoft.Extensions.Options;
using ModelContextProtocol.Protocol.Messages;
using ModelContextProtocol.Protocol.Transport;
using ModelContextProtocol.Protocol.Types;
using ModelContextProtocol.Server;
using ModelContextProtocol.Tests.Utils;
using ModelContextProtocol.Utils.Json;
using System.IO.Pipelines;
using System.Runtime.InteropServices;
using System.Text;
using System.Text.Json;
@ -33,6 +35,8 @@ public class StdioServerTransportTests : LoggedTest
[Fact]
public async Task Constructor_Should_Initialize_With_Valid_Parameters()
{
Assert.SkipWhen(RuntimeInformation.IsOSPlatform(OSPlatform.Linux), "https://github.com/modelcontextprotocol/csharp-sdk/issues/143");
// Act
await using var transport = new StdioServerTransport(_serverOptions);
@ -45,17 +49,16 @@ public class StdioServerTransportTests : LoggedTest
{
Assert.Throws<ArgumentNullException>("serverName", () => new StdioServerTransport((string)null!));
Assert.Throws<ArgumentNullException>("serverOptions", () => new StdioServerTransport((IOptions<McpServerOptions>)null!));
Assert.Throws<ArgumentNullException>("serverOptions", () => new StdioServerTransport((McpServerOptions)null!));
Assert.Throws<ArgumentNullException>("serverOptions.ServerInfo", () => new StdioServerTransport(new McpServerOptions() { ServerInfo = null! }));
Assert.Throws<ArgumentNullException>("serverName", () => new StdioServerTransport(new McpServerOptions() { ServerInfo = new() { Name = null!, Version = "" } }));
Assert.Throws<ArgumentNullException>("serverOptions.ServerInfo.Name", () => new StdioServerTransport(new McpServerOptions() { ServerInfo = new() { Name = null!, Version = "" } }));
}
[Fact]
public async Task StartListeningAsync_Should_Set_Connected_State()
public async Task Should_Start_In_Connected_State()
{
await using var transport = new StdioServerTransport(_serverOptions);
await transport.StartListeningAsync(TestContext.Current.CancellationToken);
await using var transport = new StdioServerTransport(_serverOptions.ServerInfo.Name, new Pipe().Reader.AsStream(), Stream.Null, LoggerFactory);
Assert.True(transport.IsConnected);
}
@ -70,12 +73,7 @@ public class StdioServerTransportTests : LoggedTest
new Pipe().Reader.AsStream(),
output,
LoggerFactory);
await transport.StartListeningAsync(TestContext.Current.CancellationToken);
// Ensure transport is fully initialized
await Task.Delay(100, TestContext.Current.CancellationToken);
// Verify transport is connected
Assert.True(transport.IsConnected, "Transport should be connected after StartListeningAsync");
@ -89,20 +87,10 @@ public class StdioServerTransportTests : LoggedTest
Assert.Equal(expected, result);
}
[Fact]
public async Task SendMessageAsync_Throws_Exception_If_Not_Connected()
{
await using var transport = new StdioServerTransport(_serverOptions);
var message = new JsonRpcRequest { Method = "test" };
await Assert.ThrowsAsync<McpTransportException>(() => transport.SendMessageAsync(message, TestContext.Current.CancellationToken));
}
[Fact]
public async Task DisposeAsync_Should_Dispose_Resources()
{
await using var transport = new StdioServerTransport(_serverOptions);
await using var transport = new StdioServerTransport(_serverOptions.ServerInfo.Name, Stream.Null, Stream.Null, LoggerFactory);
await transport.DisposeAsync();
@ -124,12 +112,7 @@ public class StdioServerTransportTests : LoggedTest
input,
Stream.Null,
LoggerFactory);
await transport.StartListeningAsync(TestContext.Current.CancellationToken);
// Ensure transport is fully initialized
await Task.Delay(100, TestContext.Current.CancellationToken);
// Verify transport is connected
Assert.True(transport.IsConnected, "Transport should be connected after StartListeningAsync");
@ -148,8 +131,7 @@ public class StdioServerTransportTests : LoggedTest
[Fact]
public async Task CleanupAsync_Should_Cleanup_Resources()
{
var transport = new StdioServerTransport(_serverOptions);
await transport.StartListeningAsync(TestContext.Current.CancellationToken);
var transport = new StdioServerTransport(_serverOptions.ServerInfo.Name, Stream.Null, Stream.Null, LoggerFactory);
await transport.DisposeAsync();
@ -163,24 +145,19 @@ public class StdioServerTransportTests : LoggedTest
using var output = new MemoryStream();
await using var transport = new StdioServerTransport(
_serverOptions.ServerInfo.Name,
_serverOptions.ServerInfo.Name,
new Pipe().Reader.AsStream(),
output,
output,
LoggerFactory);
await transport.StartListeningAsync(TestContext.Current.CancellationToken);
// Ensure transport is fully initialized
await Task.Delay(100, TestContext.Current.CancellationToken);
// Verify transport is connected
Assert.True(transport.IsConnected, "Transport should be connected after StartListeningAsync");
// Test 1: Chinese characters (BMP Unicode)
var chineseText = "上下文伺服器"; // "Context Server" in Chinese
var chineseMessage = new JsonRpcRequest
{
Method = "test",
var chineseMessage = new JsonRpcRequest
{
Method = "test",
Id = RequestId.FromNumber(44),
Params = new Dictionary<string, JsonElement>
{
@ -191,18 +168,18 @@ public class StdioServerTransportTests : LoggedTest
// Clear output and send message
output.SetLength(0);
await transport.SendMessageAsync(chineseMessage, TestContext.Current.CancellationToken);
// Verify Chinese characters preserved but encoded
var chineseResult = Encoding.UTF8.GetString(output.ToArray()).Trim();
var expectedChinese = JsonSerializer.Serialize(chineseMessage, McpJsonUtilities.DefaultOptions);
Assert.Equal(expectedChinese, chineseResult);
Assert.Contains(JsonSerializer.Serialize(chineseText), chineseResult);
// Test 2: Emoji (non-BMP Unicode using surrogate pairs)
var emojiText = "🔍 🚀 👍"; // Magnifying glass, rocket, thumbs up
var emojiMessage = new JsonRpcRequest
{
Method = "test",
var emojiMessage = new JsonRpcRequest
{
Method = "test",
Id = RequestId.FromNumber(45),
Params = new Dictionary<string, JsonElement>
{
@ -213,23 +190,23 @@ public class StdioServerTransportTests : LoggedTest
// Clear output and send message
output.SetLength(0);
await transport.SendMessageAsync(emojiMessage, TestContext.Current.CancellationToken);
// Verify emoji preserved - might be as either direct characters or escape sequences
var emojiResult = Encoding.UTF8.GetString(output.ToArray()).Trim();
var expectedEmoji = JsonSerializer.Serialize(emojiMessage, McpJsonUtilities.DefaultOptions);
Assert.Equal(expectedEmoji, emojiResult);
// Verify surrogate pairs in different possible formats
// Magnifying glass emoji: 🔍 (U+1F50D)
bool magnifyingGlassFound =
emojiResult.Contains("🔍") ||
bool magnifyingGlassFound =
emojiResult.Contains("🔍") ||
emojiResult.Contains("\\ud83d\\udd0d", StringComparison.OrdinalIgnoreCase);
// Rocket emoji: 🚀 (U+1F680)
bool rocketFound =
emojiResult.Contains("🚀") ||
bool rocketFound =
emojiResult.Contains("🚀") ||
emojiResult.Contains("\\ud83d\\ude80", StringComparison.OrdinalIgnoreCase);
Assert.True(magnifyingGlassFound, "Magnifying glass emoji not found in result");
Assert.True(rocketFound, "Rocket emoji not found in result");
}

View file

@ -1,4 +1,4 @@
using Microsoft.Extensions.Logging.Abstractions;
using Microsoft.Extensions.Logging;
using ModelContextProtocol.Protocol.Messages;
using ModelContextProtocol.Protocol.Transport;
using ModelContextProtocol.Utils.Json;
@ -14,8 +14,8 @@ internal sealed class StreamClientTransport : TransportBase, IClientTransport
private readonly TextReader _serverStdoutReader;
private readonly TextWriter _serverStdinWriter;
public StreamClientTransport(TextWriter serverStdinWriter, TextReader serverStdoutReader)
: base(NullLoggerFactory.Instance)
public StreamClientTransport(TextWriter serverStdinWriter, TextReader serverStdoutReader, ILoggerFactory loggerFactory)
: base(loggerFactory)
{
_serverStdoutReader = serverStdoutReader;
_serverStdinWriter = serverStdinWriter;
@ -23,7 +23,7 @@ internal sealed class StreamClientTransport : TransportBase, IClientTransport
SetConnected(true);
}
public Task ConnectAsync(CancellationToken cancellationToken = default) => Task.CompletedTask;
public Task<ITransport> ConnectAsync(CancellationToken cancellationToken = default) => Task.FromResult<ITransport>(this);
public override async Task SendMessageAsync(IJsonRpcMessage message, CancellationToken cancellationToken = default)
{
@ -37,22 +37,28 @@ internal sealed class StreamClientTransport : TransportBase, IClientTransport
private async Task ReadMessagesAsync(CancellationToken cancellationToken)
{
while (await _serverStdoutReader.ReadLineAsync(cancellationToken).ConfigureAwait(false) is string line)
try
{
if (!string.IsNullOrWhiteSpace(line))
while (await _serverStdoutReader.ReadLineAsync(cancellationToken).ConfigureAwait(false) is string line)
{
try
if (!string.IsNullOrWhiteSpace(line))
{
if (JsonSerializer.Deserialize<IJsonRpcMessage>(line.Trim(), _jsonOptions) is { } message)
try
{
if (JsonSerializer.Deserialize<IJsonRpcMessage>(line.Trim(), _jsonOptions) is { } message)
{
await WriteMessageAsync(message, cancellationToken).ConfigureAwait(false);
}
}
catch (JsonException)
{
await WriteMessageAsync(message, cancellationToken).ConfigureAwait(false);
}
}
catch (JsonException)
{
}
}
}
catch (OperationCanceledException)
{
}
}
public override async ValueTask DisposeAsync()