Extend progress notification support

Tyler-R-Kendrick · stephentoub · commit 6aa993df6129 · 2025-04-01T09:55:56.000-04:00
diff --git a/src/ModelContextProtocol/Client/McpClient.cs b/src/ModelContextProtocol/Client/McpClient.cs
@@ -42,7 +42,7 @@ public McpClient(IClientTransport clientTransport, McpClientOptions options, Mcp
 
             SetRequestHandler<CreateMessageRequestParams, CreateMessageResult>(
                 RequestMethods.SamplingCreateMessage,
-                (request, ct) => samplingHandler(request, ct));
+                (request, ct) => samplingHandler(request, new ClientTokenProgress(this, request?.Meta?.ProgressToken), ct));
         }
 
         if (options.Capabilities?.Roots is { } rootsCapability)
diff --git a/src/ModelContextProtocol/Client/McpClientExtensions.cs b/src/ModelContextProtocol/Client/McpClientExtensions.cs
@@ -531,17 +531,34 @@ internal static CreateMessageResult ToCreateMessageResult(this ChatResponse chat
     /// </summary>
     /// <param name="chatClient">The <see cref="IChatClient"/> with which to satisfy sampling requests.</param>
     /// <returns>The created handler delegate.</returns>
-    public static Func<CreateMessageRequestParams?, CancellationToken, Task<CreateMessageResult>> CreateSamplingHandler(this IChatClient chatClient)
+    public static Func<CreateMessageRequestParams?, IProgress<ProgressNotificationValue>, CancellationToken, Task<CreateMessageResult>> CreateSamplingHandler(
+        this IChatClient chatClient)
     {
         Throw.IfNull(chatClient);
 
-        return async (requestParams, cancellationToken) =>
+        return async (requestParams, progress, cancellationToken) =>
         {
             Throw.IfNull(requestParams);
 
             var (messages, options) = requestParams.ToChatClientArguments();
-            var response = await chatClient.GetResponseAsync(messages, options, cancellationToken).ConfigureAwait(false);
-            return response.ToCreateMessageResult();
+            var progressToken = requestParams.Meta?.ProgressToken;
+            int progressValue = 0;
+            var streamingResponses = chatClient.GetStreamingResponseAsync(
+                messages, options, cancellationToken);
+            List<ChatResponseUpdate> updates = [];
+            await foreach (var streamingResponse in streamingResponses)
+            {
+                updates.Add(streamingResponse);
+                if (progressToken is not null)
+                {
+                    progress.Report(new()
+                    {
+                        Progress = ++progressValue,
+                    });
+                }
+            }
+
+            return updates.ToChatResponse().ToCreateMessageResult();
         };
     }
 
diff --git a/src/ModelContextProtocol/ClientTokenProgress.cs b/src/ModelContextProtocol/ClientTokenProgress.cs
@@ -0,0 +1,28 @@
+using ModelContextProtocol.Client;
+using ModelContextProtocol.Protocol.Messages;
+
+namespace ModelContextProtocol;
+
+internal sealed class ClientTokenProgress(IMcpClient client, ProgressToken? progressToken)
+    : IProgress<ProgressNotificationValue>
+{
+    /// <inheritdoc />
+    public void Report(ProgressNotificationValue value)
+    {
+        if (progressToken is null) return;
+        _ = client.SendMessageAsync(new JsonRpcNotification()
+        {
+            Method = NotificationMethods.ProgressNotification,
+            Params = new ProgressNotification()
+            {
+                ProgressToken = progressToken.Value,
+                Progress = new()
+                {
+                    Progress = value.Progress,
+                    Total = value.Total,
+                    Message = value.Message,
+                },
+            },
+        }, CancellationToken.None);
+    }
+}
diff --git a/src/ModelContextProtocol/Protocol/Types/Capabilities.cs b/src/ModelContextProtocol/Protocol/Types/Capabilities.cs
@@ -55,7 +55,7 @@ public class SamplingCapability
 
     /// <summary>Gets or sets the handler for sampling requests.</summary>
     [JsonIgnore]
-    public Func<CreateMessageRequestParams?, CancellationToken, Task<CreateMessageResult>>? SamplingHandler { get; set; }
+    public Func<CreateMessageRequestParams?, IProgress<ProgressNotificationValue>, CancellationToken, Task<CreateMessageResult>>? SamplingHandler { get; set; }
 }
 
 /// <summary>
diff --git a/src/ModelContextProtocol/Protocol/Types/ListPromptsRequestParams.cs b/src/ModelContextProtocol/Protocol/Types/ListPromptsRequestParams.cs
@@ -4,12 +4,4 @@
 /// Sent from the client to request a list of prompts and prompt templates the server has.
 /// <see href="https://github.com/modelcontextprotocol/specification/blob/main/schema/">See the schema for details</see>
 /// </summary>
