Extend progress notification support (#163)

* Extend progress notification support

* Address feedback and fix test

---------

Co-authored-by: Stephen Toub <[email protected]>

Author: Tyler-R-Kendrick

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 @@ public McpClient(IClientTransport clientTransport, McpClientOptions options, Mcp
 
             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 @@ public McpClient(IClientTransport clientTransport, McpClientOptions options, Mcp
 
             SetRequestHandler<ListRootsRequestParams, ListRootsResult>(
                 RequestMethods.RootsList,
-                (request, ct) => rootsHandler(request, ct));
+                (request, cancellationToken) => rootsHandler(request, cancellationToken));
         }
     }
 

src/ModelContextProtocol/Client/McpClientExtensions.cs

@@ -8,9 +8,7 @@
 
 namespace ModelContextProtocol.Client;
 
-/// <summary>
-/// Provides extensions for operating on MCP clients.
-/// </summary>
+/// <summary>Provides extension methods for interacting with an <see cref="IMcpClient"/>.</summary>
 public static class McpClientExtensions
 {
     /// <summary>
@@ -531,17 +529,33 @@ 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;
+
+            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
-{
-    /// <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;

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;

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;

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;