Address feedback and fix test

Author: stephentoub

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, new ClientTokenProgress(this, request?.Meta?.ProgressToken), 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>
@@ -542,18 +540,17 @@ internal static CreateMessageResult ToCreateMessageResult(this ChatResponse chat
 
             var (messages, options) = requestParams.ToChatClientArguments();
             var progressToken = requestParams.Meta?.ProgressToken;
-            int progressValue = 0;
-            var streamingResponses = chatClient.GetStreamingResponseAsync(
-                messages, options, cancellationToken);
+
             List<ChatResponseUpdate> updates = [];
-            await foreach (var streamingResponse in streamingResponses)
+            await foreach (var update in chatClient.GetStreamingResponseAsync(messages, options, cancellationToken))
             {
-                updates.Add(streamingResponse);
+                updates.Add(update);
+
                 if (progressToken is not null)
                 {
                     progress.Report(new()
                     {
-                        Progress = ++progressValue,
+                        Progress = updates.Count,
                     });
                 }
             }

src/ModelContextProtocol/ClientTokenProgress.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/IMcpEndpoint.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/McpEndpointExtensions.cs

@@ -38,7 +38,7 @@ public JsonElement InputSchema
         {
             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/Protocol/Types/Tool.cs

@@ -111,7 +111,7 @@ private static TemporaryAIFunctionFactoryOptions CreateAIFunctionFactoryOptions(
                             if (requestContent?.Server is { } server &&
                                 requestContent?.Params?.Meta?.ProgressToken is { } progressToken)
                             {
-                                return new ServerTokenProgress(server, progressToken);
+                                return new TokenProgress(server, progressToken);
                             }
 
                             return NullProgress.Instance;

src/ModelContextProtocol/Server/AIFunctionMcpServerTool.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/IMcpServer.cs

@@ -7,7 +7,7 @@
 
 namespace ModelContextProtocol.Server;
 
-/// <inheritdoc />
+/// <summary>Provides extension methods for interacting with an <see cref="IMcpServer"/>.</summary>
 public static class McpServerExtensions
 {
     /// <summary>
@@ -169,26 +169,6 @@ 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

src/ModelContextProtocol/Server/McpServerExtensions.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()
-        {
-            Method = NotificationMethods.ProgressNotification,
-            Params = new ProgressNotification()
-            {
-                ProgressToken = progressToken,
-                Progress = new()
-                {
-                    Progress = value.Progress,
-                    Total = value.Total,
-                    Message = value.Message,
-                },
-            },
-        }, CancellationToken.None);
+        _ = endpoint.NotifyProgressAsync(progressToken, value, CancellationToken.None);
     }
 }