ClawdDotNet/tests/ClawdDotNet.Core.Tests/Mcp/McpClientTests.cs

using System.Net;
using System.Text;
using System.Text.Json;
using System.Text.Json.Nodes;
using ClawdDotNet.Core.Config;
using ClawdDotNet.Core.Mcp;
using ClawdDotNet.Core.Security;
using ClawdDotNet.Core.Staging;
using ClawdDotNet.Core.Tools;
using ClawdDotNet.Core.Tests.Infrastructure;
using Microsoft.Extensions.Logging.Abstractions;
using Shouldly;

namespace ClawdDotNet.Core.Tests.Mcp;

/// <summary>
/// #197: An external MCP server as one tool — spoken to over Streamable HTTP (JSON or an event
/// stream), on approval by default, its words cut short and its results foreign content.
/// </summary>
public sealed class McpClientTests
{
    /// <summary>A small MCP server: two tools, a session, answers as JSON or as an event stream.</summary>
    private sealed class FakeServer(bool eventStream = false) : HttpMessageHandler
    {
        public List<JsonObject> Received { get; } = [];
        public List<string?> Sessions { get; } = [];
        public List<string?> Authorization { get; } = [];
        public bool ForgetSessionOnce { get; set; }
        private int _session;

        protected override async Task<HttpResponseMessage> SendAsync(HttpRequestMessage request, CancellationToken ct)
        {
            var message = JsonNode.Parse(await request.Content!.ReadAsStringAsync(ct))!.AsObject();
            Received.Add(message);
            Sessions.Add(request.Headers.TryGetValues("Mcp-Session-Id", out var s) ? s.First() : null);
            Authorization.Add(request.Headers.Authorization?.ToString());

            var method = message["method"]!.GetValue<string>();
            if (message["id"] is null)
                return new HttpResponseMessage(HttpStatusCode.Accepted);
            if (ForgetSessionOnce && method != "initialize")
            {
                ForgetSessionOnce = false;
                return new HttpResponseMessage(HttpStatusCode.NotFound);
            }

            JsonNode result = method switch
            {
                "initialize" => new JsonObject { ["protocolVersion"] = McpClient.ProtocolVersion, ["capabilities"] = new JsonObject() },
                "tools/list" => new JsonObject
                {
                    ["tools"] = new JsonArray(
                        new JsonObject { ["name"] = "send_invoice", ["description"] = new string('d', 500), ["inputSchema"] = JsonNode.Parse("""{"type":"object","properties":{"customer":{"type":"string"}},"required":["customer"]}""") },
                        new JsonObject { ["name"] = "stock", ["description"] = "Stock level.", ["inputSchema"] = new JsonObject { ["type"] = "object" } })
                },
                "tools/call" when message["params"]!["name"]!.GetValue<string>() == "stock"
                    => new JsonObject { ["content"] = new JsonArray(new JsonObject { ["type"] = "text", ["text"] = "340 pieces" }, new JsonObject { ["type"] = "image", ["mimeType"] = "image/png", ["data"] = "AAA" }) },
                "tools/call" => new JsonObject { ["isError"] = true, ["content"] = new JsonArray(new JsonObject { ["type"] = "text", ["text"] = "customer unknown" }) },
                _ => new JsonObject()
            };
            var answer = new JsonObject { ["jsonrpc"] = "2.0", ["id"] = message["id"]!.DeepClone(), ["result"] = result }.ToJsonString();

            var response = new HttpResponseMessage(HttpStatusCode.OK);
            if (method == "initialize")
                response.Headers.Add("Mcp-Session-Id", "session-" + ++_session);
            response.Content = eventStream
                ? new StringContent($"event: message\ndata: {{\"jsonrpc\":\"2.0\",\"method\":\"notifications/progress\"}}\n\nevent: message\ndata: {answer}\n\n", Encoding.UTF8, "text/event-stream")
                : new StringContent(answer, Encoding.UTF8, "application/json");
            return response;
        }
    }

    private static McpServerConfig Config() => new()
    {
        Id = "shop", Name = "Shop", Url = "https://mcp.example/shop",
        Headers = new() { ["Authorization"] = "Bearer secret-1" }, WithoutApproval = ["stock"]
    };

    private static (McpServerTool Tool, FakeServer Server) Connect(bool eventStream = false)
    {
        var server = new FakeServer(eventStream);
        var config = Config();
        return (new McpServerTool(config, new McpClient(config, new HttpClient(server))), server);
    }

    private static AgentToolContext Context()
        => new("buyer", "inst", new Dictionary<string, object?>(), new NullStateStore(), NullLogger.Instance, CancellationToken.None);

    private static Task<ToolResult> Run(McpServerTool tool, string json)
        => tool.ExecuteAsync(JsonDocument.Parse(json).RootElement, Context(), CancellationToken.None);

    [Theory]
    [InlineData(false)]
    [InlineData(true)]
    public async Task The_tools_of_the_server_become_the_operations_of_one_tool(bool eventStream)
    {
        var (tool, server) = Connect(eventStream);

        (await tool.RefreshAsync(default)).ShouldBeTrue(tool.Problem);

        tool.Name.ShouldBe("mcp_shop");
        tool.RemoteTools.Select(t => t.Name).ShouldBe(["send_invoice", "stock"]);
        tool.InputSchema.GetProperty("properties").GetProperty("action").GetProperty("enum").GetArrayLength().ShouldBe(2);
        server.Received.Select(m => m["method"]!.GetValue<string>()).ShouldBe(["initialize", "notifications/initialized", "tools/list"]);
        server.Sessions[^1].ShouldBe("session-1", "the session the server opened is carried");
        server.Authorization.ShouldAllBe(a => a == "Bearer secret-1");
    }

