-
Notifications
You must be signed in to change notification settings - Fork 587
Add UseMcpClient #1086
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Add UseMcpClient #1086
Changes from all commits
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,225 @@ | ||
| using System.Collections.Concurrent; | ||
| using System.Diagnostics.CodeAnalysis; | ||
| using System.Runtime.CompilerServices; | ||
| using Microsoft.Extensions.AI; | ||
| using Microsoft.Extensions.Logging; | ||
| using Microsoft.Extensions.Logging.Abstractions; | ||
| using ModelContextProtocol.Client; | ||
|
|
||
| namespace ModelContextProtocol; | ||
|
|
||
| /// <summary> | ||
| /// Extension methods for adding MCP client support to chat clients. | ||
| /// </summary> | ||
| public static class McpChatClientBuilderExtensions | ||
| { | ||
| /// <summary> | ||
| /// Adds a chat client to the chat client pipeline that creates an <see cref="McpClient"/> for each <see cref="HostedMcpServerTool"/> | ||
| /// in <see cref="ChatOptions.Tools"/> and augments it with the tools from MCP servers as <see cref="AIFunction"/> instances. | ||
| /// </summary> | ||
| /// <param name="builder">The <see cref="ChatClientBuilder"/> to configure.</param> | ||
| /// <param name="httpClient">The <see cref="HttpClient"/> to use, or <see langword="null"/> to create a new instance.</param> | ||
| /// <param name="loggerFactory">The <see cref="ILoggerFactory"/> to use, or <see langword="null"/> to resolve from services.</param> | ||
| /// <returns>The <see cref="ChatClientBuilder"/> for method chaining.</returns> | ||
| /// <remarks> | ||
| /// <para> | ||
| /// When a <c>HostedMcpServerTool</c> is encountered in the tools collection, the client | ||
| /// connects to the MCP server, retrieves available tools, and expands them into callable AI functions. | ||
| /// Connections are cached by server address to avoid redundant connections. | ||
| /// </para> | ||
| /// <para> | ||
| /// Use this method as an alternative when working with chat providers that don't have built-in support for hosted MCP servers. | ||
| /// </para> | ||
| /// </remarks> | ||
| [Experimental("MEAI001")] | ||
| public static ChatClientBuilder UseMcpClient( | ||
| this ChatClientBuilder builder, | ||
| HttpClient? httpClient = null, | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. It's interesting to consider what should be accepted here and whether that has any impact on the shape of the options we expose in the MCP library. We can't just accept an HttpClientTransport, as that's tied to a specific endpoint, but otherwise we effectively would want most of the options exposed on that, I think.
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Agreed. You wouldn't want to share the ITokenCache, but the remaining ClientOAuthOptions could make a lot of sense. Maybe we should move the TokenCache out of ClientOAuthOptions so we could flow the rest of it through. Headers are a bit in-between. I could see wanting to be able to configure a mapping for known MCP server URL prefixes for things like API keys, but that becomes a lot more of a complex API. At that point maybe we'd recommend putting custom logic in a DelegatingHandler.
Collaborator
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Could we have a
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Or maybe |
||
| ILoggerFactory? loggerFactory = null) | ||
| { | ||
| return builder.Use((innerClient, services) => | ||
| { | ||
| loggerFactory ??= (ILoggerFactory)services.GetService(typeof(ILoggerFactory))!; | ||
| var chatClient = new McpChatClient(innerClient, httpClient, loggerFactory); | ||
| return chatClient; | ||
| }); | ||
| } | ||
|
|
||
| [Experimental("MEAI001")] | ||
|
Member
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Do we need to mark private types as experimental?
Collaborator
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. No, I added it to not have to suppress usage of the experimental API with I will have to switch the ID to MCPXXX in UseMcpClient so I will have to switch the source suppression here anyway. |
||
| private sealed class McpChatClient : DelegatingChatClient | ||
| { | ||
| private readonly ILoggerFactory? _loggerFactory; | ||
| private readonly ILogger _logger; | ||
| private readonly HttpClient _httpClient; | ||
| private readonly bool _ownsHttpClient; | ||
| private readonly ConcurrentDictionary<string, Task<McpClient>> _mcpClientTasks = []; | ||
|
|
||
| /// <summary> | ||
| /// Initializes a new instance of the <see cref="McpChatClient"/> class. | ||
| /// </summary> | ||
| /// <param name="innerClient">The underlying <see cref="IChatClient"/>, or the next instance in a chain of clients.</param> | ||
| /// <param name="httpClient">An optional <see cref="HttpClient"/> to use when connecting to MCP servers. If not provided, a new instance will be created.</param> | ||
| /// <param name="loggerFactory">An <see cref="ILoggerFactory"/> to use for logging information about function invocation.</param> | ||
| public McpChatClient(IChatClient innerClient, HttpClient? httpClient = null, ILoggerFactory? loggerFactory = null) | ||
| : base(innerClient) | ||
| { | ||
| _loggerFactory = loggerFactory; | ||
| _logger = (ILogger?)loggerFactory?.CreateLogger<McpChatClient>() ?? NullLogger.Instance; | ||
| _httpClient = httpClient ?? new HttpClient(); | ||
| _ownsHttpClient = httpClient is null; | ||
| } | ||
|
|
||
| /// <inheritdoc/> | ||
| public override async Task<ChatResponse> GetResponseAsync( | ||
| IEnumerable<ChatMessage> messages, ChatOptions? options = null, CancellationToken cancellationToken = default) | ||
| { | ||
| if (options?.Tools is { Count: > 0 }) | ||
| { | ||
| var downstreamTools = await BuildDownstreamAIToolsAsync(options.Tools, cancellationToken).ConfigureAwait(false); | ||
| options = options.Clone(); | ||
| options.Tools = downstreamTools; | ||
| } | ||
|
|
||
| return await base.GetResponseAsync(messages, options, cancellationToken).ConfigureAwait(false); | ||
| } | ||
|
|
||
| /// <inheritdoc/> | ||
| public override async IAsyncEnumerable<ChatResponseUpdate> GetStreamingResponseAsync(IEnumerable<ChatMessage> messages, ChatOptions? options = null, [EnumeratorCancellation] CancellationToken cancellationToken = default) | ||
| { | ||
| if (options?.Tools is { Count: > 0 }) | ||
| { | ||
| var downstreamTools = await BuildDownstreamAIToolsAsync(options.Tools, cancellationToken).ConfigureAwait(false); | ||
| options = options.Clone(); | ||
| options.Tools = downstreamTools; | ||
| } | ||
|
|
||
| await foreach (var update in base.GetStreamingResponseAsync(messages, options, cancellationToken).ConfigureAwait(false)) | ||
| { | ||
| yield return update; | ||
| } | ||
| } | ||
|
|
||
| private async Task<List<AITool>> BuildDownstreamAIToolsAsync(IList<AITool> inputTools, CancellationToken cancellationToken) | ||
| { | ||
| List<AITool> downstreamTools = []; | ||
| foreach (var tool in inputTools) | ||
| { | ||
| if (tool is not HostedMcpServerTool mcpTool) | ||
| { | ||
| // For other tools, we want to keep them in the list of tools. | ||
| downstreamTools.Add(tool); | ||
| continue; | ||
| } | ||
|
|
||
| if (!Uri.TryCreate(mcpTool.ServerAddress, UriKind.Absolute, out var parsedAddress) || | ||
| (parsedAddress.Scheme != Uri.UriSchemeHttp && parsedAddress.Scheme != Uri.UriSchemeHttps)) | ||
| { | ||
| throw new InvalidOperationException( | ||
| $"Invalid http(s) address: '{mcpTool.ServerAddress}'. MCP server address must be an absolute https(s) URL."); | ||
| } | ||
|
|
||
| // List all MCP functions from the specified MCP server. | ||
| var mcpClient = await CreateMcpClientAsync(mcpTool.ServerAddress, parsedAddress, mcpTool.ServerName, mcpTool.AuthorizationToken).ConfigureAwait(false); | ||
| var mcpFunctions = await mcpClient.ListToolsAsync(cancellationToken: cancellationToken).ConfigureAwait(false); | ||
|
|
||
| // Add the listed functions to our list of tools we'll pass to the inner client. | ||
|
Member
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. How does this approach account for the list of tools potentially changing after this initialization step?
Collaborator
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. It doesn't. We are yet to define lifespan of listed tools and how to handle protocol errors on invoked tools e.g. tool no longer exists. |
||
| foreach (var mcpFunction in mcpFunctions) | ||
| { | ||
| if (mcpTool.AllowedTools is not null && !mcpTool.AllowedTools.Contains(mcpFunction.Name)) | ||
| { | ||
| if (_logger.IsEnabled(LogLevel.Information)) | ||
| { | ||
| _logger.LogInformation("MCP function '{FunctionName}' is not allowed by the tool configuration.", mcpFunction.Name); | ||
| } | ||
| continue; | ||
| } | ||
|
|
||
| switch (mcpTool.ApprovalMode) | ||
| { | ||
| case HostedMcpServerToolNeverRequireApprovalMode: | ||
| case HostedMcpServerToolRequireSpecificApprovalMode specificApprovalMode when specificApprovalMode.NeverRequireApprovalToolNames?.Contains(mcpFunction.Name) is true: | ||
| downstreamTools.Add(mcpFunction); | ||
| break; | ||
|
|
||
| default: | ||
| // Default to always require approval if no specific mode is set. | ||
| downstreamTools.Add(new ApprovalRequiredAIFunction(mcpFunction)); | ||
| break; | ||
| } | ||
| } | ||
| } | ||
|
|
||
| return downstreamTools; | ||
| } | ||
|
|
||
| /// <inheritdoc/> | ||
| protected override void Dispose(bool disposing) | ||
| { | ||
| if (disposing) | ||
| { | ||
| // Dispose of the HTTP client if it was created by this client. | ||
| if (_ownsHttpClient) | ||
| { | ||
| _httpClient?.Dispose(); | ||
| } | ||
|
|
||
| if (_mcpClientTasks is not null) | ||
| { | ||
| // Dispose of all cached MCP clients. | ||
| foreach (var clientTask in _mcpClientTasks.Values) | ||
| { | ||
| if (clientTask.Status == TaskStatus.RanToCompletion) | ||
| { | ||
| _ = clientTask.Result.DisposeAsync(); | ||
| } | ||
| } | ||
|
|
||
| _mcpClientTasks.Clear(); | ||
| } | ||
| } | ||
|
|
||
| base.Dispose(disposing); | ||
| } | ||
|
|
||
| private async Task<McpClient> CreateMcpClientAsync(string key, Uri serverAddress, string serverName, string? authorizationToken) | ||
| { | ||
| // Note: We don't pass cancellationToken to the factory because the cached task should not be tied to any single caller's cancellation token. | ||
| // Instead, callers can cancel waiting for the task, but the connection attempt itself will complete independently. | ||
| #if NET | ||
| // Avoid closure allocation. | ||
| Task<McpClient> task = _mcpClientTasks.GetOrAdd(key, | ||
| static (_, state) => state.self.CreateMcpClientCoreAsync(state.serverAddress, state.serverName, state.authorizationToken, CancellationToken.None), | ||
| (self: this, serverAddress, serverName, authorizationToken)); | ||
| #else | ||
| Task<McpClient> task = _mcpClientTasks.GetOrAdd(key, | ||
| _ => CreateMcpClientCoreAsync(serverAddress, serverName, authorizationToken, CancellationToken.None)); | ||
| #endif | ||
|
|
||
| try | ||
| { | ||
| return await task.ConfigureAwait(false); | ||
| } | ||
| catch | ||
| { | ||
| // Remove the failed task from cache so subsequent requests can retry. | ||
| _mcpClientTasks.TryRemove(key, out _); | ||
| throw; | ||
| } | ||
| } | ||
|
|
||
| private Task<McpClient> CreateMcpClientCoreAsync(Uri serverAddress, string serverName, string? authorizationToken, CancellationToken cancellationToken) | ||
| { | ||
| var transport = new HttpClientTransport(new HttpClientTransportOptions | ||
| { | ||
| Endpoint = serverAddress, | ||
| Name = serverName, | ||
| AdditionalHeaders = authorizationToken is not null | ||
| // Update to pass all headers once https://github.com/dotnet/extensions/pull/7053 is available. | ||
| ? new Dictionary<string, string>() { { "Authorization", $"Bearer {authorizationToken}" } } | ||
| : null, | ||
| }, _httpClient, _loggerFactory); | ||
|
|
||
| return McpClient.CreateAsync(transport, cancellationToken: cancellationToken); | ||
| } | ||
| } | ||
| } | ||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
We should mark the public surface are for this as
[Experimental]. We could do so either with "MEAI001", so that it matches with the MCP types on which its based (and in which case you wouldn't need the earlier suppression) or have some new "MCP003" or something.There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
It would be nice to align this with the ID we end up using for McpServer members. cc @jeffhandley.