-public class ListPromptsRequestParams
-{
-    /// <summary>
-    /// An opaque token representing the current pagination position.
-    /// If provided, the server should return results starting after this cursor.
-    /// </summary>
-    [System.Text.Json.Serialization.JsonPropertyName("cursor")]
-    public string? Cursor { get; init; }
-}
+public class ListPromptsRequestParams : PaginatedRequestParams;
diff --git a/src/ModelContextProtocol/Protocol/Types/ListResourceTemplatesRequestParams.cs b/src/ModelContextProtocol/Protocol/Types/ListResourceTemplatesRequestParams.cs
@@ -4,12 +4,4 @@
 /// Sent from the client to request a list of resource templates the server has.
 /// <see href="https://github.com/modelcontextprotocol/specification/blob/main/schema/">See the schema for details</see>
 /// </summary>
-public class ListResourceTemplatesRequestParams
-{
-    /// <summary>
-    /// An opaque token representing the current pagination position.
-    /// If provided, the server should return results starting after this cursor.
-    /// </summary>
-    [System.Text.Json.Serialization.JsonPropertyName("cursor")]
-    public string? Cursor { get; init; }
-}
+public class ListResourceTemplatesRequestParams : PaginatedRequestParams;
diff --git a/src/ModelContextProtocol/Protocol/Types/ListResourcesRequestParams.cs b/src/ModelContextProtocol/Protocol/Types/ListResourcesRequestParams.cs
@@ -4,12 +4,4 @@
 /// Sent from the client to request a list of resources the server has.
 /// <see href="https://github.com/modelcontextprotocol/specification/blob/main/schema/">See the schema for details</see>
 /// </summary>
-public class ListResourcesRequestParams
-{
-    /// <summary>
-    /// An opaque token representing the current pagination position.
-    /// If provided, the server should return results starting after this cursor.
-    /// </summary>
-    [System.Text.Json.Serialization.JsonPropertyName("cursor")]
-    public string? Cursor { get; init; }
-}
+public class ListResourcesRequestParams : PaginatedRequestParams;
diff --git a/src/ModelContextProtocol/Protocol/Types/ListRootsRequestParams.cs b/src/ModelContextProtocol/Protocol/Types/ListRootsRequestParams.cs
@@ -6,11 +6,4 @@ namespace ModelContextProtocol.Protocol.Types;
 /// A request from the server to get a list of root URIs from the client.
 /// <see href="https://github.com/modelcontextprotocol/specification/blob/main/schema/">See the schema for details</see>
 /// </summary>
-public class ListRootsRequestParams
-{
-    /// <summary>
-    /// Optional progress token for out-of-band progress notifications.
-    /// </summary>
-    [System.Text.Json.Serialization.JsonPropertyName("progressToken")]
-    public ProgressToken? ProgressToken { get; init; }
-}
+public class ListRootsRequestParams : RequestParams;
diff --git a/src/ModelContextProtocol/Protocol/Types/ListToolsRequestParams.cs b/src/ModelContextProtocol/Protocol/Types/ListToolsRequestParams.cs
@@ -4,12 +4,4 @@
 /// Sent from the client to request a list of tools the server has.
 /// <see href="https://github.com/modelcontextprotocol/specification/blob/main/schema/">See the schema for details</see>
 /// </summary>
-public class ListToolsRequestParams
-{
-    /// <summary>
-    /// An opaque token representing the current pagination position.
-    /// If provided, the server should return results starting after this cursor.
-    /// </summary>
-    [System.Text.Json.Serialization.JsonPropertyName("cursor")]
-    public string? Cursor { get; init; }
-}
+public class ListToolsRequestParams : PaginatedRequestParams;
diff --git a/src/ModelContextProtocol/Protocol/Types/PaginatedRequest.cs b/src/ModelContextProtocol/Protocol/Types/PaginatedRequest.cs
@@ -0,0 +1,15 @@
+namespace ModelContextProtocol.Protocol.Types;
+
+/// <summary>
+/// Used as a base class for paginated requests.
+/// <see href="https://github.com/modelcontextprotocol/specification/blob/main/schema/2024-11-05/schema.json">See the schema for details</see>
+/// </summary>
+public class PaginatedRequestParams : RequestParams
+{
+    /// <summary>
+    /// An opaque token representing the current pagination position.
+    /// If provided, the server should return results starting after this cursor.
+    /// </summary>
+    [System.Text.Json.Serialization.JsonPropertyName("cursor")]
+    public string? Cursor { get; init; }
+}
diff --git a/src/ModelContextProtocol/Server/AIFunctionMcpServerTool.cs b/src/ModelContextProtocol/Server/AIFunctionMcpServerTool.cs
@@ -111,7 +111,7 @@ private static TemporaryAIFunctionFactoryOptions CreateAIFunctionFactoryOptions(
                             if (requestContent?.Server is { } server &&
                                 requestContent?.Params?.Meta?.ProgressToken is { } progressToken)
                             {
-                                return new TokenProgress(server, progressToken);
+                                return new ServerTokenProgress(server, progressToken);
                             }
 
                             return NullProgress.Instance;
diff --git a/src/ModelContextProtocol/Server/McpServer.cs b/src/ModelContextProtocol/Server/McpServer.cs
@@ -68,7 +68,6 @@ public McpServer(ITransport transport, McpServerOptions options, ILoggerFactory?
         });
 
         SetToolsHandler(options);
