modelcontextprotocol/csharp-sdk
Watch
1
0
Fork
You've already forked csharp-sdk
 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 2

Refactor transports to help enable graceful shutdown (#142)

Browse source 
- 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
1 changed files with 2015 additions and 1847 deletions
Download patch file Download diff file Expand all files Collapse all files

155
tests/ModelContextProtocol.Tests/Transport/SseClientTransportTests.cs
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);
}
}