848c9df433- 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.
329 lines
13 KiB
C#
329 lines
13 KiB
C#
using Microsoft.Extensions.AI;
|
|
using Microsoft.Extensions.DependencyInjection;
|
|
using Microsoft.Extensions.Options;
|
|
using ModelContextProtocol.Client;
|
|
using ModelContextProtocol.Protocol.Messages;
|
|
using ModelContextProtocol.Protocol.Transport;
|
|
using ModelContextProtocol.Protocol.Types;
|
|
using ModelContextProtocol.Server;
|
|
using ModelContextProtocol.Tests.Utils;
|
|
using System.ComponentModel;
|
|
using System.IO.Pipelines;
|
|
using System.Threading.Channels;
|
|
|
|
#pragma warning disable CS1998 // Async method lacks 'await' operators and will run synchronously
|
|
|
|
namespace ModelContextProtocol.Tests.Configuration;
|
|
|
|
public class McpServerBuilderExtensionsPromptsTests : LoggedTest, IAsyncDisposable
|
|
{
|
|
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;
|
|
|
|
public McpServerBuilderExtensionsPromptsTests(ITestOutputHelper testOutputHelper)
|
|
: base(testOutputHelper)
|
|
{
|
|
ServiceCollection sc = new();
|
|
sc.AddSingleton(LoggerFactory);
|
|
_builder = sc
|
|
.AddMcpServer()
|
|
.WithStreamServerTransport(_clientToServerPipe.Reader.AsStream(), _serverToClientPipe.Writer.AsStream())
|
|
.WithListPromptsHandler(async (request, cancellationToken) =>
|
|
{
|
|
var cursor = request.Params?.Cursor;
|
|
switch (cursor)
|
|
{
|
|
case null:
|
|
return new()
|
|
{
|
|
NextCursor = "abc",
|
|
Prompts = [new()
|
|
{
|
|
Name = "FirstCustomPrompt",
|
|
Description = "First prompt returned by custom handler",
|
|
}],
|
|
};
|
|
|
|
case "abc":
|
|
return new()
|
|
{
|
|
NextCursor = "def",
|
|
Prompts = [new()
|
|
{
|
|
Name = "SecondCustomPrompt",
|
|
Description = "Second prompt returned by custom handler",
|
|
}],
|
|
};
|
|
|
|
case "def":
|
|
return new()
|
|
{
|
|
NextCursor = null,
|
|
Prompts = [new()
|
|
{
|
|
Name = "FinalCustomPrompt",
|
|
Description = "Final prompt returned by custom handler",
|
|
}],
|
|
};
|
|
|
|
default:
|
|
throw new Exception("Unexpected cursor");
|
|
}
|
|
})
|
|
.WithGetPromptHandler(async (request, cancellationToken) =>
|
|
{
|
|
switch (request.Params?.Name)
|
|
{
|
|
case "FirstCustomPrompt":
|
|
case "SecondCustomPrompt":
|
|
case "FinalCustomPrompt":
|
|
return new GetPromptResult()
|
|
{
|
|
Messages = [new() { Role = Role.User, Content = new() { Text = $"hello from {request.Params.Name}", Type = "text" } }],
|
|
};
|
|
|
|
default:
|
|
throw new Exception($"Unknown prompt '{request.Params?.Name}'");
|
|
}
|
|
})
|
|
.WithPrompts<SimplePrompts>();
|
|
|
|
|
|
sc.AddSingleton(new ObjectWithId());
|
|
_serviceProvider = sc.BuildServiceProvider();
|
|
|
|
var server = _serviceProvider.GetRequiredService<IMcpServer>();
|
|
_cts = CancellationTokenSource.CreateLinkedTokenSource(TestContext.Current.CancellationToken);
|
|
_serverTask = server.RunAsync(cancellationToken: _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 McpServerConfig()
|
|
{
|
|
Id = "TestServer",
|
|
Name = "TestServer",
|
|
TransportType = "ignored",
|
|
},
|
|
options,
|
|
createTransportFunc: (_, _) => new StreamClientTransport(
|
|
serverInput: _clientToServerPipe.Writer.AsStream(),
|
|
serverOutput: _serverToClientPipe.Reader.AsStream(),
|
|
LoggerFactory),
|
|
loggerFactory: LoggerFactory,
|
|
cancellationToken: TestContext.Current.CancellationToken);
|
|
}
|
|
|
|
[Fact]
|
|
public void Adds_Prompts_To_Server()
|
|
{
|
|
var serverOptions = _serviceProvider.GetRequiredService<IOptions<McpServerOptions>>().Value;
|
|
var prompts = serverOptions?.Capabilities?.Prompts?.PromptCollection;
|
|
Assert.NotNull(prompts);
|
|
Assert.NotEmpty(prompts);
|
|
}
|
|
|
|
[Fact]
|
|
public async Task Can_List_And_Call_Registered_Prompts()
|
|
{
|
|
IMcpClient client = await CreateMcpClientForServer();
|
|
|
|
var prompts = await client.ListPromptsAsync(TestContext.Current.CancellationToken);
|
|
Assert.Equal(6, prompts.Count);
|
|
|
|
var prompt = prompts.First(t => t.Name == nameof(SimplePrompts.ReturnsChatMessages));
|
|
Assert.Equal("Returns chat messages", prompt.Description);
|
|
|
|
var result = await prompt.GetAsync(new Dictionary<string, object?>() { ["message"] = "hello" }, cancellationToken: TestContext.Current.CancellationToken);
|
|
var chatMessages = result.ToChatMessages();
|
|
|
|
Assert.NotNull(chatMessages);
|
|
Assert.NotEmpty(chatMessages);
|
|
Assert.Equal(2, chatMessages.Count);
|
|
Assert.Equal("The prompt is: hello", chatMessages[0].Text);
|
|
Assert.Equal("Summarize.", chatMessages[1].Text);
|
|
|
|
prompt = prompts.First(t => t.Name == "SecondCustomPrompt");
|
|
Assert.Equal("Second prompt returned by custom handler", prompt.Description);
|
|
result = await prompt.GetAsync(cancellationToken: TestContext.Current.CancellationToken);
|
|
chatMessages = result.ToChatMessages();
|
|
Assert.NotNull(chatMessages);
|
|
Assert.Single(chatMessages);
|
|
Assert.Equal("hello from SecondCustomPrompt", chatMessages[0].Text);
|
|
}
|
|
|
|
[Fact]
|
|
public async Task Can_Be_Notified_Of_Prompt_Changes()
|
|
{
|
|
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);
|
|
|
|
var serverOptions = _serviceProvider.GetRequiredService<IOptions<McpServerOptions>>().Value;
|
|
var serverPrompts = serverOptions.Capabilities?.Prompts?.PromptCollection;
|
|
Assert.NotNull(serverPrompts);
|
|
|
|
var newPrompt = McpServerPrompt.Create([McpServerPrompt(Name = "NewPrompt")] () => "42");
|
|
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");
|
|
|
|
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);
|
|
Assert.DoesNotContain(prompts, t => t.Name == "NewPrompt");
|
|
}
|
|
|
|
[Fact]
|
|
public async Task Throws_When_Prompt_Fails()
|
|
{
|
|
IMcpClient client = await CreateMcpClientForServer();
|
|
|
|
await Assert.ThrowsAsync<McpException>(async () => await client.GetPromptAsync(
|
|
nameof(SimplePrompts.ThrowsException),
|
|
cancellationToken: TestContext.Current.CancellationToken));
|
|
}
|
|
|
|
[Fact]
|
|
public async Task Throws_Exception_On_Unknown_Prompt()
|
|
{
|
|
IMcpClient client = await CreateMcpClientForServer();
|
|
|
|
var e = await Assert.ThrowsAsync<McpException>(async () => await client.GetPromptAsync(
|
|
"NotRegisteredPrompt",
|
|
cancellationToken: TestContext.Current.CancellationToken));
|
|
|
|
Assert.Contains("'NotRegisteredPrompt'", e.Message);
|
|
}
|
|
|
|
[Fact]
|
|
public async Task Throws_Exception_Missing_Parameter()
|
|
{
|
|
IMcpClient client = await CreateMcpClientForServer();
|
|
|
|
var e = await Assert.ThrowsAsync<McpException>(async () => await client.GetPromptAsync(
|
|
nameof(SimplePrompts.ReturnsChatMessages),
|
|
cancellationToken: TestContext.Current.CancellationToken));
|
|
|
|
Assert.Contains("Missing required parameter", e.Message);
|
|
}
|
|
|
|
[Fact]
|
|
public void WithPrompts_InvalidArgs_Throws()
|
|
{
|
|
Assert.Throws<ArgumentNullException>("promptTypes", () => _builder.WithPrompts((IEnumerable<Type>)null!));
|
|
|
|
IMcpServerBuilder nullBuilder = null!;
|
|
Assert.Throws<ArgumentNullException>("builder", () => nullBuilder.WithPrompts<object>());
|
|
Assert.Throws<ArgumentNullException>("builder", () => nullBuilder.WithPrompts(Array.Empty<Type>()));
|
|
Assert.Throws<ArgumentNullException>("builder", () => nullBuilder.WithPromptsFromAssembly());
|
|
}
|
|
|
|
[Fact]
|
|
public void Empty_Enumerables_Is_Allowed()
|
|
{
|
|
_builder.WithPrompts(promptTypes: []); // no exception
|
|
_builder.WithPrompts<object>(); // no exception even though no prompts exposed
|
|
_builder.WithPromptsFromAssembly(typeof(AIFunction).Assembly); // no exception even though no prompts exposed
|
|
}
|
|
|
|
[Fact]
|
|
public void Register_Prompts_From_Current_Assembly()
|
|
{
|
|
ServiceCollection sc = new();
|
|
sc.AddMcpServer().WithPromptsFromAssembly();
|
|
IServiceProvider services = sc.BuildServiceProvider();
|
|
|
|
Assert.Contains(services.GetServices<McpServerPrompt>(), t => t.ProtocolPrompt.Name == nameof(SimplePrompts.ReturnsChatMessages));
|
|
}
|
|
|
|
[Fact]
|
|
public void Register_Prompts_From_Multiple_Sources()
|
|
{
|
|
ServiceCollection sc = new();
|
|
sc.AddMcpServer()
|
|
.WithPrompts<SimplePrompts>()
|
|
.WithPrompts<MorePrompts>();
|
|
IServiceProvider services = sc.BuildServiceProvider();
|
|
|
|
Assert.Contains(services.GetServices<McpServerPrompt>(), t => t.ProtocolPrompt.Name == nameof(SimplePrompts.ReturnsChatMessages));
|
|
Assert.Contains(services.GetServices<McpServerPrompt>(), t => t.ProtocolPrompt.Name == nameof(SimplePrompts.ThrowsException));
|
|
Assert.Contains(services.GetServices<McpServerPrompt>(), t => t.ProtocolPrompt.Name == nameof(SimplePrompts.ReturnsString));
|
|
Assert.Contains(services.GetServices<McpServerPrompt>(), t => t.ProtocolPrompt.Name == nameof(MorePrompts.AnotherPrompt));
|
|
}
|
|
|
|
[McpServerPromptType]
|
|
public sealed class SimplePrompts(ObjectWithId? id = null)
|
|
{
|
|
[McpServerPrompt, Description("Returns chat messages")]
|
|
public static ChatMessage[] ReturnsChatMessages([Description("The first parameter")] string message) =>
|
|
[
|
|
new(ChatRole.User, $"The prompt is: {message}"),
|
|
new(ChatRole.User, "Summarize."),
|
|
];
|
|
|
|
|
|
[McpServerPrompt, Description("Returns chat messages")]
|
|
public static ChatMessage[] ThrowsException([Description("The first parameter")] string message) =>
|
|
throw new FormatException("uh oh");
|
|
|
|
[McpServerPrompt, Description("Returns chat messages")]
|
|
public string ReturnsString([Description("The first parameter")] string message) =>
|
|
$"The prompt is: {message}. The id is {id}.";
|
|
}
|
|
|
|
[McpServerToolType]
|
|
public sealed class MorePrompts
|
|
{
|
|
[McpServerPrompt]
|
|
public static PromptMessage AnotherPrompt() =>
|
|
new PromptMessage
|
|
{
|
|
Role = Role.User,
|
|
Content = new() { Text = "hello", Type = "text" },
|
|
};
|
|
}
|
|
|
|
public class ObjectWithId
|
|
{
|
|
public string Id { get; set; } = Guid.NewGuid().ToString("N");
|
|
}
|
|
}
|