-
         SetInitializeHandler(options);
         SetCompletionHandler(options);
         SetPingHandler();
diff --git a/src/ModelContextProtocol/Server/McpServerExtensions.cs b/src/ModelContextProtocol/Server/McpServerExtensions.cs
@@ -169,6 +169,26 @@ public static Task<ListRootsResult> RequestRootsAsync(
             cancellationToken);
     }
 
+    /// <summary>
+    /// Requests the client to list the roots it exposes.
+    /// </summary>
+    /// <param name="server">The server issueing the request.</param>
+    /// <param name="notification">The notification to send.</param>
+    /// <param name="cancellationToken">A token to cancel the operation.</param>
+    /// <returns>A task containing the response from the client.</returns>
+    public static Task NotifyProgressAsync(
+        this IMcpServer server,
+        ProgressNotification notification,
+        CancellationToken cancellationToken = default)
+    {
+        Throw.IfNull(server);
+        return server.SendMessageAsync(new JsonRpcNotification()
+        {
+            Method = NotificationMethods.ProgressNotification,
+            Params = notification,
+        }, cancellationToken);
+    }
+
     /// <summary>Provides an <see cref="IChatClient"/> implementation that's implemented via client sampling.</summary>
     /// <param name="server"></param>
     private sealed class SamplingChatClient(IMcpServer server) : IChatClient
