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

Fix and enhance cancellation operations across MCP Sessions. (#179)

Browse source 
* Propagate CancellationToken request cancellation to remote endpoint

* Refactor existing tests to have a shared base class

* Add cancellation tests

---------

Co-authored-by: Stephen Toub <stoub@microsoft.com>
This commit is contained in:
Tyler Kendrick 2025-04-08 15:12:12 -05:00 • committed by GitHub
commit 03e3094dcc
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 334 additions and 336 deletions
Download patch file Download diff file Expand all files Collapse all files

69
tests/ModelContextProtocol.Tests/Protocol/CancellationTests.cs Normal file
View file

@ -0,0 +1,69 @@
using Microsoft.Extensions.DependencyInjection;
using ModelContextProtocol.Client;
using ModelContextProtocol.Protocol.Messages;
using ModelContextProtocol.Server;
namespace ModelContextProtocol.Tests;
public class CancellationTests : ClientServerTestBase
{
public CancellationTests(ITestOutputHelper testOutputHelper)
: base(testOutputHelper)
{
}
protected override void ConfigureServices(ServiceCollection services, IMcpServerBuilder mcpServerBuilder)
{
services.AddSingleton(McpServerTool.Create(WaitForCancellation));
}
private static async Task WaitForCancellation(CancellationToken cancellationToken)
{
try
{
await Task.Delay(-1, cancellationToken);
throw new InvalidOperationException("Unexpected completion without exception");
}
catch (OperationCanceledException)
{
return;
}
}
[Fact]
public async Task PrecancelRequest_CancelsBeforeSending()
{
var client = await CreateMcpClientForServer();
bool gotCancellation = false;
await using (Server.RegisterNotificationHandler(NotificationMethods.CancelledNotification, (notification, cancellationToken) =>
{
gotCancellation = true;
return Task.CompletedTask;
}))
{
await Assert.ThrowsAsync<OperationCanceledException>(() => client.ListToolsAsync(cancellationToken: new CancellationToken(true)));
}
Assert.False(gotCancellation);
}
[Fact]
public async Task CancellationPropagation_RequestingCancellationCancelsPendingRequest()
{
var client = await CreateMcpClientForServer();
var tools = await client.ListToolsAsync(cancellationToken: TestContext.Current.CancellationToken);
var waitTool = tools.First(t => t.Name == nameof(WaitForCancellation));
CancellationTokenSource cts = new();
var waitTask = waitTool.InvokeAsync(cancellationToken: cts.Token);
Assert.False(waitTask.IsCompleted);
await Task.Delay(1, TestContext.Current.CancellationToken);
Assert.False(waitTask.IsCompleted);
cts.Cancel();
await Assert.ThrowsAnyAsync<OperationCanceledException>(async () => await waitTask);
}
}

63
tests/ModelContextProtocol.Tests/Protocol/NotificationHandlerTests.cs
View file

@ -1,61 +1,12 @@
using Microsoft.Extensions.AI;
using Microsoft.Extensions.DependencyInjection;
using ModelContextProtocol.Client;
using ModelContextProtocol.Protocol.Transport;
using ModelContextProtocol.Server;
using ModelContextProtocol.Tests.Utils;
using System.IO.Pipelines;
using ModelContextProtocol.Client;
namespace ModelContextProtocol.Tests;
public class NotificationHandlerTests : LoggedTest, IAsyncDisposable
public class NotificationHandlerTests : ClientServerTestBase
{
private readonly Pipe _clientToServerPipe = new();
private readonly Pipe _serverToClientPipe = new();
private readonly ServiceProvider _serviceProvider;
private readonly IMcpServerBuilder _builder;
private readonly CancellationTokenSource _cts;
private readonly Task _serverTask;
private readonly IMcpServer _server;
public NotificationHandlerTests(ITestOutputHelper testOutputHelper)
: base(testOutputHelper)
{
ServiceCollection sc = new();
sc.AddSingleton(LoggerFactory);
_builder = sc
.AddMcpServer()
.WithStreamServerTransport(_clientToServerPipe.Reader.AsStream(), _serverToClientPipe.Writer.AsStream());
_serviceProvider = sc.BuildServiceProvider();
_cts = CancellationTokenSource.CreateLinkedTokenSource(TestContext.Current.CancellationToken);
_server = _serviceProvider.GetRequiredService<IMcpServer>();
_serverTask = _server.RunAsync(_cts.Token);
}
public async ValueTask DisposeAsync()
{
await _cts.CancelAsync();
_clientToServerPipe.Writer.Complete();
_serverToClientPipe.Writer.Complete();
await _serverTask;
await _serviceProvider.DisposeAsync();
_cts.Dispose();
Dispose();
}
private async Task<IMcpClient> CreateMcpClientForServer(McpClientOptions? options = null)
{
return await McpClientFactory.CreateAsync(
new StreamClientTransport(
serverInput: _clientToServerPipe.Writer.AsStream(),
_serverToClientPipe.Reader.AsStream(),
LoggerFactory),
loggerFactory: LoggerFactory,
cancellationToken: TestContext.Current.CancellationToken);
}
[Fact]
@ -77,7 +28,7 @@ public class NotificationHandlerTests : LoggedTest, IAsyncDisposable
return Task.CompletedTask;
}))
{
await _server.SendNotificationAsync(NotificationName, TestContext.Current.CancellationToken);
await Server.SendNotificationAsync(NotificationName, TestContext.Current.CancellationToken);
await tcs.Task;
}
}
@ -113,7 +64,7 @@ public class NotificationHandlerTests : LoggedTest, IAsyncDisposable
try
{
await _server.SendNotificationAsync(NotificationName, TestContext.Current.CancellationToken);
await Server.SendNotificationAsync(NotificationName, TestContext.Current.CancellationToken);
await tcs.Task;
}
finally
@ -153,7 +104,7 @@ public class NotificationHandlerTests : LoggedTest, IAsyncDisposable
try
{
await _server.SendNotificationAsync(NotificationName, TestContext.Current.CancellationToken);
await Server.SendNotificationAsync(NotificationName, TestContext.Current.CancellationToken);
await tcs.Task;
}
finally
@ -182,7 +133,7 @@ public class NotificationHandlerTests : LoggedTest, IAsyncDisposable
await releaseHandler.Task;
});
await _server.SendNotificationAsync(NotificationName, TestContext.Current.CancellationToken);
await Server.SendNotificationAsync(NotificationName, TestContext.Current.CancellationToken);
await handlerRunning.Task;
var disposals = new ValueTask[numberOfDisposals];
@ -231,7 +182,7 @@ public class NotificationHandlerTests : LoggedTest, IAsyncDisposable
handlerRunning.SetResult(true);
});
await _server.SendNotificationAsync(NotificationName, TestContext.Current.CancellationToken);
await Server.SendNotificationAsync(NotificationName, TestContext.Current.CancellationToken);
await handlerRunning.Task;
}
}