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

Co-authored-by: Tarek Mahmoud Sayed <tarekms@ntdev.microsoft.com>

Author: tarekgh

docs/concepts/tools/tools.md

@@ -340,3 +340,50 @@ Rules and constraints:
 - Values containing non-ASCII characters, control characters, or leading/trailing whitespace are Base64-encoded using the `=?base64?{value}?=` wrapper.
 - Header names must be case-insensitively unique within the tool's input schema.
 - Header validation is enforced only for protocol versions that support the HTTP Standardization feature (currently `DRAFT-2026-v1` and later).
+
+### Pre-loading tool definitions on the client
+
+By default, `Mcp-Param-*` headers are sent only for tools discovered via <xref:ModelContextProtocol.Client.McpClient.ListToolsAsync*>. If a client already has tool schema information (for example, from a previous session, hardcoded configuration, or an out-of-band source), it can pre-load those definitions so that headers are sent immediately—without a round trip to the server.
+
+```csharp
+// Build the tool definition with x-mcp-header annotations
+var tool = new Tool
+{
+    Name = "execute_sql",
+    InputSchema = JsonDocument.Parse("""
+        {
+            "type": "object",
+            "properties": {
+                "region": {
+                    "type": "string",
+                    "x-mcp-header": "Region"
+                },
+                "query": {
+                    "type": "string"
+                }
+            }
+        }
+        """).RootElement.Clone(),
+};
+
+// Pre-load the tool definition — no ListToolsAsync needed
+client.AddKnownTools([tool]);
+
+// This call now sends an Mcp-Param-Region header automatically
+var result = await client.CallToolAsync("execute_sql",
+    new Dictionary<string, object?> { ["region"] = "us-west-2", ["query"] = "SELECT 1" });
+```
+
+Known tools survive <xref:ModelContextProtocol.Client.McpClient.ListToolsAsync*> cache clears—they remain in the cache even when the server's tool list is refreshed. If the server returns a tool with the same name, the server's definition overwrites the cached one, but the tool keeps its known status.
+
+To remove known tools, use <xref:ModelContextProtocol.Client.McpClient.RemoveKnownTools*> for specific tools or <xref:ModelContextProtocol.Client.McpClient.ClearKnownTools*> to remove all:
+
+```csharp
+// Remove specific known tools by name
+client.RemoveKnownTools(["execute_sql"]);
+
+// Or remove all known tools at once
+client.ClearKnownTools();
+```
+
+All tools passed to <xref:ModelContextProtocol.Client.McpClient.AddKnownTools*> are validated for correct `x-mcp-header` annotations. If any tool in the batch fails validation, an <xref:System.ArgumentException> is thrown and no tools are added (all-or-nothing).

src/ModelContextProtocol.Core/Client/McpClient.cs

@@ -70,4 +70,84 @@ protected McpClient()
     /// </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.");
+    }
 }

src/ModelContextProtocol.Core/Client/McpClientImpl.cs

@@ -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 McpClientImpl(ITransport transport, string endpointName, McpClientOptio
 
         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 void ResumeSession(ResumeClientSessionOptions resumeOptions)
         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 @@ public override Task<JsonRpcResponse> SendRequestAsync(JsonRpcRequest request, C
         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 @@ public override async ValueTask DisposeAsync()
     [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);