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

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
5 changed files with 1016 additions and 6 deletions

View file

@ -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).

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.");
}
}

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

View file

@ -0,0 +1,364 @@
using Microsoft.AspNetCore.Builder;
using Microsoft.AspNetCore.Http;
using Microsoft.AspNetCore.Http.Json;
using Microsoft.Extensions.DependencyInjection;
using ModelContextProtocol.AspNetCore.Tests.Utils;
using ModelContextProtocol.Client;
using ModelContextProtocol.Protocol;
using ModelContextProtocol.Tests.Utils;
using System.Collections.Concurrent;
using System.Text.Json;
namespace ModelContextProtocol.AspNetCore.Tests;
/// <summary>
/// Tests that <see cref="McpClient.AddKnownTools"/> allows sending Mcp-Param-* headers
/// without a prior <see cref="McpClient.ListToolsAsync"/> call.
/// </summary>
public class AddKnownToolsHeaderTests(ITestOutputHelper outputHelper) : KestrelInMemoryTest(outputHelper), IAsyncDisposable
{
private WebApplication? _app;
/// <summary>
/// Captured headers from tools/call requests, keyed by JSON-RPC request id.
/// </summary>
private readonly ConcurrentDictionary<string, Dictionary<string, string>> _capturedHeaders = new();
private async Task StartAsync()
{
Builder.Services.Configure<JsonOptions>(options =>
{
options.SerializerOptions.TypeInfoResolverChain.Add(McpJsonUtilities.DefaultOptions.TypeInfoResolver!);
});
_app = Builder.Build();
_app.MapPost("/mcp", (JsonRpcMessage message, HttpContext context) =>
{
if (message is not JsonRpcRequest request)
{
return Results.Accepted();
}
if (request.Method == "initialize")
{
return Results.Json(new JsonRpcResponse
{
Id = request.Id,
Result = JsonSerializer.SerializeToNode(new InitializeResult
{
ProtocolVersion = "DRAFT-2026-v1",
Capabilities = new() { Tools = new() },
ServerInfo = new Implementation { Name = "header-capture-test", Version = "1.0" },
}, McpJsonUtilities.DefaultOptions)
});
}
if (request.Method == "tools/call")
{
// Capture all Mcp-Param-* headers from the incoming HTTP request
var paramHeaders = new Dictionary<string, string>(StringComparer.OrdinalIgnoreCase);
foreach (var header in context.Request.Headers)
{
if (header.Key.StartsWith("Mcp-Param-", StringComparison.OrdinalIgnoreCase))
{
paramHeaders[header.Key] = header.Value.ToString();
}
}
_capturedHeaders[request.Id.ToString()!] = paramHeaders;
var parameters = JsonSerializer.Deserialize(request.Params, McpJsonUtilities.DefaultOptions.GetTypeInfo(typeof(CallToolRequestParams))) as CallToolRequestParams;
return Results.Json(new JsonRpcResponse
{
Id = request.Id,
Result = JsonSerializer.SerializeToNode(new CallToolResult
{
Content = [new TextContentBlock { Text = $"ok" }],
}, McpJsonUtilities.DefaultOptions),
});
}
if (request.Method == "tools/list")
{
return Results.Json(new JsonRpcResponse
{
Id = request.Id,
Result = JsonSerializer.SerializeToNode(new ListToolsResult
{
Tools = [],
}, McpJsonUtilities.DefaultOptions),
});
}
return Results.Accepted();
});
await _app.StartAsync(TestContext.Current.CancellationToken);
HttpClient.DefaultRequestHeaders.Accept.Add(new("application/json"));
HttpClient.DefaultRequestHeaders.Accept.Add(new("text/event-stream"));
}
public async ValueTask DisposeAsync()
{
if (_app is not null)
{
await _app.DisposeAsync();
}
base.Dispose();
}
private static Tool CreateToolWithHeaders()
{
var schemaJson = """
{
"type": "object",
"properties": {
"region": {
"type": "string",
"x-mcp-header": "Region"
},
"priority": {
"type": "integer",
"x-mcp-header": "Priority"
}
},
"required": ["region", "priority"]
}
""";
return new Tool
{
Name = "my_tool",
InputSchema = JsonDocument.Parse(schemaJson).RootElement.Clone(),
};
}
[Fact]
public async Task AddKnownTools_ThenCallTool_SendsMcpParamHeaders_WithoutListToolsAsync()
{
await StartAsync();
await using var transport = new HttpClientTransport(new()
{
Endpoint = new("http://localhost:5000/mcp"),
TransportMode = HttpTransportMode.StreamableHttp,
}, HttpClient, LoggerFactory);
await using var client = await McpClient.CreateAsync(transport, loggerFactory: LoggerFactory,
cancellationToken: TestContext.Current.CancellationToken);
// Register the tool WITHOUT calling ListToolsAsync first — this is the core scenario from issue #1577
client.AddKnownTools([CreateToolWithHeaders()]);
// Call the tool
var result = await client.CallToolAsync(
"my_tool",
new Dictionary<string, object?> { ["region"] = "us-west-2", ["priority"] = 42 },
cancellationToken: TestContext.Current.CancellationToken);
Assert.NotNull(result);
// Verify that Mcp-Param-* headers were captured by the server
Assert.Single(_capturedHeaders);
var headers = _capturedHeaders.Values.First();
Assert.True(headers.ContainsKey("Mcp-Param-Region"), "Expected Mcp-Param-Region header to be sent");
Assert.Equal("us-west-2", headers["Mcp-Param-Region"]);
Assert.True(headers.ContainsKey("Mcp-Param-Priority"), "Expected Mcp-Param-Priority header to be sent");
Assert.Equal("42", headers["Mcp-Param-Priority"]);
}
[Fact]
public async Task CallToolWithoutRegisterOrList_DoesNotSendMcpParamHeaders()
{
await StartAsync();
await using var transport = new HttpClientTransport(new()
{
Endpoint = new("http://localhost:5000/mcp"),
TransportMode = HttpTransportMode.StreamableHttp,
}, HttpClient, LoggerFactory);
await using var client = await McpClient.CreateAsync(transport, loggerFactory: LoggerFactory,
cancellationToken: TestContext.Current.CancellationToken);
// Call the tool without AddKnownTools or ListToolsAsync — no Mcp-Param-* headers should be sent
var result = await client.CallToolAsync(
"my_tool",
new Dictionary<string, object?> { ["region"] = "us-west-2", ["priority"] = 42 },
cancellationToken: TestContext.Current.CancellationToken);
Assert.NotNull(result);
// Verify that NO Mcp-Param-* headers were sent
Assert.Single(_capturedHeaders);
var headers = _capturedHeaders.Values.First();
Assert.Empty(headers);
// Verify that a cache miss warning IS logged for HTTP transport
Assert.Contains(MockLoggerProvider.LogMessages, log =>
log.LogLevel == Microsoft.Extensions.Logging.LogLevel.Warning &&
log.Message.Contains("not found in cache during tools/call"));
}
[Fact]
public async Task AddKnownTools_SurvivesListToolsAsync_HeadersStillSent()
{
await StartAsync();
await using var transport = new HttpClientTransport(new()
{
Endpoint = new("http://localhost:5000/mcp"),
TransportMode = HttpTransportMode.StreamableHttp,
}, HttpClient, LoggerFactory);
await using var client = await McpClient.CreateAsync(transport, loggerFactory: LoggerFactory,
cancellationToken: TestContext.Current.CancellationToken);
// Register the tool first
client.AddKnownTools([CreateToolWithHeaders()]);
// Call ListToolsAsync — server returns empty list, but registered tool should survive
await client.ListToolsAsync(cancellationToken: TestContext.Current.CancellationToken);
// Call the registered tool — Mcp-Param-* headers should still be sent
var result = await client.CallToolAsync(
"my_tool",
new Dictionary<string, object?> { ["region"] = "eu-central-1", ["priority"] = 99 },
cancellationToken: TestContext.Current.CancellationToken);
Assert.NotNull(result);
// Verify headers were sent
Assert.Single(_capturedHeaders);
var headers = _capturedHeaders.Values.First();
Assert.True(headers.ContainsKey("Mcp-Param-Region"), "Expected Mcp-Param-Region header after ListToolsAsync");
Assert.Equal("eu-central-1", headers["Mcp-Param-Region"]);
Assert.True(headers.ContainsKey("Mcp-Param-Priority"), "Expected Mcp-Param-Priority header after ListToolsAsync");
Assert.Equal("99", headers["Mcp-Param-Priority"]);
}
[Fact]
public async Task RemoveKnownTools_ThenCallTool_NoMcpParamHeaders()
{
await StartAsync();
await using var transport = new HttpClientTransport(new()
{
Endpoint = new("http://localhost:5000/mcp"),
TransportMode = HttpTransportMode.StreamableHttp,
}, HttpClient, LoggerFactory);
await using var client = await McpClient.CreateAsync(transport, loggerFactory: LoggerFactory,
cancellationToken: TestContext.Current.CancellationToken);
// Register then remove — headers should no longer be sent
client.AddKnownTools([CreateToolWithHeaders()]);
client.RemoveKnownTools(["my_tool"]);
var result = await client.CallToolAsync(
"my_tool",
new Dictionary<string, object?> { ["region"] = "us-east-1", ["priority"] = 1 },
cancellationToken: TestContext.Current.CancellationToken);
Assert.NotNull(result);
// Verify no Mcp-Param-* headers were sent after removal
Assert.Single(_capturedHeaders);
var headers = _capturedHeaders.Values.First();
Assert.Empty(headers);
}
private static Tool CreateToolWithSingleHeader(string toolName, string headerName)
{
var schemaJson = $$"""
{
"type": "object",
"properties": {
"value": {
"type": "string",
"x-mcp-header": "{{headerName}}"
}
},
"required": ["value"]
}
""";
return new Tool
{
Name = toolName,
InputSchema = JsonDocument.Parse(schemaJson).RootElement.Clone(),
};
}
[Fact]
public async Task AddKnownTools_ServerReturnsEmptyList_RegisteredToolStillUsedForHeaders()
{
// Staleness test: register foo → server returns [] → ListToolsAsync → call foo → headers still sent
await StartAsync();
await using var transport = new HttpClientTransport(new()
{
Endpoint = new("http://localhost:5000/mcp"),
TransportMode = HttpTransportMode.StreamableHttp,
}, HttpClient, LoggerFactory);
await using var client = await McpClient.CreateAsync(transport, loggerFactory: LoggerFactory,
cancellationToken: TestContext.Current.CancellationToken);
// Register tool, then ListToolsAsync returns empty list from server
client.AddKnownTools([CreateToolWithHeaders()]);
await client.ListToolsAsync(cancellationToken: TestContext.Current.CancellationToken);
// Call the registered tool — headers should still be sent (sticky registration)
var result = await client.CallToolAsync(
"my_tool",
new Dictionary<string, object?> { ["region"] = "ap-southeast-1", ["priority"] = 5 },
cancellationToken: TestContext.Current.CancellationToken);
Assert.NotNull(result);
Assert.Single(_capturedHeaders);
var headers = _capturedHeaders.Values.First();
Assert.True(headers.ContainsKey("Mcp-Param-Region"), "Expected Mcp-Param-Region after server returned empty list");
Assert.Equal("ap-southeast-1", headers["Mcp-Param-Region"]);
Assert.True(headers.ContainsKey("Mcp-Param-Priority"), "Expected Mcp-Param-Priority after server returned empty list");
Assert.Equal("5", headers["Mcp-Param-Priority"]);
}
[Fact]
public async Task AddKnownTools_ReRegisterOverwrite_LastWriteWinsHeaders()
{
// Last-write-wins: register foo with schema A → register foo with schema B → call → headers reflect schema B
await StartAsync();
await using var transport = new HttpClientTransport(new()
{
Endpoint = new("http://localhost:5000/mcp"),
TransportMode = HttpTransportMode.StreamableHttp,
}, HttpClient, LoggerFactory);
await using var client = await McpClient.CreateAsync(transport, loggerFactory: LoggerFactory,
cancellationToken: TestContext.Current.CancellationToken);
// Register with header "SchemaA", then overwrite with "SchemaB"
client.AddKnownTools([CreateToolWithSingleHeader("my_tool", "SchemaA")]);
client.AddKnownTools([CreateToolWithSingleHeader("my_tool", "SchemaB")]);
var result = await client.CallToolAsync(
"my_tool",
new Dictionary<string, object?> { ["value"] = "test" },
cancellationToken: TestContext.Current.CancellationToken);
Assert.NotNull(result);
Assert.Single(_capturedHeaders);
var headers = _capturedHeaders.Values.First();
// SchemaA header should NOT be present
Assert.False(headers.ContainsKey("Mcp-Param-SchemaA"), "SchemaA header should have been overwritten");
// SchemaB header SHOULD be present (last write wins)
Assert.True(headers.ContainsKey("Mcp-Param-SchemaB"), "Expected Mcp-Param-SchemaB from overwritten registration");
Assert.Equal("test", headers["Mcp-Param-SchemaB"]);
}
}

