Extend progress notification support (#163)

* Extend progress notification support

* Address feedback and fix test

---------

Co-authored-by: Stephen Toub <stoub@microsoft.com>

src/ModelContextProtocol/Client/IMcpClient.cs

@@ -1,12 +1,11 @@
-using ModelContextProtocol.Protocol.Messages;
+using ModelContextProtocol.Protocol.Types;
-using ModelContextProtocol.Protocol.Types;
 
 namespace ModelContextProtocol.Client;
 
 /// <summary>
 /// Represents an instance of an MCP client connecting to a specific server.
 /// </summary>
-public interface IMcpClient : IAsyncDisposable
+public interface IMcpClient : IMcpEndpoint
 {
     /// <summary>
     /// Gets the capabilities supported by the server.
@@ -24,40 +23,4 @@ public interface IMcpClient : IAsyncDisposable
     /// It can be thought of like a "hint" to the model. For example, this information MAY be added to the system prompt.
     /// </summary>
     string? ServerInstructions { get; }
-
-    /// <summary>
-    /// Adds a handler for server notifications of a specific method.
-    /// </summary>
-    /// <param name="method">The notification method to handle.</param>
-    /// <param name="handler">The async handler function to process notifications.</param>
-    /// <remarks>
-    /// <para>
-    /// Each method may have multiple handlers. Adding a handler for a method that already has one
-    /// will not replace the existing handler.
-    /// </para>
-    /// <para>
-    /// <see cref="NotificationMethods"> provides constants for common notification methods.</see>
-    /// </para>
-    /// </remarks>
-    void AddNotificationHandler(string method, Func<JsonRpcNotification, Task> handler);
-
-    /// <summary>
-    /// Sends a generic JSON-RPC request to the server.
-    /// </summary>
-    /// <typeparam name="TResult">The expected response type.</typeparam>
-    /// <param name="request">The JSON-RPC request to send.</param>
-    /// <param name="cancellationToken">A token to cancel the operation.</param>
-    /// <returns>A task containing the server's response.</returns>
-    /// <remarks>
-    /// It is recommended to use the capability-specific methods that use this one in their implementation.
-    /// Use this method for custom requests or those not yet covered explicitly.
-    /// </remarks>
-    Task<TResult> SendRequestAsync<TResult>(JsonRpcRequest request, CancellationToken cancellationToken = default) where TResult : class;
-
-    /// <summary>
-    /// Sends a message to the server.
-    /// </summary>
-    /// <param name="message">The message.</param>
-    /// <param name="cancellationToken">A token to cancel the operation.</param>
-    Task SendMessageAsync(IJsonRpcMessage message, CancellationToken cancellationToken = default);
 }

src/ModelContextProtocol/Client/McpClient.cs

@@ -42,7 +42,10 @@ internal sealed class McpClient : McpJsonRpcEndpoint, IMcpClient
 
             SetRequestHandler<CreateMessageRequestParams, CreateMessageResult>(
                 RequestMethods.SamplingCreateMessage,
-                (request, ct) => samplingHandler(request, ct));
+                (request, cancellationToken) => samplingHandler(
+                    request,
+                    request?.Meta?.ProgressToken is { } token ? new TokenProgress(this, token) : NullProgress.Instance,
+                    cancellationToken));
         }
 
         if (options.Capabilities?.Roots is { } rootsCapability)
@@ -54,7 +57,7 @@ internal sealed class McpClient : McpJsonRpcEndpoint, IMcpClient
 
             SetRequestHandler<ListRootsRequestParams, ListRootsResult>(
                 RequestMethods.RootsList,
-                (request, ct) => rootsHandler(request, ct));
+                (request, cancellationToken) => rootsHandler(request, cancellationToken));
         }
     }
 

src/ModelContextProtocol/Client/McpClientExtensions.cs

@@ -8,9 +8,7 @@ using System.Runtime.CompilerServices;
 
 namespace ModelContextProtocol.Client;
 
-/// <summary>
+/// <summary>Provides extension methods for interacting with an <see cref="IMcpClient"/>.</summary>
-/// Provides extensions for operating on MCP clients.
-/// </summary>
 public static class McpClientExtensions
 {
     /// <summary>
@@ -531,17 +529,33 @@ public static class McpClientExtensions
     /// </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);
+            var progressToken = requestParams.Meta?.ProgressToken;
-            return response.ToCreateMessageResult();
+
+            List<ChatResponseUpdate> updates = [];
+            await foreach (var update in chatClient.GetStreamingResponseAsync(messages, options, cancellationToken))
+            {
+                updates.Add(update);
+
+                if (progressToken is not null)
+                {
+                    progress.Report(new()
+                    {
+                        Progress = updates.Count,
+                    });
+                }
+            }
+
+            return updates.ToChatResponse().ToCreateMessageResult();
         };
     }
 

