Merge branch 'progress-tokens' of https://github.com/Tyler-R-Kendrick/mcp-csharp-sdk into progress-tokens

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

@@ -46,7 +46,7 @@ public McpClient(IClientTransport clientTransport, McpClientOptions options, Mcp
                 {
                     var progressToken = request?.Meta?.ProgressToken;
                     return samplingHandler(
-                        request, progressToken is not null
+                        request, progressToken is { } token
                         ? new ClientTokenProgress(this, progressToken.Value)
                         : NullProgress.Instance, ct);
                 });
@@ -61,7 +61,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

@@ -3,25 +3,15 @@
 
 namespace ModelContextProtocol;
 
-internal sealed class ClientTokenProgress(IMcpClient client, ProgressToken progressToken)
-    : IProgress<ProgressNotificationValue>
+/// <summary>
+/// Provides an <see cref="IProgress{ProgressNotificationValue}"/> tied to a specific progress token and that will issue
+/// progress notifications on the supplied endpoint.
+/// </summary>
+internal sealed class TokenProgress(IMcpEndpoint endpoint, ProgressToken progressToken) : IProgress<ProgressNotificationValue>
 {
     /// <inheritdoc />
     public void Report(ProgressNotificationValue value)
     {
-        _ = client.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);
     }
 }

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/Tool.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/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/McpServerExtensions.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>

tests/ModelContextProtocol.Tests/Server/McpServerTests.cs

@@ -639,81 +639,43 @@ public Task RunAsync(CancellationToken cancellationToken = default) =>
     [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()
+        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(),
+                ProgressToken = new("abc"),
                 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);
-            }
-        };
+        }, TestContext.Current.CancellationToken);
 
-        await transport.SendMessageAsync(new JsonRpcNotification
-        {
-            Method = method,
-            Params = parameters,
-        });
+        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);
 
-        var response = await receivedMessage.Task.WaitAsync(TimeSpan.FromSeconds(1));
+        await server.DisposeAsync();
+        await serverTask;
     }
 }