modelcontextprotocol/csharp-sdk
Watch
1
0
Fork
You've already forked csharp-sdk
 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 1

Add RegisterTools API to McpClient for pre-populating tool cache (#1590)

Browse source 
Co-authored-by: Tarek Mahmoud Sayed <tarekms@ntdev.microsoft.com>
This commit is contained in:
Tarek Mahmoud Sayed 2026-05-26 13:31:31 -07:00 • committed by GitHub
commit 157f855e51
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 1016 additions and 6 deletions
Download patch file Download diff file Expand all files Collapse all files

80
src/ModelContextProtocol.Core/Client/McpClient.cs
View file

@ -70,4 +70,84 @@ public abstract partial class McpClient : McpSession
/// </para>
/// </remarks>
public abstract Task<ClientCompletionDetails> Completion { get; }
/// <summary>
/// Registers one or more tool definitions in the client's tool cache, enabling the transport
/// to send <c>Mcp-Param-*</c> headers for those tools without requiring a prior <see cref="McpClient.ListToolsAsync(RequestOptions?, CancellationToken)"/> call.
/// </summary>
/// <param name="tools">The tool definitions to register.</param>
/// <remarks>
/// <para>
/// This method allows callers who already have tool schema information (e.g., from a previous session,
/// hardcoded configuration, or an out-of-band source) to provide it directly to the client. Once registered,
/// any <see cref="McpClient.CallToolAsync(string, IReadOnlyDictionary{string, object?}?, IProgress{ProgressNotificationValue}?, RequestOptions?, CancellationToken)"/>
/// call for a registered tool will automatically include <c>Mcp-Param-*</c> HTTP headers based on
/// the tool's <c>x-mcp-header</c> schema annotations, exactly as if the tool had been discovered
/// via <see cref="McpClient.ListToolsAsync(RequestOptions?, CancellationToken)"/>.
/// </para>
/// <para>
/// <b>Cache interaction behavior:</b>
/// <list type="bullet">
/// <item>Registered tools are added to the same internal tool cache used by <see cref="McpClient.ListToolsAsync(RequestOptions?, CancellationToken)"/>.</item>
/// <item>Calling <see cref="McpClient.ListToolsAsync(RequestOptions?, CancellationToken)"/> after <see cref="AddKnownTools"/> preserves
/// manually registered tools — only server-discovered tools are cleared and repopulated.</item>
/// <item>If the server returns a tool with the same name as a manually registered tool, the server's
/// definition overwrites the registered one in the cache, but the tool retains its known status
/// and will survive subsequent cache clears. This registration is sticky for the lifetime of the
/// <see cref="McpClient"/>; use <see cref="RemoveKnownTools"/> or <see cref="ClearKnownTools"/> to
/// explicitly drop known tools that are no longer needed.</item>
/// <item>Tools can be registered at any time — before or after <see cref="McpClient.ListToolsAsync(RequestOptions?, CancellationToken)"/>,
/// and across multiple calls.</item>
/// <item>Re-registering a tool with the same name overwrites the previous definition in the cache (last write wins).</item>
/// </list>
/// </para>
/// <para>
/// Tools with invalid <c>x-mcp-header</c> annotations cause an <see cref="ArgumentException"/> to be thrown.
/// No tools are added to the cache if any tool in the batch fails validation (all-or-nothing).
/// </para>
/// </remarks>
/// <exception cref="ArgumentNullException"><paramref name="tools"/> is <see langword="null"/>.</exception>
/// <exception cref="ArgumentException">One or more tools have invalid <c>x-mcp-header</c> annotations.</exception>
public virtual void AddKnownTools(IEnumerable<Tool> tools)
{
Throw.IfNull(tools);
throw new NotSupportedException($"{GetType().Name} does not support adding known tools.");
}
/// <summary>
/// Removes one or more previously registered tool definitions from the client's tool cache by name.
/// </summary>
/// <param name="toolNames">The names of the tools to remove.</param>
/// <remarks>
/// <para>
/// This removes the specified tools from both the known-tools set and the internal tool cache.
/// After removal, those tools will no longer survive <see cref="McpClient.ListToolsAsync(RequestOptions?, CancellationToken)"/>
/// cache clears, and <c>Mcp-Param-*</c> headers will no longer be sent for them unless the server
/// re-discovers them via <see cref="McpClient.ListToolsAsync(RequestOptions?, CancellationToken)"/>.
/// </para>
/// <para>
/// Removing a tool name that was not previously added via <see cref="AddKnownTools"/> is a no-op.
/// </para>
/// </remarks>
/// <exception cref="ArgumentNullException"><paramref name="toolNames"/> is <see langword="null"/>.</exception>
public virtual void RemoveKnownTools(IEnumerable<string> toolNames)
{
Throw.IfNull(toolNames);
throw new NotSupportedException($"{GetType().Name} does not support removing known tools.");
}
/// <summary>
/// Removes all previously registered tool definitions from the client's tool cache.
/// </summary>
/// <remarks>
/// <para>
/// This clears all tools that were added via <see cref="AddKnownTools"/> from both the known-tools
/// set and the internal tool cache. Server-discovered tools that are not also known tools are not affected
/// and will remain in the cache until the next <see cref="McpClient.ListToolsAsync(RequestOptions?, CancellationToken)"/> call.
/// </para>
/// </remarks>
public virtual void ClearKnownTools()
{
throw new NotSupportedException($"{GetType().Name} does not support clearing known tools.");
}
}

