McpServerBuilderExtensionsToolsTests.cs

using System.ComponentModel;
using System.Text.Json;
using ModelContextProtocol.Server;
using Microsoft.Extensions.DependencyInjection;
using ModelContextProtocol.Protocol.Transport;
using System.IO.Pipelines;
using ModelContextProtocol.Client;
using ModelContextProtocol.Configuration;
using ModelContextProtocol.Tests.Transport;
using System.Text.RegularExpressions;
using Microsoft.Extensions.AI;
using System.Threading.Channels;
using ModelContextProtocol.Protocol.Messages;

namespace ModelContextProtocol.Tests.Configuration;

public class McpServerBuilderExtensionsToolsTests : IAsyncDisposable
{
    private Pipe _clientToServerPipe = new();
    private Pipe _serverToClientPipe = new();
    private readonly IMcpServerBuilder _builder;
    private readonly IMcpServer _server;

    public McpServerBuilderExtensionsToolsTests()
    {
        ServiceCollection sc = new();
        sc.AddSingleton<IServerTransport>(new StdioServerTransport("TestServer", _clientToServerPipe.Reader.AsStream(), _serverToClientPipe.Writer.AsStream()));
        _builder = sc.AddMcpServer().WithTools<EchoTool>();
        _server = sc.BuildServiceProvider().GetRequiredService<IMcpServer>();
    }

    public ValueTask DisposeAsync()
    {
        _clientToServerPipe.Writer.Complete();
        _serverToClientPipe.Writer.Complete();
        return _server.DisposeAsync();
    }

    private async Task<IMcpClient> CreateMcpClientForServer()
    {
        await _server.StartAsync(TestContext.Current.CancellationToken);

        var stdin = new StreamReader(_serverToClientPipe.Reader.AsStream());
        var stdout = new StreamWriter(_clientToServerPipe.Writer.AsStream());

        var serverConfig = new McpServerConfig()
        {
            Id = "TestServer",
            Name = "TestServer",
            TransportType = "ignored",
        };

        return await McpClientFactory.CreateAsync(
            serverConfig,
            createTransportFunc: (_, _) => new StreamClientTransport(stdin, stdout),
            cancellationToken: TestContext.Current.CancellationToken);
    }

    [Fact]
    public void Adds_Tools_To_Server()
    {
        var tools = _server.ServerOptions?.Capabilities?.Tools?.ToolCollection;
        Assert.NotNull(tools);
        Assert.NotEmpty(tools);
    }

    [Fact]
    public async Task Can_List_Registered_Tools()
    {
        IMcpClient client = await CreateMcpClientForServer();

        var tools = await client.ListToolsAsync(TestContext.Current.CancellationToken);
        Assert.Equal(10, tools.Count);

        McpClientTool echoTool = tools.First(t => t.Name == "Echo");
        Assert.Equal("Echo", echoTool.Name);
        Assert.Equal("Echoes the input back to the client.", echoTool.Description);
        Assert.Equal("object", echoTool.JsonSchema.GetProperty("type").GetString());
        Assert.Equal(JsonValueKind.Object, echoTool.JsonSchema.GetProperty("properties").GetProperty("message").ValueKind);
        Assert.Equal("the echoes message", echoTool.JsonSchema.GetProperty("properties").GetProperty("message").GetProperty("description").GetString());
        Assert.Equal(1, echoTool.JsonSchema.GetProperty("required").GetArrayLength());

        McpClientTool doubleEchoTool = tools.First(t => t.Name == "double_echo");
        Assert.Equal("double_echo", doubleEchoTool.Name);
        Assert.Equal("Echoes the input back to the client.", doubleEchoTool.Description);
    }

