Overhaul tool handling (#89)

* Overhaul tool handling

- Renames [McpTool{Type}] to [McpServerTool{Type}], in order to distinguish representations of tools on the server from tools on the client.
- Enables [McpServerTool] methods to have arguments injected from DI, as well as a specific set of types from the implementation directly, like IMcpServer.
- Renames WithTools to WithToolsFromAssembly.
- All of the WithToolsXx methods now publish each individual McpServerTool into DI; other code can do so as well. The options setup code gathers all of the tools from DI and combines them into a collection which is then stored in McpServerOptions and used to construct the server.
- The server tools specified via DI as well as manually-provided handlers, using CallToolHandler as a fallback in case the requested tool doesn't exist in the tools collection, and ListToolsHandler augments the information for whatever tools exist.
- The tools are stored in McpServerOptions in a new McpServerToolCollection type, which is a thread-safe container that exposes add/removal notification. Adding/removing tools triggers a change notification that in turn sends a notification to the client about a tools update.
- The ServerOptions are exposed from the server instance.
- Removed cursor-based APIs from McpClientExtensions.
- Changed McpClientExtensions APIs to return `Task<IList<...>` rather than `IAsyncEnumerable<...>`.

I'm sure this will need to evolve further before we're "done", but it's a significant improvement from where we are now. One area I'm not super happy with is the notifying collection; that might work ok for a stdio server, where you can grab the created collection from the server's options and mutate it, but I don't know how well that's going to work for sse servers.

* Make McpServerTool{Collection} extensible

* Fix README.md example

* Add notifications test

* Address feedback
This commit is contained in:
Stephen Toub 2025-03-25 15:32:46 -04:00 • committed by GitHub
commit caca2ae61d
No known key found for this signature in database
GPG key ID: B5690EEEBB952194

View file

@ -1,12 +1,9 @@
using ModelContextProtocol.Configuration;
using ModelContextProtocol.Protocol.Types;
using ModelContextProtocol.Server;
using ModelContextProtocol.Utils;
using ModelContextProtocol.Utils.Json;
using Microsoft.Extensions.AI;
using System.Diagnostics.CodeAnalysis;
using System.Reflection;
using System.Text.Json;
using Microsoft.Extensions.DependencyInjection;
namespace ModelContextProtocol;
@ -23,23 +20,17 @@ public static partial class McpServerBuilderExtensions
/// <typeparam name="TTool">The tool type.</typeparam>
/// <param name="builder">The builder instance.</param>
/// <exception cref="ArgumentNullException"><paramref name="builder"/> is <see langword="null"/>.</exception>
public static IMcpServerBuilder WithTool<[DynamicallyAccessedMembers(DynamicallyAccessedMemberTypes.PublicMethods)] TTool>(this IMcpServerBuilder builder)
public static IMcpServerBuilder WithTools<[DynamicallyAccessedMembers(DynamicallyAccessedMemberTypes.PublicMethods | DynamicallyAccessedMemberTypes.NonPublicMethods)] TTool>(
this IMcpServerBuilder builder)
{
Throw.IfNull(builder);
List<AIFunction> functions = [];
PopulateFunctions(typeof(TTool), functions);
return WithTools(builder, functions);
}
/// <summary>
/// Adds all tools marked with <see cref="McpToolTypeAttribute"/> from the current assembly to the server.
/// </summary>
/// <param name="builder">The builder instance.</param>
/// <exception cref="ArgumentNullException"><paramref name="builder"/> is <see langword="null"/>.</exception>
[RequiresUnreferencedCode(RequiresUnreferencedCodeMessage)]
public static IMcpServerBuilder WithTools(this IMcpServerBuilder builder)
{
return WithToolsFromAssembly(builder, Assembly.GetCallingAssembly());
foreach (var toolMethod in GetToolMethods(typeof(TTool)))
{
builder.Services.AddSingleton(services => McpServerTool.Create(toolMethod, services: services));
}
return builder;
}
/// <summary>
@ -55,147 +46,42 @@ public static partial class McpServerBuilderExtensions
Throw.IfNull(builder);
Throw.IfNull(toolTypes);
List<AIFunction> functions = [];
foreach (var toolType in toolTypes)
{
if (toolType is null)
if (toolType is not null)
{
throw new ArgumentNullException(nameof(toolTypes), $"A tool type provided by the enumerator was null.");
foreach (var toolMethod in GetToolMethods(toolType))
{
builder.Services.AddSingleton(services => McpServerTool.Create(toolMethod, services: services));
}
}
PopulateFunctions(toolType, functions);
}
return WithTools(builder, functions);
}
/// <summary>
/// Adds tools to the server.
/// </summary>
/// <param name="builder">The builder instance.</param>
/// <param name="functions"><see cref="AIFunction"/> instances to use as the tools.</param>
/// <exception cref="ArgumentNullException"><paramref name="builder"/> is <see langword="null"/>.</exception>
/// <exception cref="ArgumentNullException"><paramref name="functions"/> is <see langword="null"/>.</exception>
public static IMcpServerBuilder WithTools(this IMcpServerBuilder builder, params IEnumerable<AIFunction> functions)
{
Throw.IfNull(builder);
Throw.IfNull(functions);
List<Tool> tools = [];
Dictionary<string, Func<RequestContext<CallToolRequestParams>, CancellationToken, Task<CallToolResponse>>> callbacks = [];
foreach (AIFunction function in functions)
{
if (function is null)
{
throw new ArgumentNullException(nameof(functions), $"A function provided by the enumerator was null.");
}
tools.Add(new()
{
Name = function.Name,
Description = function.Description,
InputSchema = function.JsonSchema,
});
callbacks.Add(function.Name, async (request, cancellationToken) =>
{
cancellationToken.ThrowIfCancellationRequested();
object? result;
try
{
result = await function.InvokeAsync((request.Params?.Arguments ?? [])!, cancellationToken).ConfigureAwait(false);
}
catch (Exception e) when (e is not OperationCanceledException)
{
return new CallToolResponse()
{
IsError = true,
Content = [new() { Text = e.Message, Type = "text" }],
};
}
switch (result)
{
case JsonElement je when je.ValueKind == JsonValueKind.Null:
return new() { Content = [] };
case JsonElement je when je.ValueKind == JsonValueKind.Array:
return new() { Content = je.EnumerateArray().Select(x => new Content() { Text = x.ToString(), Type = "text" }).ToList() };
default:
return new() { Content = [new() { Text = result?.ToString(), Type = "text" }] };
}
});
}
builder.WithListToolsHandler((_, _) => Task.FromResult(new ListToolsResult() { Tools = tools }));
builder.WithCallToolHandler(async (request, cancellationToken) =>
{
if (request.Params is null || !callbacks.TryGetValue(request.Params.Name, out var callback))
{
throw new McpServerException($"Unknown tool '{request.Params?.Name}'");
}
return await callback(request, cancellationToken).ConfigureAwait(false);
});
return builder;
}
/// <summary>
/// Adds types marked with the <see cref="McpToolTypeAttribute"/> attribute from the given assembly as tools to the server.
/// Adds types marked with the <see cref="McpServerToolTypeAttribute"/> attribute from the given assembly as tools to the server.
/// </summary>
/// <param name="builder">The builder instance.</param>
/// <param name="assembly">The assembly to load the types from. Null to get the current assembly</param>
/// <param name="toolAssembly">The assembly to load the types from. Null to get the current assembly</param>
/// <exception cref="ArgumentNullException"><paramref name="builder"/> is <see langword="null"/>.</exception>
[RequiresUnreferencedCode(RequiresUnreferencedCodeMessage)]
public static IMcpServerBuilder WithToolsFromAssembly(this IMcpServerBuilder builder, Assembly? assembly = null)
public static IMcpServerBuilder WithToolsFromAssembly(this IMcpServerBuilder builder, Assembly? toolAssembly = null)
{
assembly ??= Assembly.GetCallingAssembly();
Throw.IfNull(builder);
List<Type> toolTypes = [];
toolAssembly ??= Assembly.GetCallingAssembly();
foreach (var type in assembly.GetTypes())
{
if (type.GetCustomAttribute<McpToolTypeAttribute>() is null)
{
continue;
}
foreach (var method in type.GetMethods(BindingFlags.Public | BindingFlags.Static))
{
if (method.GetCustomAttribute<McpToolAttribute>() is not null)
{
toolTypes.Add(type);
break;
}
}
}
return toolTypes.Count > 0 ?
WithTools(builder, toolTypes) :
builder;
return builder.WithTools(
from t in toolAssembly.GetTypes()
where t.GetCustomAttribute<McpServerToolTypeAttribute>() is not null
select t);
}
private static void PopulateFunctions(
[DynamicallyAccessedMembers(DynamicallyAccessedMemberTypes.PublicMethods)] Type toolType,
List<AIFunction> functions)
{
foreach (var method in toolType.GetMethods(BindingFlags.Public | BindingFlags.Static))
{
if (method.GetCustomAttribute<McpToolAttribute>() is not { } attribute)
{
continue;
}
functions.Add(AIFunctionFactory.Create(method, target: null, new()
{
Name = attribute.Name ?? method.Name,
}));
}
}
private static IEnumerable<MethodInfo> GetToolMethods(
[DynamicallyAccessedMembers(DynamicallyAccessedMemberTypes.PublicMethods | DynamicallyAccessedMemberTypes.NonPublicMethods)] Type toolType) =>
from method in toolType.GetMethods(BindingFlags.Public | BindingFlags.NonPublic | BindingFlags.Static)
where method.GetCustomAttribute<McpServerToolAttribute>() is not null
select method;
}