Move notification handler registrations to capabilities (#207)

* Move notification handler registrations to capabilities

Currently request handlers are set on the capability objects, but notification handlers are set after construction via an AddNotificationHandler method on the IMcpEndpoint interface. This moves handler specification to be at construction as well. This makes it more consistent with request handlers, simplifies the IMcpEndpoint interface to just be about message sending, and avoids a concurrency bug that could occur if someone tried to add a handler while the endpoint was processing notifications.

* Address more feedback and further cleanup
This commit is contained in:
Stephen Toub 2025-04-04 13:11:07 -04:00 • committed by GitHub
commit d21d933876
No known key found for this signature in database
GPG key ID: B5690EEEBB952194

View file

@ -141,15 +141,16 @@ public class McpServerBuilderExtensionsToolsTests : LoggedTest, IAsyncDisposable
Dispose();
}
private async Task<IMcpClient> CreateMcpClientForServer()
private async Task<IMcpClient> CreateMcpClientForServer(McpClientOptions? options = null)
{
return await McpClientFactory.CreateAsync(
new McpServerConfig()
{
Id = "TestServer",
Name = "TestServer",
TransportType = "ignored",
},
new McpServerConfig()
{
Id = "TestServer",
Name = "TestServer",
TransportType = "ignored",
},
options,
createTransportFunc: (_, _) => new StreamClientTransport(
serverInput: _clientToServerPipe.Writer.AsStream(),
_serverToClientPipe.Reader.AsStream(),
@ -242,18 +243,23 @@ public class McpServerBuilderExtensionsToolsTests : LoggedTest, IAsyncDisposable
[Fact]
public async Task Can_Be_Notified_Of_Tool_Changes()
{
IMcpClient client = await CreateMcpClientForServer();
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;
})],
},
});
var tools = await client.ListToolsAsync(cancellationToken: TestContext.Current.CancellationToken);
Assert.Equal(16, tools.Count);
Channel<JsonRpcNotification> listChanged = Channel.CreateUnbounded<JsonRpcNotification>();
client.AddNotificationHandler(NotificationMethods.ToolListChangedNotification, notification =>
{
listChanged.Writer.TryWrite(notification);
return Task.CompletedTask;
});
var notificationRead = listChanged.Reader.ReadAsync(TestContext.Current.CancellationToken);
Assert.False(notificationRead.IsCompleted);
@ -622,12 +628,17 @@ public class McpServerBuilderExtensionsToolsTests : LoggedTest, IAsyncDisposable
{
ConcurrentQueue<ProgressNotification> notifications = new();
IMcpClient client = await CreateMcpClientForServer();
client.AddNotificationHandler(NotificationMethods.ProgressNotification, notification =>
IMcpClient client = await CreateMcpClientForServer(new()
{
ProgressNotification pn = JsonSerializer.Deserialize<ProgressNotification>(notification.Params)!;
notifications.Enqueue(pn);
return Task.CompletedTask;
Capabilities = new()
{
NotificationHandlers = [new(NotificationMethods.ProgressNotification, notification =>
{
ProgressNotification pn = JsonSerializer.Deserialize<ProgressNotification>(notification.Params)!;
notifications.Enqueue(pn);
return Task.CompletedTask;
})],
},
});
var tools = await client.ListToolsAsync(cancellationToken: TestContext.Current.CancellationToken);