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

Allow notification handlers post-creation temporary registration (#223)

Browse source 
- We still allow notification handlers to be configured pre-creation, eliminating any race conditions that may result from missed notifications. But now we also allow for post-creation notification handler registration.
- Handlers may now also be removed. This permits temporary registrations that may be created in a scoped manner.
- Registration handling is thread-safe.
- Handlers are now cancelable.
- Exceptions from notification handlers trigger all normal exception handling in the dispatch pipeline.

This also:
- Adds a WithStreamServerTransport as a counterpart to WithStdioServerTransport. The latter was used in several tests in a hacky way, and that's now simplified via the former's existence.
This commit is contained in:
Stephen Toub 2025-04-06 21:17:00 -04:00 • committed by GitHub
commit 848c9df433
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 704 additions and 172 deletions
Download patch file Download diff file Expand all files Collapse all files

88
tests/ModelContextProtocol.Tests/Configuration/McpServerBuilderExtensionsToolsTests.cs
View file

@ -35,7 +35,7 @@ public class McpServerBuilderExtensionsToolsTests : LoggedTest, IAsyncDisposable
sc.AddSingleton(LoggerFactory);
_builder = sc
.AddMcpServer()
.WithStdioServerTransport()
.WithStreamServerTransport(_clientToServerPipe.Reader.AsStream(), _serverToClientPipe.Writer.AsStream())
.WithListToolsHandler(async (request, cancellationToken) =>
{
var cursor = request.Params?.Cursor;
@ -117,8 +117,6 @@ public class McpServerBuilderExtensionsToolsTests : LoggedTest, IAsyncDisposable
})
.WithTools<EchoTool>();
// Call WithStdioServerTransport to get the IMcpServer registration, then overwrite default transport with a pipe transport.
sc.AddSingleton<ITransport>(new StreamServerTransport(_clientToServerPipe.Reader.AsStream(), _serverToClientPipe.Writer.AsStream(), loggerFactory: LoggerFactory));
sc.AddSingleton(new ObjectWithId());
_serviceProvider = sc.BuildServiceProvider();
@ -243,23 +241,12 @@ public class McpServerBuilderExtensionsToolsTests : LoggedTest, IAsyncDisposable
[Fact]
public async Task Can_Be_Notified_Of_Tool_Changes()
{
Channel<JsonRpcNotification> listChanged = Channel.CreateUnbounded<JsonRpcNotification>();
IMcpClient client = await CreateMcpClientForServer(new()
{
Capabilities = new()
{
NotificationHandlers = [new(NotificationMethods.ToolListChangedNotification, notification =>
{
listChanged.Writer.TryWrite(notification);
return Task.CompletedTask;
})],
},
});
IMcpClient client = await CreateMcpClientForServer();
var tools = await client.ListToolsAsync(cancellationToken: TestContext.Current.CancellationToken);
Assert.Equal(16, tools.Count);
Channel<JsonRpcNotification> listChanged = Channel.CreateUnbounded<JsonRpcNotification>();
var notificationRead = listChanged.Reader.ReadAsync(TestContext.Current.CancellationToken);
Assert.False(notificationRead.IsCompleted);
@ -268,17 +255,24 @@ public class McpServerBuilderExtensionsToolsTests : LoggedTest, IAsyncDisposable
Assert.NotNull(serverTools);
var newTool = McpServerTool.Create([McpServerTool(Name = "NewTool")] () => "42");
serverTools.Add(newTool);
await notificationRead;
await using (client.RegisterNotificationHandler(NotificationMethods.ToolListChangedNotification, (notification, cancellationToken) =>
{
listChanged.Writer.TryWrite(notification);
return Task.CompletedTask;
}))
{
serverTools.Add(newTool);
await notificationRead;
tools = await client.ListToolsAsync(cancellationToken: TestContext.Current.CancellationToken);
Assert.Equal(17, tools.Count);
Assert.Contains(tools, t => t.Name == "NewTool");
tools = await client.ListToolsAsync(cancellationToken: TestContext.Current.CancellationToken);
Assert.Equal(17, tools.Count);
Assert.Contains(tools, t => t.Name == "NewTool");
notificationRead = listChanged.Reader.ReadAsync(TestContext.Current.CancellationToken);
Assert.False(notificationRead.IsCompleted);
serverTools.Remove(newTool);
await notificationRead;
notificationRead = listChanged.Reader.ReadAsync(TestContext.Current.CancellationToken);
Assert.False(notificationRead.IsCompleted);
serverTools.Remove(newTool);
await notificationRead;
}
tools = await client.ListToolsAsync(cancellationToken: TestContext.Current.CancellationToken);
Assert.Equal(16, tools.Count);
@ -626,20 +620,7 @@ public class McpServerBuilderExtensionsToolsTests : LoggedTest, IAsyncDisposable
[Fact]
public async Task HandlesIProgressParameter()
{
ConcurrentQueue<ProgressNotification> notifications = new();
IMcpClient client = await CreateMcpClientForServer(new()
{
Capabilities = new()
{
NotificationHandlers = [new(NotificationMethods.ProgressNotification, notification =>
{
ProgressNotification pn = JsonSerializer.Deserialize<ProgressNotification>(notification.Params)!;
notifications.Enqueue(pn);
return Task.CompletedTask;
})],
},
});
IMcpClient client = await CreateMcpClientForServer();
var tools = await client.ListToolsAsync(cancellationToken: TestContext.Current.CancellationToken);
Assert.NotNull(tools);
@ -647,17 +628,26 @@ public class McpServerBuilderExtensionsToolsTests : LoggedTest, IAsyncDisposable
McpClientTool progressTool = tools.First(t => t.Name == nameof(EchoTool.SendsProgressNotifications));
var result = await client.SendRequestAsync<CallToolRequestParams, CallToolResponse>(
RequestMethods.ToolsCall,
new CallToolRequestParams
{
Name = progressTool.ProtocolTool.Name,
Meta = new() { ProgressToken = new("abc123") },
},
cancellationToken: TestContext.Current.CancellationToken);
ConcurrentQueue<ProgressNotification> notifications = new();
await using (client.RegisterNotificationHandler(NotificationMethods.ProgressNotification, (notification, cancellationToken) =>
{
ProgressNotification pn = JsonSerializer.Deserialize<ProgressNotification>(notification.Params)!;
notifications.Enqueue(pn);
return Task.CompletedTask;
}))
{
var result = await client.SendRequestAsync<CallToolRequestParams, CallToolResponse>(
RequestMethods.ToolsCall,
new CallToolRequestParams
{
Name = progressTool.ProtocolTool.Name,
Meta = new() { ProgressToken = new("abc123") },
},
cancellationToken: TestContext.Current.CancellationToken);
Assert.Contains("done", JsonSerializer.Serialize(result));
SpinWait.SpinUntil(() => notifications.Count == 10, TimeSpan.FromSeconds(10));
Assert.Contains("done", JsonSerializer.Serialize(result));
SpinWait.SpinUntil(() => notifications.Count == 10, TimeSpan.FromSeconds(10));
}
ProgressNotification[] array = notifications.OrderBy(n => n.Progress.Progress).ToArray();
Assert.Equal(10, array.Length);