    [Fact]
    public async Task Can_Be_Notified_Of_Tool_Changes()
    {
        IMcpClient client = await CreateMcpClientForServer();

        var tools = await client.ListToolsAsync(TestContext.Current.CancellationToken);
        Assert.Equal(10, tools.Count);

        Channel<JsonRpcNotification> listChanged = Channel.CreateUnbounded<JsonRpcNotification>();
        client.AddNotificationHandler("notifications/tools/list_changed", notification =>
        {
            listChanged.Writer.TryWrite(notification);
            return Task.CompletedTask;
        });

        var notificationRead = listChanged.Reader.ReadAsync(TestContext.Current.CancellationToken);
        Assert.False(notificationRead.IsCompleted);

        var serverTools = _server.ServerOptions.Capabilities?.Tools?.ToolCollection;
        Assert.NotNull(serverTools);

        var newTool = McpServerTool.Create([McpServerTool(name: "NewTool")] () => "42");
        serverTools.Add(newTool);
        await notificationRead;

        tools = await client.ListToolsAsync(TestContext.Current.CancellationToken);
        Assert.Equal(11, tools.Count);
        Assert.Contains(tools, t => t.Name == "NewTool");

        notificationRead = listChanged.Reader.ReadAsync(TestContext.Current.CancellationToken);
        Assert.False(notificationRead.IsCompleted);
        serverTools.Remove(newTool);
        await notificationRead;

        tools = await client.ListToolsAsync(TestContext.Current.CancellationToken);
        Assert.Equal(10, tools.Count);
        Assert.DoesNotContain(tools, t => t.Name == "NewTool");
    }

    [Fact]
    public async Task Can_Call_Registered_Tool()
    {
        IMcpClient client = await CreateMcpClientForServer();

        var result = await client.CallToolAsync(
            "Echo",
            new Dictionary<string, object?>() { ["message"] = "Peter" }, 
            TestContext.Current.CancellationToken);

        Assert.NotNull(result);
        Assert.NotNull(result.Content);
        Assert.NotEmpty(result.Content);

        Assert.Equal("hello Peter", result.Content[0].Text);
        Assert.Equal("text", result.Content[0].Type);
    }

    [Fact]
    public async Task Can_Call_Registered_Tool_With_Array_Result()
    {
        IMcpClient client = await CreateMcpClientForServer();

        var result = await client.CallToolAsync(
            "EchoArray",
            new Dictionary<string, object?>() { ["message"] = "Peter" },
            TestContext.Current.CancellationToken);

        Assert.NotNull(result.Content);
        Assert.NotEmpty(result.Content);

        Assert.Equal("hello Peter", result.Content[0].Text);
        Assert.Equal("hello2 Peter", result.Content[1].Text);
    }

    [Fact]
    public async Task Can_Call_Registered_Tool_With_Null_Result()
    {
        IMcpClient client = await CreateMcpClientForServer();

        var result = await client.CallToolAsync(
            "ReturnNull",
            cancellationToken: TestContext.Current.CancellationToken);

        Assert.NotNull(result);
        Assert.NotNull(result.Content);
        Assert.Empty(result.Content);
    }

    [Fact]
    public async Task Can_Call_Registered_Tool_With_Json_Result()
    {
        IMcpClient client = await CreateMcpClientForServer();

        var result = await client.CallToolAsync(
            "ReturnJson",
            cancellationToken: TestContext.Current.CancellationToken);

        Assert.NotNull(result);
        Assert.NotNull(result.Content);
        Assert.NotEmpty(result.Content);

        Assert.Equal("""{"SomeProp":false}""", Regex.Replace(result.Content[0].Text ?? string.Empty, "\\s+", ""));
        Assert.Equal("text", result.Content[0].Type);
    }

    [Fact]
    public async Task Can_Call_Registered_Tool_With_Int_Result()
    {
        IMcpClient client = await CreateMcpClientForServer();

        var result = await client.CallToolAsync(
            "ReturnInteger",
            cancellationToken: TestContext.Current.CancellationToken);

        Assert.NotNull(result.Content);
        Assert.NotEmpty(result.Content);

        Assert.Equal("5", result.Content[0].Text);
        Assert.Equal("text", result.Content[0].Type);
    }