101
src/ModelContextProtocol.Core/Client/McpClientImpl.cs
View file

@ -24,6 +24,7 @@ internal sealed partial class McpClientImpl : McpClient
private readonly SemaphoreSlim _disposeLock = new(1, 1);
private readonly McpTaskCancellationTokenProvider? _taskCancellationTokenProvider;
private readonly ConcurrentDictionary<string, Tool> _toolCache = new(StringComparer.Ordinal);
private readonly ConcurrentDictionary<string, byte> _registeredToolNames = new(StringComparer.Ordinal);
private ServerCapabilities? _serverCapabilities;
private Implementation? _serverInfo;
@ -72,7 +73,23 @@ internal sealed partial class McpClientImpl : McpClient
ToolDiscovered = tool => _toolCache[tool.Name] = tool;
ToolRejected = (tool, reason) => LogToolRejected(tool.Name, reason);
ToolCacheClearing = () => _toolCache.Clear();
ToolCacheClearing = () =>
{
if (_registeredToolNames.IsEmpty)
{
_toolCache.Clear();
return;
}
// Only remove server-discovered tools; preserve manually registered tools.
foreach (var key in _toolCache.Keys)
{
if (!_registeredToolNames.ContainsKey(key))
{
_toolCache.TryRemove(key, out _);
}
}
};
}
private void RegisterHandlers(McpClientOptions options, NotificationHandlers notificationHandlers, RequestHandlers requestHandlers)
@ -637,6 +654,69 @@ internal sealed partial class McpClientImpl : McpClient
LogClientSessionResumed(_endpointName);
}
/// <inheritdoc/>
public override void AddKnownTools(IEnumerable<Tool> tools)
{
Throw.IfNull(tools);
var snapshot = tools as IReadOnlyCollection<Tool> ?? [.. tools];
List<string>? rejections = null;
foreach (var tool in snapshot)
{
Throw.IfNull(tool);
if (!McpHeaderExtractor.ValidateToolSchema(tool, out var rejectionReason))
{
ToolRejected?.Invoke(tool, rejectionReason!);
(rejections ??= []).Add($"{tool.Name}: {rejectionReason}");
}
}
if (rejections is { Count: > 0 })
{
throw new ArgumentException(
"One or more tools failed x-mcp-header validation: " + string.Join("; ", rejections),
nameof(tools));
}
foreach (var tool in snapshot)
{
_registeredToolNames[tool.Name] = 0;
_toolCache[tool.Name] = tool;
}
}
/// <inheritdoc/>
public override void RemoveKnownTools(IEnumerable<string> toolNames)
{
Throw.IfNull(toolNames);
var snapshot = toolNames as IReadOnlyCollection<string> ?? [.. toolNames];
foreach (var name in snapshot)
{
Throw.IfNull(name);
}
foreach (var name in snapshot)
{
_registeredToolNames.TryRemove(name, out _);
_toolCache.TryRemove(name, out _);
}
}
/// <inheritdoc/>
public override void ClearKnownTools()
{
foreach (var name in _registeredToolNames.Keys)
{
_toolCache.TryRemove(name, out _);
}
_registeredToolNames.Clear();
}
/// <inheritdoc/>
public override Task<JsonRpcResponse> SendRequestAsync(JsonRpcRequest request, CancellationToken cancellationToken = default)
{
@ -645,12 +725,18 @@ internal sealed partial class McpClientImpl : McpClient
if (request.Method == RequestMethods.ToolsCall &&
request.Params is System.Text.Json.Nodes.JsonObject paramsObj &&
paramsObj.TryGetPropertyValue("name", out var nameNode) &&
nameNode?.GetValue<string>() is { } toolName &&
_toolCache.TryGetValue(toolName, out var tool))
nameNode?.GetValue<string>() is { } toolName)
{
request.Context ??= new();
request.Context.Items ??= new Dictionary<string, object?>();
request.Context.Items[McpHttpHeaders.ToolContextKey] = tool;
if (_toolCache.TryGetValue(toolName, out var tool))
{
request.Context ??= new();
request.Context.Items ??= new Dictionary<string, object?>();
request.Context.Items[McpHttpHeaders.ToolContextKey] = tool;
}
else if (_transport is StreamableHttpClientSessionTransport)
{
LogToolCacheMiss(toolName);
}
}
return _sessionHandler.SendRequestAsync(request, cancellationToken);
@ -707,6 +793,9 @@ internal sealed partial class McpClientImpl : McpClient
[LoggerMessage(Level = LogLevel.Information, Message = "{EndpointName} client resumed existing session.")]
private partial void LogClientSessionResumed(string endpointName);
[LoggerMessage(Level = LogLevel.Warning, Message = "Tool '{ToolName}' not found in cache during tools/call. Mcp-Param-* headers will not be sent. Call AddKnownTools or ListToolsAsync to populate the cache.")]
private partial void LogToolCacheMiss(string toolName);
[LoggerMessage(Level = LogLevel.Warning, Message = "Tool '{ToolName}' excluded from tools/list: {Reason}")]
private partial void LogToolRejected(string toolName, string reason);