src/ModelContextProtocol/IMcpEndpoint.cs

@@ -0,0 +1,35 @@
+using ModelContextProtocol.Protocol.Messages;
+
+namespace ModelContextProtocol;
+
+/// <summary>Represents a client or server MCP endpoint.</summary>
+public interface IMcpEndpoint : IAsyncDisposable
+{
+    /// <summary>Sends a generic JSON-RPC request to the connected endpoint.</summary>
+    /// <typeparam name="TResult">The expected response type.</typeparam>
+    /// <param name="request">The JSON-RPC request to send.</param>
+    /// <param name="cancellationToken">A token to cancel the operation.</param>
+    /// <returns>A task containing the client's response.</returns>
+    Task<TResult> SendRequestAsync<TResult>(JsonRpcRequest request, CancellationToken cancellationToken = default) where TResult : class;
+
+    /// <summary>Sends a message to the connected endpoint.</summary>
+    /// <param name="message">The message.</param>
+    /// <param name="cancellationToken">A token to cancel the operation.</param>
+    Task SendMessageAsync(IJsonRpcMessage message, CancellationToken cancellationToken = default);
+
+    /// <summary>
+    /// Adds a handler for server notifications of a specific method.
+    /// </summary>
+    /// <param name="method">The notification method to handle.</param>
+    /// <param name="handler">The async handler function to process notifications.</param>
+    /// <remarks>
+    /// <para>
+    /// Each method may have multiple handlers. Adding a handler for a method that already has one
+    /// will not replace the existing handler.
+    /// </para>
+    /// <para>
+    /// <see cref="NotificationMethods"> provides constants for common notification methods.</see>
+    /// </para>
+    /// </remarks>
+    void AddNotificationHandler(string method, Func<JsonRpcNotification, Task> handler);
+}

src/ModelContextProtocol/McpEndpointExtensions.cs

@@ -0,0 +1,34 @@
+using ModelContextProtocol.Protocol.Messages;
+using ModelContextProtocol.Utils;
+
+namespace ModelContextProtocol;
+
+/// <summary>Provides extension methods for interacting with an <see cref="IMcpEndpoint"/>.</summary>
+public static class McpEndpointExtensions
+{
+    /// <summary>Notifies the connected endpoint of progress.</summary>
+    /// <param name="endpoint">The endpoint issueing the notification.</param>
+    /// <param name="progressToken">The <see cref="ProgressToken"/> identifying the operation.</param>
+    /// <param name="progress">The progress update to send.</param>
+    /// <param name="cancellationToken">A token to cancel the operation.</param>
+    /// <returns>A task representing the completion of the operation.</returns>
+    /// <exception cref="ArgumentNullException"><paramref name="endpoint"/> is <see langword="null"/>.</exception>
+    public static Task NotifyProgressAsync(
+        this IMcpEndpoint endpoint,
+        ProgressToken progressToken,
+        ProgressNotificationValue progress, 
+        CancellationToken cancellationToken = default)
+    {
+        Throw.IfNull(endpoint);
+
+        return endpoint.SendMessageAsync(new JsonRpcNotification()
+        {
+            Method = NotificationMethods.ProgressNotification,
+            Params = new ProgressNotification()
+            {
+                ProgressToken = progressToken,
+                Progress = progress,
+            },
+        }, cancellationToken);
+    }
+}

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>

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
+public class ListPromptsRequestParams : PaginatedRequestParams;
-{
-    /// <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; }
-}

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
+public class ListResourceTemplatesRequestParams : PaginatedRequestParams;
-{
-    /// <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; }
-}

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
+public class ListResourcesRequestParams : PaginatedRequestParams;
-{
-    /// <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; }
-}

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
+public class ListRootsRequestParams : RequestParams;
-{
-    /// <summary>
-    /// Optional progress token for out-of-band progress notifications.
-    /// </summary>
-    [System.Text.Json.Serialization.JsonPropertyName("progressToken")]
-    public ProgressToken? ProgressToken { get; init; }
-}

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
+public class ListToolsRequestParams : PaginatedRequestParams;
-{
-    /// <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; }
-}

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; }
+}

src/ModelContextProtocol/Protocol/Types/Tool.cs