    [Fact]
    public async Task Can_Call_Registered_Tool_And_Pass_ComplexType()
    {
        IMcpClient client = await CreateMcpClientForServer();

        var result = await client.CallToolAsync(
            "EchoComplex",
            new Dictionary<string, object?>() { ["complex"] = JsonDocument.Parse("""{"Name": "Peter", "Age": 25}""").RootElement },
            cancellationToken: TestContext.Current.CancellationToken);

        Assert.NotNull(result);
        Assert.NotNull(result.Content);
        Assert.NotEmpty(result.Content);

        Assert.Equal("Peter", result.Content[0].Text);
        Assert.Equal("text", result.Content[0].Type);
    }

    [Fact]
    public async Task Returns_IsError_Content_When_Tool_Fails()
    {
        IMcpClient client = await CreateMcpClientForServer();

        var result = await client.CallToolAsync(
            "ThrowException",
            cancellationToken: TestContext.Current.CancellationToken);

        Assert.True(result.IsError);
        Assert.NotNull(result.Content);
        Assert.NotEmpty(result.Content);
        Assert.Contains("Test error", result.Content[0].Text);
    }

    [Fact]
    public async Task Throws_Exception_On_Unknown_Tool()
    {
        IMcpClient client = await CreateMcpClientForServer();

        var e = await Assert.ThrowsAsync<McpClientException>(async () => await client.CallToolAsync(
            "NotRegisteredTool",
            cancellationToken: TestContext.Current.CancellationToken));

        Assert.Contains("'NotRegisteredTool'", e.Message);
    }

    [Fact(Skip = "https://github.com/dotnet/extensions/issues/6124")]
    public async Task Throws_Exception_Missing_Parameter()
    {
        IMcpClient client = await CreateMcpClientForServer();

        var e = await Assert.ThrowsAsync<McpClientException>(async () => await client.CallToolAsync(
            "Echo",
            cancellationToken: TestContext.Current.CancellationToken));

        Assert.Equal("Missing required argument 'message'.", e.Message);
    }

    [Fact]
    public void WithTools_InvalidArgs_Throws()
    {
        Assert.Throws<ArgumentNullException>("toolTypes", () => _builder.WithTools((IEnumerable<Type>)null!));

        IMcpServerBuilder nullBuilder = null!;
        Assert.Throws<ArgumentNullException>("builder", () => nullBuilder.WithTools<object>());
        Assert.Throws<ArgumentNullException>("builder", () => nullBuilder.WithTools(Array.Empty<Type>()));
        Assert.Throws<ArgumentNullException>("builder", () => nullBuilder.WithToolsFromAssembly());
    }

    [Fact]
    public void Empty_Enumerables_Is_Allowed()
    {
        _builder.WithTools(toolTypes: []); // no exception
        _builder.WithTools<object>(); // no exception even though no tools exposed
        _builder.WithToolsFromAssembly(typeof(AIFunction).Assembly); // no exception even though no tools exposed
    }

    [Fact]
    public void Register_Tools_From_Current_Assembly()
    {
        ServiceCollection sc = new();
        sc.AddMcpServer().WithToolsFromAssembly();
        IServiceProvider services = sc.BuildServiceProvider();

        Assert.Contains(services.GetServices<McpServerTool>(), t => t.ProtocolTool.Name == "Echo");
    }