View file

@ -0,0 +1,430 @@
using Microsoft.Extensions.DependencyInjection;
using ModelContextProtocol.Client;
using ModelContextProtocol.Protocol;
using ModelContextProtocol.Server;
using System.Text.Json;
namespace ModelContextProtocol.Tests.Client;
public class McpClientAddKnownToolsTests : ClientServerTestBase
{
private const string ServerToolName = "ServerTool";
private const string ServerToolName2 = "ServerTool2";
public McpClientAddKnownToolsTests(ITestOutputHelper outputHelper)
: base(outputHelper)
{
}
protected override void ConfigureServices(ServiceCollection services, IMcpServerBuilder mcpServerBuilder)
{
mcpServerBuilder.WithTools([
McpServerTool.Create(
(string input) => $"echo {input}",
new() { Name = ServerToolName }),
McpServerTool.Create(
(string input) => $"echo2 {input}",
new() { Name = ServerToolName2 }),
]);
}
private static Tool CreateTool(string name, string? headerAnnotation = null)
{
string schemaJson = headerAnnotation is not null
? $$"""
{
"type": "object",
"properties": {
"param1": {
"type": "string",
"x-mcp-header": "{{headerAnnotation}}"
}
}
}
"""
: """
{
"type": "object",
"properties": {
"param1": {
"type": "string"
}
}
}
""";
return new Tool
{
Name = name,
InputSchema = JsonDocument.Parse(schemaJson).RootElement.Clone(),
};
}
private static Tool CreateInvalidTool(string name)
{
// Colon in header name is invalid
var schemaJson = """
{
"type": "object",
"properties": {
"param1": {
"type": "string",
"x-mcp-header": "Invalid:Header"
}
}
}
""";
return new Tool
{
Name = name,
InputSchema = JsonDocument.Parse(schemaJson).RootElement.Clone(),
};
}
[Fact]
public async Task AddKnownTools_ThenListToolsAsync_ServerToolsStillReturned()
{
// Arrange
await using var client = await CreateMcpClientForServer();
var registeredTool = CreateTool("MyRegisteredTool", "X-Custom");
// Act — register without calling ListToolsAsync first, then list
client.AddKnownTools([registeredTool]);
var tools = await client.ListToolsAsync(cancellationToken: TestContext.Current.CancellationToken);
// Assert — ListToolsAsync returns server tools (registered tools stay in cache for header generation
// but are not returned by ListToolsAsync which only returns server-reported tools)
Assert.Contains(tools, t => t.Name == ServerToolName);
Assert.Contains(tools, t => t.Name == ServerToolName2);
Assert.Equal(2, tools.Count);
}
[Fact]
public async Task AddKnownTools_ThenMultipleListToolsAsync_ServerToolsAlwaysRepopulated()
{
// Arrange
await using var client = await CreateMcpClientForServer();
client.AddKnownTools([CreateTool("MyRegisteredTool", "X-Custom")]);
// Act — ListToolsAsync clears non-registered tools and repopulates from server
var tools1 = await client.ListToolsAsync(cancellationToken: TestContext.Current.CancellationToken);
var tools2 = await client.ListToolsAsync(cancellationToken: TestContext.Current.CancellationToken);
// Assert — server tools repopulated correctly after each clear
Assert.Contains(tools1, t => t.Name == ServerToolName);
Assert.Contains(tools1, t => t.Name == ServerToolName2);
Assert.Contains(tools2, t => t.Name == ServerToolName);
Assert.Contains(tools2, t => t.Name == ServerToolName2);
}
[Fact]
public async Task ListToolsAsync_ThenRegisterTool_ServerToolsStillRepopulated()
{
// Arrange
await using var client = await CreateMcpClientForServer();
// Act — list first, then register, then list again
var tools = await client.ListToolsAsync(cancellationToken: TestContext.Current.CancellationToken);
Assert.Equal(2, tools.Count);
client.AddKnownTools([CreateTool("MyRegisteredTool", "X-Custom")]);
var tools2 = await client.ListToolsAsync(cancellationToken: TestContext.Current.CancellationToken);
// Assert — server tools still repopulated after registering
Assert.Contains(tools2, t => t.Name == ServerToolName);
Assert.Contains(tools2, t => t.Name == ServerToolName2);
}
[Fact]
public async Task RegisterTool_ListToolsAsync_RegisterTool_ServerToolsIntact()
{
// Arrange
await using var client = await CreateMcpClientForServer();
// Act — register, list, register again
client.AddKnownTools([CreateTool("FirstRegistered", "X-First")]);
await client.ListToolsAsync(cancellationToken: TestContext.Current.CancellationToken);
client.AddKnownTools([CreateTool("SecondRegistered", "X-Second")]);
// Another ListToolsAsync — server tools should still be repopulated
var tools = await client.ListToolsAsync(cancellationToken: TestContext.Current.CancellationToken);
Assert.Contains(tools, t => t.Name == ServerToolName);
Assert.Contains(tools, t => t.Name == ServerToolName2);
}
[Fact]
public async Task AddKnownTools_WithSameNameAsServerTool_ServerDefinitionReturned()
{
// Arrange
await using var client = await CreateMcpClientForServer();
var registeredTool = CreateTool(ServerToolName, "X-Override");
// Act — register a tool with the same name as a server tool, then list
client.AddKnownTools([registeredTool]);
var tools = await client.ListToolsAsync(cancellationToken: TestContext.Current.CancellationToken);
// Assert — server's definition is returned by ListToolsAsync
Assert.Contains(tools, t => t.Name == ServerToolName);
// After another ListToolsAsync, the tool is still present (pinned as registered + server tool)
var tools2 = await client.ListToolsAsync(cancellationToken: TestContext.Current.CancellationToken);
Assert.Contains(tools2, t => t.Name == ServerToolName);
}
[Fact]
public async Task AddKnownTools_WithInvalidSchema_ThrowsArgumentException()
{
// Arrange
await using var client = await CreateMcpClientForServer();
var invalidTool = CreateInvalidTool("BadTool");
var validTool = CreateTool("GoodTool", "X-Good");
// Act & Assert — all-or-nothing: neither tool should be added
var ex = Assert.Throws<ArgumentException>(() => client.AddKnownTools([invalidTool, validTool]));
Assert.Contains("BadTool", ex.Message);
// Server tools still work normally
var tools = await client.ListToolsAsync(cancellationToken: TestContext.Current.CancellationToken);
Assert.Contains(tools, t => t.Name == ServerToolName);
Assert.Contains(tools, t => t.Name == ServerToolName2);
}
[Fact]
public async Task AddKnownTools_DuplicateRegistration_DoesNotBreakCache()
{
// Arrange
await using var client = await CreateMcpClientForServer();
var tool1 = CreateTool("MyTool", "X-First");
var tool2 = CreateTool("MyTool", "X-Second");
// Act — register same name twice; second should overwrite
client.AddKnownTools([tool1]);
client.AddKnownTools([tool2]);
// Assert — cache clearing still works; server tools repopulated
await client.ListToolsAsync(cancellationToken: TestContext.Current.CancellationToken);
var tools = await client.ListToolsAsync(cancellationToken: TestContext.Current.CancellationToken);
Assert.Contains(tools, t => t.Name == ServerToolName);
Assert.Contains(tools, t => t.Name == ServerToolName2);
}
[Fact]
public async Task AddKnownTools_NullArgument_ThrowsArgumentNullException()
{
await using var client = await CreateMcpClientForServer();
Assert.Throws<ArgumentNullException>(() => client.AddKnownTools(null!));
}
[Fact]
public async Task MultipleListToolsAsync_WithRegisteredTools_ServerToolsAlwaysRepopulated()
{
// Arrange
await using var client = await CreateMcpClientForServer();
client.AddKnownTools([CreateTool("PinnedTool", "X-Pinned")]);
// Act — call ListToolsAsync multiple times
await client.ListToolsAsync(cancellationToken: TestContext.Current.CancellationToken);
await client.ListToolsAsync(cancellationToken: TestContext.Current.CancellationToken);
var tools = await client.ListToolsAsync(cancellationToken: TestContext.Current.CancellationToken);
// Assert — server tools repopulated each time despite registered tool in cache
Assert.Contains(tools, t => t.Name == ServerToolName);
Assert.Contains(tools, t => t.Name == ServerToolName2);
}
[Fact]
public async Task AddKnownTools_WithNoHeaderAnnotation_StillAccepted()
{
// Arrange — a tool without x-mcp-header is still valid and should be cached
await using var client = await CreateMcpClientForServer();
var tool = CreateTool("PlainTool");
// Act — register a tool with no x-mcp-header; should not throw
client.AddKnownTools([tool]);
// Assert — server tools still repopulated after cache clears
await client.ListToolsAsync(cancellationToken: TestContext.Current.CancellationToken);
var tools = await client.ListToolsAsync(cancellationToken: TestContext.Current.CancellationToken);
Assert.Contains(tools, t => t.Name == ServerToolName);
Assert.Contains(tools, t => t.Name == ServerToolName2);
}
[Fact]
public async Task AddKnownTools_ThenCallTool_RegisteredToolUsedForCacheLookup()
{
// Arrange — register a tool with the same name as a server tool so CallToolAsync succeeds server-side
await using var client = await CreateMcpClientForServer();
var tool = CreateTool(ServerToolName, "X-Custom");
// Act — register without ListToolsAsync, then call the tool directly
client.AddKnownTools([tool]);
// The tool is in the cache, so SendRequestAsync will find it for header attachment.
// The server has a tool with this name, so the call succeeds.
var result = await client.CallToolAsync(
ServerToolName,
new Dictionary<string, object?> { ["input"] = "test" },
cancellationToken: TestContext.Current.CancellationToken);
// Assert — call succeeded (tool was found in cache, request was processed by server)
Assert.NotNull(result);
Assert.Contains(result.Content, c => c is TextContentBlock text && text.Text == "echo test");
}
[Fact]
public async Task RemoveKnownTools_RemovedToolNoLongerSurvivesListToolsAsync()
{
// Arrange
await using var client = await CreateMcpClientForServer();
client.AddKnownTools([CreateTool("MyTool", "X-Custom")]);
// Act — remove the known tool
client.RemoveKnownTools(["MyTool"]);
// Assert — server tools still repopulated, removed tool doesn't interfere
var tools = await client.ListToolsAsync(cancellationToken: TestContext.Current.CancellationToken);
Assert.Contains(tools, t => t.Name == ServerToolName);
Assert.Contains(tools, t => t.Name == ServerToolName2);
}
[Fact]
public async Task RemoveKnownTools_NonExistentName_IsNoOp()
{
await using var client = await CreateMcpClientForServer();
// Should not throw
client.RemoveKnownTools(["NonExistentTool"]);
var tools = await client.ListToolsAsync(cancellationToken: TestContext.Current.CancellationToken);
Assert.Contains(tools, t => t.Name == ServerToolName);
}
[Fact]
public async Task RemoveKnownTools_NullArgument_ThrowsArgumentNullException()
{
await using var client = await CreateMcpClientForServer();
Assert.Throws<ArgumentNullException>(() => client.RemoveKnownTools(null!));
}
[Fact]
public async Task RemoveKnownTools_PartialRemove_OtherToolsSurvive()
{
// Arrange
await using var client = await CreateMcpClientForServer();
client.AddKnownTools([
CreateTool("ToolA", "X-A"),
CreateTool("ToolB", "X-B"),
]);
// Act — remove only ToolA
client.RemoveKnownTools(["ToolA"]);
// Assert — ToolB still survives cache clears, server tools repopulated
await client.ListToolsAsync(cancellationToken: TestContext.Current.CancellationToken);
var tools = await client.ListToolsAsync(cancellationToken: TestContext.Current.CancellationToken);
Assert.Contains(tools, t => t.Name == ServerToolName);
Assert.Contains(tools, t => t.Name == ServerToolName2);
}
[Fact]
public async Task ClearKnownTools_RemovesAllKnownToolsFromCache()
{
// Arrange
await using var client = await CreateMcpClientForServer();
client.AddKnownTools([
CreateTool("ToolA", "X-A"),
CreateTool("ToolB", "X-B"),
]);
// Act
client.ClearKnownTools();
// Assert — server tools still work after clearing known tools
var tools = await client.ListToolsAsync(cancellationToken: TestContext.Current.CancellationToken);
Assert.Contains(tools, t => t.Name == ServerToolName);
Assert.Contains(tools, t => t.Name == ServerToolName2);
}
[Fact]
public async Task ClearKnownTools_WhenEmpty_IsNoOp()
{
await using var client = await CreateMcpClientForServer();
// Should not throw when nothing is registered
client.ClearKnownTools();
var tools = await client.ListToolsAsync(cancellationToken: TestContext.Current.CancellationToken);
Assert.Contains(tools, t => t.Name == ServerToolName);
}
[Fact]
public async Task ClearKnownTools_ThenAddKnownTools_WorksCorrectly()
{
// Arrange
await using var client = await CreateMcpClientForServer();
client.AddKnownTools([CreateTool("ToolA", "X-A")]);
// Act — clear then add new tools
client.ClearKnownTools();
client.AddKnownTools([CreateTool("ToolC", "X-C")]);
// Assert — server tools repopulated, new tool doesn't interfere
var tools = await client.ListToolsAsync(cancellationToken: TestContext.Current.CancellationToken);
Assert.Contains(tools, t => t.Name == ServerToolName);
}
[Fact]
public async Task AddKnownTools_PartialFailure_NothingRegistered()
{
// Arrange — [valid, invalid, valid] should register nothing (all-or-nothing)
await using var client = await CreateMcpClientForServer();
var valid1 = CreateTool("Valid1", "X-One");
var invalid = CreateInvalidTool("BadTool");
var valid2 = CreateTool("Valid2", "X-Two");
// Act & Assert — throws, no tools registered
Assert.Throws<ArgumentException>(() => client.AddKnownTools([valid1, invalid, valid2]));
// Server tools still work; none of the valid tools were cached
var tools = await client.ListToolsAsync(cancellationToken: TestContext.Current.CancellationToken);
Assert.Contains(tools, t => t.Name == ServerToolName);
Assert.Equal(2, tools.Count);
}
[Fact]
public async Task AddKnownTools_NullElementInMiddle_NothingRegistered()
{
// Arrange — null element at index 1; elements before it should not be cached
await using var client = await CreateMcpClientForServer();
var valid = CreateTool("Valid", "X-Valid");
// Act & Assert — throws ArgumentNullException on null element, nothing cached
Assert.Throws<ArgumentNullException>(() => client.AddKnownTools([valid, null!, CreateTool("Other", "X-Other")]));
// Server tools still work; valid tool was NOT cached due to atomicity
var tools = await client.ListToolsAsync(cancellationToken: TestContext.Current.CancellationToken);
Assert.Equal(2, tools.Count);
}
[Fact]
public async Task CallToolWithoutCache_PipeTransport_DoesNotLogWarning()
{
// Arrange — pipe transport should NOT log a cache miss warning
await using var client = await CreateMcpClientForServer();
// Act — call a server tool without populating cache via ListToolsAsync
var result = await client.CallToolAsync(
ServerToolName,
new Dictionary<string, object?> { ["input"] = "test" },
cancellationToken: TestContext.Current.CancellationToken);
// Assert — call succeeds and no cache miss warning is logged (pipe transport, not HTTP)
Assert.NotNull(result);
Assert.DoesNotContain(MockLoggerProvider.LogMessages, log =>
log.LogLevel == Microsoft.Extensions.Logging.LogLevel.Warning &&
log.Message.Contains("not found in cache during tools/call"));
}
}