    [Fact]
    public async Task A_call_returns_the_text_and_names_what_it_left_out()
    {
        var (tool, _) = Connect();
        await tool.RefreshAsync(default);

        var result = await Run(tool, """{"action":"stock","arguments":{"item":"A4"}}""");

        result.Success.ShouldBeTrue(result.ErrorMessage);
        result.Content.ShouldContain("340 pieces");
        result.Content.ShouldContain("[image content omitted: image/png]");
    }

    [Fact]
    public async Task An_error_of_the_server_is_an_error_of_the_tool()
    {
        var (tool, _) = Connect();
        await tool.RefreshAsync(default);

        var result = await Run(tool, """{"action":"send_invoice","arguments":{"customer":"x"}}""");

        result.Success.ShouldBeFalse();
        result.ErrorMessage!.ShouldContain("customer unknown");
    }

    [Fact]
    public async Task A_forgotten_session_is_opened_again()
    {
        var (tool, server) = Connect();
        await tool.RefreshAsync(default);
        server.ForgetSessionOnce = true;

        (await Run(tool, """{"action":"stock"}""")).Success.ShouldBeTrue();

        server.Received.Count(m => m["method"]!.GetValue<string>() == "initialize").ShouldBe(2);
        server.Sessions[^1].ShouldBe("session-2");
    }

    [Fact]
    public async Task Preflight_names_unknown_tools_and_missing_arguments_before_anything_is_staged()
    {
        var (tool, _) = Connect();
        await tool.RefreshAsync(default);

        tool.Preflight(JsonDocument.Parse("""{"action":"delete_all"}""").RootElement, Context())!.ShouldContain("has no tool 'delete_all'");
        tool.Preflight(JsonDocument.Parse("""{"action":"send_invoice","arguments":{}}""").RootElement, Context())!.ShouldContain("customer");
        tool.Preflight(JsonDocument.Parse("""{"action":"send_invoice","arguments":{"customer":"c"}}""").RootElement, Context()).ShouldBeNull();
    }

    [Fact]
    public async Task Foreign_tools_stand_on_approval_unless_trusted_and_their_words_are_cut_short()
    {
        var (tool, _) = Connect();
        await tool.RefreshAsync(default);

        tool.Vorgabe("send_invoice", Context()).ShouldBe(StagingDecision.Approve);
        tool.Vorgabe("stock", Context()).ShouldBeNull("trusted by the human");
        tool.ReadsForeignContent.ShouldBeTrue();
        tool.Description.ShouldContain("the server's own words");
        tool.Description.ShouldNotContain(new string('d', McpServerTool.MaxRemoteDescription + 1));
    }

    [Fact]
    public async Task An_unreachable_server_keeps_its_problem_and_no_tools()
    {
        var config = Config();
        var tool = new McpServerTool(config, new McpClient(config, new HttpClient(new FailingHandler())));

        (await tool.RefreshAsync(default)).ShouldBeFalse();
        tool.Problem.ShouldNotBeNull();
        tool.Description.ShouldContain("Not reachable");
    }

    private sealed class FailingHandler : HttpMessageHandler
    {
        protected override Task<HttpResponseMessage> SendAsync(HttpRequestMessage request, CancellationToken ct)
            => Task.FromResult(new HttpResponseMessage(HttpStatusCode.Unauthorized));
    }

    [Fact]
    public async Task The_connections_follow_the_configuration()
    {
        var registry = new ToolRegistry();
        var servers = new List<McpServerConfig> { Config() };
        await using var connections = new McpConnections(registry, () => servers, NullLogger.Instance, () => new FakeServer());

        await connections.SyncAsync(default);
        registry.Get("mcp_shop").ShouldNotBeNull();
        connections.Find("shop")!.RemoteTools.Count.ShouldBe(2);

        servers.Clear();
        await connections.SyncAsync(default);
        registry.Get("mcp_shop").ShouldBeNull("a removed server is no tool any more");
    }

    [Fact]
    public void The_event_stream_yields_the_answer_with_our_id_only()
    {
        var body = "data: {\"jsonrpc\":\"2.0\",\"id\":7,\"result\":{\"x\":1}}\n\ndata: {\"jsonrpc\":\"2.0\",\"id\":8,\"result\":{\"x\":2}}\n\n";

        McpClient.FromEventStream(body, 8)!["result"]!["x"]!.GetValue<int>().ShouldBe(2);
        McpClient.FromEventStream(body, 9).ShouldBeNull();
    }

    [Fact]
    public void The_headers_of_a_server_are_kept_protected()
    {
        var keys = Path.Combine(Path.GetTempPath(), "clawd-tests", Guid.NewGuid().ToString("N"));
        SecretKeyStore.UseDirectory(keys);
        try
        {
            var instance = new InstanceConfig();
            instance.McpServers.Add(Config());

            ConfigSecrets.Protect(instance);
            instance.McpServers[0].Headers["Authorization"].ShouldNotContain("secret-1");
            ConfigSecrets.Unprotect(instance);
            instance.McpServers[0].Headers["Authorization"].ShouldBe("Bearer secret-1");
        }
        finally
        {
            SecretKeyStore.UseDirectory(null);
            try { Directory.Delete(keys, recursive: true); } catch { /* cleanup is secondary */ }
        }
    }
}

internal sealed class NullStateStore : ClawdDotNet.Core.State.IStateStore
{
    public Task<string?> GetAsync(string key, CancellationToken ct) => Task.FromResult<string?>(null);
    public Task SetAsync(string key, string value, CancellationToken ct) => Task.CompletedTask;
    public Task DeleteAsync(string key, CancellationToken ct) => Task.CompletedTask;
}