    [Fact]
    public async Task Recognizes_Parameter_Types()
    {
        IMcpClient client = await CreateMcpClientForServer();

        var tools = await client.ListToolsAsync(TestContext.Current.CancellationToken);

        Assert.NotNull(tools);
        Assert.NotEmpty(tools);

        var tool = tools.First(t => t.Name == "TestTool");
        Assert.Equal("TestTool", tool.Name);
        Assert.Empty(tool.Description!);
        Assert.Equal("object", tool.JsonSchema.GetProperty("type").GetString());

        Assert.Contains("integer", tool.JsonSchema.GetProperty("properties").GetProperty("number").GetProperty("type").GetString());
        Assert.Contains("number", tool.JsonSchema.GetProperty("properties").GetProperty("otherNumber").GetProperty("type").GetString());
        Assert.Contains("boolean", tool.JsonSchema.GetProperty("properties").GetProperty("someCheck").GetProperty("type").GetString());
        Assert.Contains("string", tool.JsonSchema.GetProperty("properties").GetProperty("someDate").GetProperty("type").GetString());
        Assert.Contains("string", tool.JsonSchema.GetProperty("properties").GetProperty("someOtherDate").GetProperty("type").GetString());
        Assert.Contains("array", tool.JsonSchema.GetProperty("properties").GetProperty("data").GetProperty("type").GetString());
        Assert.Contains("object", tool.JsonSchema.GetProperty("properties").GetProperty("complexObject").GetProperty("type").GetString());
    }

    [Fact]
    public void Register_Tools_From_Multiple_Sources()
    {
        ServiceCollection sc = new();
        sc.AddMcpServer()
            .WithTools<EchoTool>()
            .WithTools<AnotherToolType>()
            .WithTools(typeof(ToolTypeWithNoAttribute));
        IServiceProvider services = sc.BuildServiceProvider();

        Assert.Contains(services.GetServices<McpServerTool>(), t => t.ProtocolTool.Name == "double_echo");
        Assert.Contains(services.GetServices<McpServerTool>(), t => t.ProtocolTool.Name == "DifferentName");
        Assert.Contains(services.GetServices<McpServerTool>(), t => t.ProtocolTool.Name == "MethodB");
        Assert.Contains(services.GetServices<McpServerTool>(), t => t.ProtocolTool.Name == "MethodC");
        Assert.Contains(services.GetServices<McpServerTool>(), t => t.ProtocolTool.Name == "MethodD");
    }

    [McpServerToolType]
    public sealed class EchoTool
    {
        [McpServerTool, Description("Echoes the input back to the client.")]
        public static string Echo([Description("the echoes message")] string message)
        {
            return "hello " + message;
        }

        [McpServerTool("double_echo"), Description("Echoes the input back to the client.")]
        public static string Echo2(string message)
        {
            return "hello hello" + message;
        }

        [McpServerTool]
        public static string TestTool(int number, double otherNumber, bool someCheck, DateTime someDate, DateTimeOffset someOtherDate, string[] data, ComplexObject complexObject)
        {
            return "hello hello";
        }

        [McpServerTool]
        public static string[] EchoArray(string message)
        {
            return ["hello " + message, "hello2 " + message];
        }

        [McpServerTool]
        public static string? ReturnNull()
        {
            return null;
        }

        [McpServerTool]
        public static JsonElement ReturnJson()
        {
            return JsonDocument.Parse("{\"SomeProp\": false}").RootElement;
        }

        [McpServerTool]
        public static int ReturnInteger()
        {
            return 5;
        }

        [McpServerTool]
        public static string ThrowException()
        {
            throw new InvalidOperationException("Test error");
        }

        [McpServerTool]
        public static int ReturnCancellationToken(CancellationToken cancellationToken)
        {
            return cancellationToken.GetHashCode();
        }

        [McpServerTool]
        public static string EchoComplex(ComplexObject complex)
        {
            return complex.Name!;
        }
    }

    [McpServerToolType]
    internal class AnotherToolType
    {
        [McpServerTool("DifferentName")]
        private static string MethodA(int a) => a.ToString();

        [McpServerTool]
        internal static string MethodB(string b) => b.ToString();

        [McpServerTool]
        protected static string MethodC(long c) => c.ToString();
    }

    internal class ToolTypeWithNoAttribute
    {
        [McpServerTool]
        public static string MethodD(string d) => d.ToString();
    }

    public class ComplexObject
    {
        public string? Name { get; set; }
        public int Age { get; set; }
    }
}