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:
parent
f135355bf9
commit
848c9df433
1 changed files with 704 additions and 172 deletions
|
|
@ -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);
|
||||
|
|
|
|||
Loading…
Reference in a new issue