Extend progress notification support (#163)

* Extend progress notification support

* Address feedback and fix test

---------

Co-authored-by: Stephen Toub <stoub@microsoft.com>
This commit is contained in:
Tyler Kendrick 2025-04-01 10:43:50 -05:00 • committed by GitHub
commit 25bcb440a0
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 175 additions and 154 deletions

View file

@ -355,7 +355,7 @@ public class ClientIntegrationTests : LoggedTest, IClassFixture<ClientIntegratio
{
Sampling = new()
{
SamplingHandler = (_, _) =>
SamplingHandler = (_, _, _) =>
{
samplingHandlerCalls++;
return Task.FromResult(new CreateMessageResult
@ -511,6 +511,9 @@ public class ClientIntegrationTests : LoggedTest, IClassFixture<ClientIntegratio
[Fact(Skip = "Requires OpenAI API Key", SkipWhen = nameof(NoOpenAIKeySet))]
public async Task SamplingViaChatClient_RequestResponseProperlyPropagated()
{
var samplingHandler = new OpenAIClient(s_openAIKey)
.AsChatClient("gpt-4o-mini")
.CreateSamplingHandler();
await using var client = await McpClientFactory.CreateAsync(_fixture.EverythingServerConfig, new()
{
ClientInfo = new() { Name = nameof(SamplingViaChatClient_RequestResponseProperlyPropagated), Version = "1.0.0" },
@ -518,7 +521,7 @@ public class ClientIntegrationTests : LoggedTest, IClassFixture<ClientIntegratio
{
Sampling = new()
{
SamplingHandler = new OpenAIClient(s_openAIKey).AsChatClient("gpt-4o-mini").CreateSamplingHandler(),
SamplingHandler = samplingHandler,
},
},
}, cancellationToken: TestContext.Current.CancellationToken);

View file

@ -635,4 +635,47 @@ public class McpServerTests : LoggedTest
public Task RunAsync(CancellationToken cancellationToken = default) =>
throw new NotImplementedException();
}
[Fact]
public async Task NotifyProgress_Should_Be_Handled()
{
await using TestServerTransport transport = new();
var options = CreateOptions();
var notificationReceived = new TaskCompletionSource<JsonRpcNotification>();
var server = McpServerFactory.Create(transport, options, LoggerFactory, _serviceProvider);
server.AddNotificationHandler(NotificationMethods.ProgressNotification, notification =>
{
notificationReceived.SetResult(notification);
return Task.CompletedTask;
});
Task serverTask = server.RunAsync(TestContext.Current.CancellationToken);
await transport.SendMessageAsync(new JsonRpcNotification
{
Method = NotificationMethods.ProgressNotification,
Params = new ProgressNotification()
{
ProgressToken = new("abc"),
Progress = new()
{
Progress = 50,
Total = 100,
Message = "Progress message",
},
},
}, TestContext.Current.CancellationToken);
var notification = await notificationReceived.Task;
var progress = (ProgressNotification)notification.Params!;
Assert.Equal("\"abc\"", progress.ProgressToken.ToString());
Assert.Equal(50, progress.Progress.Progress);
Assert.Equal(100, progress.Progress.Total);
Assert.Equal("Progress message", progress.Progress.Message);
await server.DisposeAsync();
await serverTask;
}
}

View file

@ -116,7 +116,7 @@ public class SseIntegrationTests(ITestOutputHelper outputHelper) : LoggedTest(ou
{
Sampling = new()
{
SamplingHandler = (_, _) =>
SamplingHandler = (_, _, _) =>
{
samplingHandlerCalls++;
return Task.FromResult(new CreateMessageResult

View file

@ -215,7 +215,7 @@ public class SseServerIntegrationTests : LoggedTest, IClassFixture<SseServerInte
var options = SseServerIntegrationTestFixture.CreateDefaultClientOptions();
options.Capabilities ??= new();
options.Capabilities.Sampling ??= new();
options.Capabilities.Sampling.SamplingHandler = async (_, _) =>
options.Capabilities.Sampling.SamplingHandler = async (_, _, _) =>
{
samplingHandlerCalls++;
return new CreateMessageResult