Extend progress notification support (#163)

* Extend progress notification support

* Address feedback and fix test

---------

Co-authored-by: Stephen Toub <stoub@microsoft.com>
This commit is contained in:
Tyler Kendrick 2025-04-01 10:43:50 -05:00 • committed by GitHub
commit 25bcb440a0
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 175 additions and 154 deletions

View file

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

View file

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

View file

@ -8,9 +8,7 @@ using System.Runtime.CompilerServices;
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 @@ 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);
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();
};
}