@@ -38,7 +38,7 @@ public class Tool
         {
             if (!McpJsonUtilities.IsValidMcpToolSchema(value))
             {
-                throw new ArgumentException("The specified document is not a valid MPC tool JSON schema.", nameof(InputSchema));
+                throw new ArgumentException("The specified document is not a valid MCP tool JSON schema.", nameof(InputSchema));
             }
 
             _inputSchema = value;

src/ModelContextProtocol/Server/IMcpServer.cs

@@ -1,12 +1,11 @@
-using ModelContextProtocol.Protocol.Messages;
+using ModelContextProtocol.Protocol.Types;
-using ModelContextProtocol.Protocol.Types;
 
 namespace ModelContextProtocol.Server;
 
 /// <summary>
 /// Represents a server that can communicate with a client using the MCP protocol.
 /// </summary>
-public interface IMcpServer : IAsyncDisposable
+public interface IMcpServer : IMcpEndpoint
 {
     /// <summary>
     /// Gets the capabilities supported by the client.
@@ -26,42 +25,8 @@ public interface IMcpServer : IAsyncDisposable
     /// </summary>
     IServiceProvider? Services { get; }
 
-    /// <summary>
-    /// Adds a handler for client notifications of a specific method.
-    /// </summary>
-    /// <param name="method">The notification method to handle.</param>
-    /// <param name="handler">The async handler function to process notifications.</param>
-    /// <remarks>
-    /// <para>
-    /// Each method may have multiple handlers. Adding a handler for a method that already has one
-    /// will not replace the existing handler.
-    /// </para>
-    /// <para>
-    /// <see cref="NotificationMethods"> provides constants for common notification methods.</see>
-    /// </para>
-    /// </remarks>
-    void AddNotificationHandler(string method, Func<JsonRpcNotification, Task> handler);
-
     /// <summary>
     /// Runs the server, listening for and handling client requests.
     /// </summary>
     Task RunAsync(CancellationToken cancellationToken = default);
-
-    /// <summary>
-    /// Sends a generic JSON-RPC request to the client.
-    /// NB! This is a temporary method that is available to send not yet implemented feature messages.
-    /// Once all MCP features are implemented this will be made private, as it is purely a convenience for those who wish to implement features ahead of the library.
-    /// </summary>
-    /// <typeparam name="TResult">The expected response type.</typeparam>
-    /// <param name="request">The JSON-RPC request to send.</param>
-    /// <param name="cancellationToken">A token to cancel the operation.</param>
-    /// <returns>A task containing the client's response.</returns>
-    Task<TResult> SendRequestAsync<TResult>(JsonRpcRequest request, CancellationToken cancellationToken = default) where TResult : class;
-
-    /// <summary>
-    /// Sends a message to the client.
-    /// </summary>
-    /// <param name="message">The message.</param>
-    /// <param name="cancellationToken">A token to cancel the operation.</param>
-    Task SendMessageAsync(IJsonRpcMessage message, CancellationToken cancellationToken = default);
 }

src/ModelContextProtocol/Server/McpServer.cs

@@ -68,7 +68,6 @@ internal sealed class McpServer : McpJsonRpcEndpoint, IMcpServer
         });
 
         SetToolsHandler(options);
-
         SetInitializeHandler(options);
         SetCompletionHandler(options);
         SetPingHandler();

src/ModelContextProtocol/Server/McpServerExtensions.cs

@@ -7,7 +7,7 @@ using System.Text;
 
 namespace ModelContextProtocol.Server;
 
-/// <inheritdoc />
+/// <summary>Provides extension methods for interacting with an <see cref="IMcpServer"/>.</summary>
 public static class McpServerExtensions
 {
     /// <summary>

src/ModelContextProtocol/TokenProgress.cs

@@ -1,30 +1,16 @@
 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.
+/// progress notifications on the supplied endpoint.
 /// </summary>
-internal sealed class TokenProgress(IMcpServer server, ProgressToken progressToken) : IProgress<ProgressNotificationValue>
+internal sealed class TokenProgress(IMcpEndpoint endpoint, ProgressToken progressToken) : IProgress<ProgressNotificationValue>
 {
     /// <inheritdoc />
     public void Report(ProgressNotificationValue value)
     {
-        _ = server.SendMessageAsync(new JsonRpcNotification()
+        _ = endpoint.NotifyProgressAsync(progressToken, value, CancellationToken.None);
-        {
-            Method = NotificationMethods.ProgressNotification,
-            Params = new ProgressNotification()
-            {
-                ProgressToken = progressToken,
-                Progress = new()
-                {
-                    Progress = value.Progress,
-                    Total = value.Total,
-                    Message = value.Message,
-                },
-            },
-        }, CancellationToken.None);
     }
 }

tests/ModelContextProtocol.Tests/ClientIntegrationTests.cs

@@ -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);

tests/ModelContextProtocol.Tests/Server/McpServerTests.cs

@@ -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;
+    }
 }

tests/ModelContextProtocol.Tests/SseIntegrationTests.cs

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