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

- 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

View file

@ -31,7 +31,7 @@ public class McpServerBuilderExtensionsPromptsTests : LoggedTest, IAsyncDisposab
sc.AddSingleton(LoggerFactory);
_builder = sc
.AddMcpServer()
.WithStdioServerTransport()
.WithStreamServerTransport(_clientToServerPipe.Reader.AsStream(), _serverToClientPipe.Writer.AsStream())
.WithListPromptsHandler(async (request, cancellationToken) =>
{
var cursor = request.Params?.Cursor;
@ -93,8 +93,6 @@ public class McpServerBuilderExtensionsPromptsTests : LoggedTest, IAsyncDisposab
.WithPrompts<SimplePrompts>();
// 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();
@ -176,23 +174,12 @@ public class McpServerBuilderExtensionsPromptsTests : LoggedTest, IAsyncDisposab
[Fact]
public async Task Can_Be_Notified_Of_Prompt_Changes()
{
Channel<JsonRpcNotification> listChanged = Channel.CreateUnbounded<JsonRpcNotification>();
IMcpClient client = await CreateMcpClientForServer(new()
{
Capabilities = new()
{
NotificationHandlers = [new("notifications/prompts/list_changed", notification =>
{
listChanged.Writer.TryWrite(notification);
return Task.CompletedTask;
})],
},
});
IMcpClient client = await CreateMcpClientForServer();
var prompts = await client.ListPromptsAsync(TestContext.Current.CancellationToken);
Assert.Equal(6, prompts.Count);
Channel<JsonRpcNotification> listChanged = Channel.CreateUnbounded<JsonRpcNotification>();
var notificationRead = listChanged.Reader.ReadAsync(TestContext.Current.CancellationToken);
Assert.False(notificationRead.IsCompleted);
@ -201,17 +188,24 @@ public class McpServerBuilderExtensionsPromptsTests : LoggedTest, IAsyncDisposab
Assert.NotNull(serverPrompts);
var newPrompt = McpServerPrompt.Create([McpServerPrompt(Name = "NewPrompt")] () => "42");
serverPrompts.Add(newPrompt);
await notificationRead;
await using (client.RegisterNotificationHandler("notifications/prompts/list_changed", (notification, cancellationToken) =>
{
listChanged.Writer.TryWrite(notification);
return Task.CompletedTask;
}))
{
serverPrompts.Add(newPrompt);
await notificationRead;
prompts = await client.ListPromptsAsync(TestContext.Current.CancellationToken);
Assert.Equal(7, prompts.Count);
Assert.Contains(prompts, t => t.Name == "NewPrompt");
prompts = await client.ListPromptsAsync(TestContext.Current.CancellationToken);
Assert.Equal(7, prompts.Count);
Assert.Contains(prompts, t => t.Name == "NewPrompt");
notificationRead = listChanged.Reader.ReadAsync(TestContext.Current.CancellationToken);
Assert.False(notificationRead.IsCompleted);
serverPrompts.Remove(newPrompt);
await notificationRead;
notificationRead = listChanged.Reader.ReadAsync(TestContext.Current.CancellationToken);
Assert.False(notificationRead.IsCompleted);
serverPrompts.Remove(newPrompt);
await notificationRead;
}
prompts = await client.ListPromptsAsync(TestContext.Current.CancellationToken);
Assert.Equal(6, prompts.Count);