diff --git a/src/ModelContextProtocol/ServerTokenProgress.cs b/src/ModelContextProtocol/ServerTokenProgress.cs
@@ -0,0 +1,27 @@
+﻿using ModelContextProtocol.Protocol.Messages;
+using ModelContextProtocol.Server;
+
+namespace ModelContextProtocol;
+
+/// <summary>
+/// Provides an <see cref="IProgress{ProgressNotificationValue}"/> tied to a specific progress token and that will issue
+/// progress notifications to the supplied endpoint.
+/// </summary>
+internal sealed class ServerTokenProgress(IMcpServer server, ProgressToken progressToken)
+    : IProgress<ProgressNotificationValue>
+{
+    /// <inheritdoc />
+    public void Report(ProgressNotificationValue value)
+    {
+        _ = server.NotifyProgressAsync(new()
+        {
+            ProgressToken = progressToken,
+            Progress = new()
+            {
+                Progress = value.Progress,
+                Total = value.Total,
+                Message = value.Message,
+            },
+        }, CancellationToken.None);
+    }
+}
diff --git a/tests/ModelContextProtocol.Tests/ClientIntegrationTests.cs b/tests/ModelContextProtocol.Tests/ClientIntegrationTests.cs
@@ -355,7 +355,7 @@ public async Task Sampling_Stdio(string clientId)
             {
                 Sampling = new()
                 {
-                    SamplingHandler = (_, _) =>
+                    SamplingHandler = (_, _, _) =>
                     {
                         samplingHandlerCalls++;
                         return Task.FromResult(new CreateMessageResult
@@ -511,14 +511,17 @@ public async Task ListToolsAsync_UsingEverythingServer_ToolsAreProperlyCalled()
     [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" },
             Capabilities = new()
             {
                 Sampling = new()
                 {
-                    SamplingHandler = new OpenAIClient(s_openAIKey).AsChatClient("gpt-4o-mini").CreateSamplingHandler(),
+                    SamplingHandler = samplingHandler,
                 },
             },
         }, cancellationToken: TestContext.Current.CancellationToken);
diff --git a/tests/ModelContextProtocol.Tests/Server/McpServerTests.cs b/tests/ModelContextProtocol.Tests/Server/McpServerTests.cs
@@ -635,4 +635,85 @@ public Task SendMessageAsync(IJsonRpcMessage message, CancellationToken cancella
         public Task RunAsync(CancellationToken cancellationToken = default) =>
             throw new NotImplementedException();
     }
+
+    [Fact]
+    public async Task NotifyProgress_Should_Be_Handled()
+    {
+        var taskCompletionSource = new TaskCompletionSource();
+        bool notificationHandled = false;
+        await Notifications_Are_Handled(
+            serverCapabilities: null,
+            method: NotificationMethods.ProgressNotification,
+            parameters: new ProgressNotification()
+            {
+                ProgressToken = new(),
+                Progress = new()
+                {
+                    Progress = 50,
+                    Total = 100,
+                    Message = "Progress message",
+                },
+            },
+            configureOptions: null,
+            configureServer: server =>
+            {
+                server.AddNotificationHandler(NotificationMethods.ProgressNotification,
+                    (notification) =>
+                    {
+                        notificationHandled = true;
+                        var progress = (ProgressNotificationValue?)notification.Params;
+                        Assert.NotNull(progress);
+                        var progressValue = progress.Value;
+                        taskCompletionSource.SetResult();
+                        Assert.Equal(50, progressValue.Progress);
+                        Assert.Equal(100, progressValue.Total);
+                        Assert.Equal("Progress message", progressValue.Message);
+                        return Task.CompletedTask;
+                    });
+            },
+            assertResult: async response =>
+            {
+                //Note: awaiting here so handlers are guaranteed to be called first.
+                await taskCompletionSource.Task.WaitAsync(TimeSpan.FromSeconds(1));
+                Assert.True(notificationHandled);
+            });
+    }
+
+    private async Task Notifications_Are_Handled(
+        ServerCapabilities? serverCapabilities,
+        string method, object? parameters,
+        Action<McpServerOptions>? configureOptions,
+        Action<IMcpServer>? configureServer,
+        Action<JsonRpcNotification> assertResult)
+    {
+        await using TestServerTransport transport = new();
+        var options = CreateOptions(serverCapabilities);
+        configureOptions?.Invoke(options);
+
+        await using var server = McpServerFactory.Create(
+            transport, options, LoggerFactory, _serviceProvider);
+
+        configureServer?.Invoke(server);
+        await server.RunAsync();
+
+        TaskCompletionSource<JsonRpcNotification> receivedMessage = new();
+
+        transport.OnMessageSent = (message) =>
+        {
+            Assert.NotNull(message);
+            if (message is JsonRpcNotification notification && notification.Method == method)
+            {
+                assertResult(notification);
+                receivedMessage.SetResult(notification);
+            }
+        };
+
+        await transport.SendMessageAsync(new JsonRpcNotification
+        {
+            Method = method,
+            Params = parameters,
+        });
+
+        var response = await receivedMessage.Task.WaitAsync(TimeSpan.FromSeconds(1));
+    }
 }
diff --git a/tests/ModelContextProtocol.Tests/SseIntegrationTests.cs b/tests/ModelContextProtocol.Tests/SseIntegrationTests.cs
@@ -116,7 +116,7 @@ public async Task Sampling_Sse_EverythingServer()
             {
                 Sampling = new()
                 {
-                    SamplingHandler = (_, _) =>
+                    SamplingHandler = (_, _, _) =>
                     {
                         samplingHandlerCalls++;
                         return Task.FromResult(new CreateMessageResult
diff --git a/tests/ModelContextProtocol.Tests/SseServerIntegrationTests.cs b/tests/ModelContextProtocol.Tests/SseServerIntegrationTests.cs
@@ -215,7 +215,7 @@ public async Task Sampling_Sse_TestServer()
         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

Original file line number	Diff line number	Diff line change
`@@ -42,7 +42,7 @@ public McpClient(IClientTransport clientTransport, McpClientOptions options, Mcp`
`42`	`42`
`43`	`43`	`SetRequestHandler<CreateMessageRequestParams, CreateMessageResult>(`
`44`	`44`	`RequestMethods.SamplingCreateMessage,`
`45`		`- (request, ct) => samplingHandler(request, ct));`
	`45`	`+ (request, ct) => samplingHandler(request, new ClientTokenProgress(this, request?.Meta?.ProgressToken), ct));`
`46`	`46`	`}`
`47`	`47`
`48`	`48`	`if (options.Capabilities?.Roots is { } rootsCapability)`
Original file line number	Diff line number	Diff line change
`@@ -55,7 +55,7 @@ public class SamplingCapability`
`55`	`55`
`56`	`56`	`/// <summary>Gets or sets the handler for sampling requests.</summary>`
`57`	`57`	`[JsonIgnore]`
`58`		`- public Func<CreateMessageRequestParams?, CancellationToken, Task<CreateMessageResult>>? SamplingHandler { get; set; }`
	`58`	`+ public Func<CreateMessageRequestParams?, IProgress<ProgressNotificationValue>, CancellationToken, Task<CreateMessageResult>>? SamplingHandler { get; set; }`
`59`	`59`	`}`
`60`	`60`
`61`	`61`	`/// <summary>`