Add UseMcpClient #1086
Add UseMcpClient #1086
Conversation
There was a problem hiding this comment.
The tests are in this project because it wasn't possible for me to mock the McpClient so I had to use integration tests.
There was a problem hiding this comment.
Can you elaborate on the challenges with mocking McpClient? It'd be good to understand and fix blockers.
There was a problem hiding this comment.
Not sure if those are the same challenges, but mocking McpClient used to confuse AI assistants (Claude, GitHub Copilot, etc.) as they couldn't figure it out and did some really odd stuff with some of the common mocking frameworks. I seem to remember coding up an initial example of doing it and that made it click for the model, but I don't know if they still struggle with the mocking.
Apologies, if there's an actual issue with mocking McpClient currently and I hope this doesn't come across in the wrong way. But it caught me off guard more than once, and I recall the mocking actually being a little tricky with 0.3.0-preview4.
There was a problem hiding this comment.
I would also like to understand this. The new code we're trying to test calls McpClient.CreateAsync directly, so the only way to mock it would be to adding indirection there for testing or to mock a lower level like the MCP transport level or HttpClient level.
Since the HttpClient is already pluggable, I recommend creating a new test class based on KestrelInMemoryTest rather than continuing to expand on SseServerIntegrationTestFixture. The SseServerIntegrationTestFixture is mostly the way it is because it's testing the legacy transport and I haven't gotten around to updating it.
KestrelInMemoryTest already has a MockLoggerProvider and XunitLoggerProvider wired up like all LoggedTests, and you can run multiple servers with the same in-memory Kestrel transport as demonstrated by OAuthTestBase. This allows you to define all your server-side and client-side code inside a single test method so we don't end up with a massive test server project with different code scattered about that are relevant to different tests far away in the source code.
stephentoub left a comment
stephentoub left a comment There was a problem hiding this comment.
Nice. Seems like this generally works out?
| /// Use this method as an alternative when working with chat providers that don't have built-in support for hosted MCP servers. | ||
| /// </para> | ||
| /// </remarks> | ||
| public static ChatClientBuilder UseMcpClient( |
There was a problem hiding this comment.
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.
| /// </remarks> | ||
| public static ChatClientBuilder UseMcpClient( | ||
| this ChatClientBuilder builder, | ||
| HttpClient? httpClient = null, |
There was a problem hiding this comment.
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.
There was a problem hiding this comment.
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.
There was a problem hiding this comment.
Could we have a Func<HostedMcpServerTool, HttpClientTransportOptions?>? transportOptionsProvider = null? Someone could return options for specific servers, their own defaults, or null.
| }); | ||
| } | ||
| | ||
| private class McpChatClient : DelegatingChatClient |
There was a problem hiding this comment.
nit:
| private class McpChatClient : DelegatingChatClient | |
| private sealed class McpChatClient : DelegatingChatClient |
| private readonly ILogger _logger; | ||
| private readonly HttpClient _httpClient; | ||
| private readonly bool _ownsHttpClient; | ||
| private ConcurrentDictionary<string, Task<McpClient>>? _mcpClientTasks = null; |
There was a problem hiding this comment.
Presumably this middleware is only being used if there's an expectation that HostedMcpServerTool will be used, in which case making this lazy is probably not worthwhile. I suggest just making it be:
| private ConcurrentDictionary<string, Task<McpClient>>? _mcpClientTasks = null; | |
| private readonly ConcurrentDictionary<string, Task<McpClient>> _mcpClientTasks = []; |
and then you can remove the null checks on it later.
| } | ||
| } | ||
| | ||
| private async Task<List<AITool>?> BuildDownstreamAIToolsAsync(IList<AITool>? inputTools, CancellationToken cancellationToken) |
There was a problem hiding this comment.
inputTools can be non-nullable, right? The call sites both validated it's not null. And then you can remove the subsequent ?? []
| case HostedMcpServerToolAlwaysRequireApprovalMode alwaysRequireApproval: | ||
| downstreamTools.Add(new ApprovalRequiredAIFunction(mcpFunction)); | ||
| break; | ||
| case HostedMcpServerToolNeverRequireApprovalMode neverRequireApproval: | ||
| downstreamTools.Add(mcpFunction); | ||
| break; | ||
| case HostedMcpServerToolRequireSpecificApprovalMode specificApprovalMode when specificApprovalMode.AlwaysRequireApprovalToolNames?.Contains(mcpFunction.Name) is true: | ||
| downstreamTools.Add(new ApprovalRequiredAIFunction(mcpFunction)); | ||
| break; | ||
| 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; |
There was a problem hiding this comment.
| case HostedMcpServerToolAlwaysRequireApprovalMode alwaysRequireApproval: | |
| downstreamTools.Add(new ApprovalRequiredAIFunction(mcpFunction)); | |
| break; | |
| case HostedMcpServerToolNeverRequireApprovalMode neverRequireApproval: | |
| downstreamTools.Add(mcpFunction); | |
| break; | |
| case HostedMcpServerToolRequireSpecificApprovalMode specificApprovalMode when specificApprovalMode.AlwaysRequireApprovalToolNames?.Contains(mcpFunction.Name) is true: | |
| downstreamTools.Add(new ApprovalRequiredAIFunction(mcpFunction)); | |
| break; | |
| 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; | |
| case HostedMcpServerToolNeverRequireApprovalMode neverRequireApproval: | |
| 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; |
| #if NETSTANDARD2_0 | ||
| if (clientTask.Status == TaskStatus.RanToCompletion) | ||
| #else | ||
| if (clientTask.IsCompletedSuccessfully) | ||
| #endif |
There was a problem hiding this comment.
You can just keep it simple and always do the is TaskStatus.RanToCompletion path.
| | ||
| // 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. | ||
| return _mcpClientTasks.GetOrAdd(serverAddress.ToString(), _ => CreateMcpClientCoreAsync(serverAddress, serverName, authorizationToken, CancellationToken.None)); |
There was a problem hiding this comment.
Probably not a huge deal given the other overheads, but this is always going to allocate a closure to capture serverAddress, serverName, and authorizationToken.
| | ||
| private async Task<McpClient> CreateMcpClientCoreAsync(Uri serverAddress, string serverName, string? authorizationToken, CancellationToken cancellationToken) | ||
| { | ||
| var serverAddressKey = serverAddress.ToString(); |
There was a problem hiding this comment.
If you pass in the key that's provided to the GetOrAdd delegate, you won't need to recompute this
| catch | ||
| { | ||
| // Remove the failed task from cache so subsequent requests can retry | ||
| _mcpClientTasks?.TryRemove(serverAddressKey, out _); |
There was a problem hiding this comment.
It may not have been added to the cache yet, if all the awaits until this point are for things that completed synchronously, or given other race conditions.
| This is a really awesome extension. ❤ |
4 participants
Contributes to dotnet/extensions#6492
cc @westey-m