BareGit
Login
agent-smith / master / tests / test_mcp_jsonrpc.cpp 
#include <atomic>
#include <thread>

#include <agent_smith/agent_smith.hpp>
#include <agent_smith/detail/gather_all.hpp>
#include <gmock/gmock.h>
#include <gtest/gtest.h>
#include <mw/http_client_async_mock.hpp>

namespace agent_smith
{
namespace
{
using namespace std::chrono_literals;
using Json = nlohmann::json;

mcp::ServerConfig server(std::string mode = "normal",
                         std::string alias = "test")
{
    return {.alias = std::move(alias),
            .transport = mcp::StdioServerConfig{{AGENT_SMITH_TEST_PYTHON,
                                                 AGENT_SMITH_MCP_FIXTURE,
                                                 std::move(mode)}},
            .request_timeout = 2s};
}

TEST(McpJsonrpcTest, InitializesDiscoversAndRendersText)
{
    Runtime runtime;
    std::atomic<int> notifications{0};
    auto result = syncWait(mcp::Client::connect(
        runtime, server(),
        [&notifications](const mcp::Event& event)
        {
            if(std::holds_alternative<mcp::Notification>(event))
            {
                ++notifications;
            }
        }));
    ASSERT_TRUE(result) << result.error().msg();
    auto& client = **result;
    EXPECT_EQ(client.info().name, "agent-smith-example");
    ASSERT_EQ(client.tools().size(), 1U);
    EXPECT_EQ(client.tools()[0].name, "add");
    ASSERT_EQ(client.resources().size(), 1U);
    EXPECT_EQ(client.resources()[0].mime_type, "text/plain");
    ASSERT_EQ(client.prompts().size(), 1U);
    auto sum = syncWait(client.callTool("add", {{"a", 2}, {"b", 3}}, 1s));
    ASSERT_TRUE(sum) << sum.error().msg();
    EXPECT_EQ(sum->content, "5");
    EXPECT_EQ(sum->structured["sum"], 5);
    auto resource = syncWait(client.readResource("example://help"));
    ASSERT_TRUE(resource);
    EXPECT_EQ(*resource, "The add tool adds two numbers.");
    auto prompt = syncWait(client.getPrompt("greet", {{"name", "Alice"}}));
    ASSERT_TRUE(prompt);
    ASSERT_EQ(prompt->messages.size(), 1U);
    EXPECT_EQ(std::get<UserMessage>(prompt->messages[0]).content,
              "Say hello to Alice.");
    EXPECT_GE(notifications.load(), 1);
    EXPECT_TRUE(syncWait(client.shutdown()));
    EXPECT_TRUE(syncWait(client.shutdown()));
}

TEST(McpJsonrpcTest, RoutesOutOfOrderResponsesConcurrently)
{
    Runtime runtime;
    auto connected = syncWait(mcp::Client::connect(runtime, server()));
    ASSERT_TRUE(connected);
    std::vector<mw::Task<mw::E<ToolResult>>> tasks;
    tasks.push_back(
        (*connected)->callTool("ordered", {{"label", "first"}}, 1s));
    tasks.push_back(
        (*connected)->callTool("ordered", {{"label", "second"}}, 1s));
    auto results = syncWait(internal::gatherAll(std::move(tasks)));
    ASSERT_EQ(results.size(), 2U);
    ASSERT_TRUE(results[0]);
    ASSERT_TRUE(results[1]);
    EXPECT_EQ(results[0]->content, "first");
    EXPECT_EQ(results[1]->content, "second");
}

TEST(McpJsonrpcTest, DistinguishesToolFailureAndRpcError)
{
    Runtime runtime;
    auto connected = syncWait(mcp::Client::connect(runtime, server()));
    ASSERT_TRUE(connected);
    auto failure = syncWait((*connected)->callTool("add", Json::object(), 1s));
    ASSERT_TRUE(failure);
    EXPECT_TRUE(failure->is_error);
    auto remote =
        syncWait((*connected)->callTool("rpc_error", Json::object(), 1s));
    ASSERT_FALSE(remote);
    const auto* error = remote.error().as<McpError>();
    ASSERT_NE(error, nullptr);
    EXPECT_EQ(error->rpc_code, -32001);
    EXPECT_EQ(error->server, "test");
    EXPECT_EQ(error->data["reason"], "test");
}

TEST(McpJsonrpcTest, TimesOutAndNotifiesCancellation)
{
    Runtime runtime;
    auto connected = syncWait(mcp::Client::connect(runtime, server()));
    ASSERT_TRUE(connected);
    auto hung = syncWait((*connected)->callTool("hang", Json::object(), 30ms));
    ASSERT_FALSE(hung);
    EXPECT_NE(hung.error().as<TimeoutError>(), nullptr);
    auto observed =
        syncWait((*connected)->callTool("cancelled", Json::object(), 1s));
    ASSERT_TRUE(observed);
    EXPECT_EQ(Json::parse(observed->content).size(), 1U);
    auto next =
        syncWait((*connected)->callTool("add", {{"a", 1}, {"b", 4}}, 1s));
    ASSERT_TRUE(next);
    EXPECT_EQ(next->content, "5");
}

TEST(McpJsonrpcTest, CancelsAnInFlightRequest)
{
    Runtime runtime;
    auto connected = syncWait(mcp::Client::connect(runtime, server()));
    ASSERT_TRUE(connected);
    CancellationSource cancel;
    std::jthread thread(
        [cancel]() mutable
        {
            std::this_thread::sleep_for(30ms);
            cancel.cancel();
        });
    auto hung = syncWait(
        (*connected)->callTool("hang", Json::object(), 1s, cancel.token()));
    ASSERT_FALSE(hung);
    EXPECT_NE(hung.error().as<CancelledError>(), nullptr);
    auto observed =
        syncWait((*connected)->callTool("cancelled", Json::object(), 1s));
    ASSERT_TRUE(observed);
    EXPECT_EQ(Json::parse(observed->content).size(), 1U);
}

TEST(McpJsonrpcTest, FiltersToolsWithDenyTakingPrecedence)
{
    Runtime runtime;
    auto config = server("pagination");
    config.allow_tools = {"add", "second"};
    config.deny_tools = {"add"};
    auto connected = syncWait(mcp::Client::connect(runtime, std::move(config)));
    ASSERT_TRUE(connected);
    ASSERT_EQ((*connected)->tools().size(), 1U);
    EXPECT_EQ((*connected)->tools()[0].name, "second");
}

TEST(McpJsonrpcTest, OptionalMethodsMayBeMissingButMalformedListsFail)
{
    Runtime runtime;
    auto connected =
        syncWait(mcp::Client::connect(runtime, server("optional_missing")));
    ASSERT_TRUE(connected);
    EXPECT_TRUE((*connected)->resources().empty());
    EXPECT_TRUE((*connected)->prompts().empty());
    auto malformed =
        syncWait(mcp::Client::connect(runtime, server("optional_bad")));
    ASSERT_FALSE(malformed);
    EXPECT_NE(malformed.error().as<McpError>(), nullptr);
}

TEST(McpJsonrpcTest, RequiredToolsMustBeAvailable)
{
    Runtime runtime;
    EXPECT_FALSE(syncWait(mcp::Client::connect(runtime, server("no_tools"))));
    auto config = server("no_tools");
    config.tools_required = false;
    auto connected = syncWait(mcp::Client::connect(runtime, std::move(config)));
    ASSERT_TRUE(connected);
    EXPECT_TRUE((*connected)->tools().empty());
}

TEST(McpJsonrpcTest, RejectsInvalidProtocolAndPagination)
{
    Runtime runtime;
    for(const auto& mode : {"bad_version", "bad_json", "cycle"})
    {
        auto connected = syncWait(mcp::Client::connect(runtime, server(mode)));
        ASSERT_FALSE(connected) << mode;
        EXPECT_NE(connected.error().as<McpError>(), nullptr);
    }
}

TEST(McpJsonrpcTest, InitializationHasDeadlineAndLaunchErrorsAreImmediate)
{
    Runtime runtime;
    auto config = server("hang_init");
    config.request_timeout = 100ms;
    auto connected = syncWait(mcp::Client::connect(runtime, std::move(config)));
    ASSERT_FALSE(connected);
    EXPECT_NE(connected.error().as<TimeoutError>(), nullptr);
    config = server();
    config.transport = mcp::StdioServerConfig{{"/no/such/mcp/executable"}};
    auto failed = syncWait(mcp::Client::connect(runtime, std::move(config)));
    ASSERT_FALSE(failed);
    ASSERT_NE(failed.error().as<McpError>(), nullptr);
    EXPECT_TRUE(failed.error().as<McpError>()->transport);
}

TEST(McpJsonrpcTest, RejectsUnrepresentableContentAndMalformedEnvelopes)
{
    Runtime runtime;
    auto connected = syncWait(mcp::Client::connect(runtime, server("binary")));
    ASSERT_TRUE(connected);
    EXPECT_FALSE(syncWait((*connected)->readResource("example://help")));
    EXPECT_FALSE(
        syncWait((*connected)->callTool("multimodal", Json::object(), 1s)));
    EXPECT_FALSE(
        syncWait((*connected)->callTool("bad_envelope", Json::object(), 1s)));
    auto prompt = syncWait(mcp::Client::connect(runtime, server("bad_prompt")));
    ASSERT_TRUE(prompt);
    EXPECT_FALSE(syncWait((*prompt)->getPrompt("greet", Json::object())));
}

TEST(McpJsonrpcTest, DoesNotReplayWrittenToolCallsAndReconnectsNextOperation)
{
    Runtime runtime;
    std::atomic<int> reconnects{0};
    auto connected = syncWait(mcp::Client::connect(
        runtime, server(),
        [&reconnects](const mcp::Event& event)
        {
            const auto* transition = std::get_if<mcp::LifecycleEvent>(&event);
            if(transition &&
               transition->kind == mcp::LifecycleEventKind::RECONNECTING)
            {
                ++reconnects;
            }
        }));
    ASSERT_TRUE(connected);
    auto failed =
        syncWait((*connected)->callTool("disconnect", Json::object(), 1s));
    ASSERT_FALSE(failed);
    EXPECT_EQ(reconnects.load(), 0);
    auto next =
        syncWait((*connected)->callTool("add", {{"a", 1}, {"b", 2}}, 1s));
    ASSERT_TRUE(next) << next.error().msg();
    EXPECT_EQ(next->content, "3");
    EXPECT_EQ(reconnects.load(), 1);
}

TEST(McpJsonrpcTest, PrefixesAdaptersAndRetainsClientsAcrossServers)
{
    Runtime runtime;
    ToolRegistry registry;
    for(const auto& alias : {"one", "two"})
    {
        auto connected =
            syncWait(mcp::Client::connect(runtime, server("normal", alias)));
        ASSERT_TRUE(connected);
        auto client = std::shared_ptr<mcp::Client>(std::move(*connected));
        for(auto& tool : mcp::adaptTools(client))
        {
            EXPECT_TRUE(registry.add(std::move(tool)));
        }
    }
    EXPECT_EQ(registry.size(), 4U);
    ASSERT_NE(registry.find("mcp_one__add"), nullptr);
    auto result = syncWait(registry.find("mcp_one__add")
                               ->invoke({{"a", 8}, {"b", 9}},
                                        {.budget = 1s, .runtime = &runtime}));
    ASSERT_TRUE(result);
    EXPECT_EQ(result->content, "17");
    auto resource = syncWait(registry.find("mcp_two__read_resource")
                                 ->invoke({{"uri", "example://help"}},
                                          {.budget = 1s, .runtime = &runtime}));
    ASSERT_TRUE(resource);
    EXPECT_EQ(resource->content, "The add tool adds two numbers.");
}

TEST(McpJsonrpcTest, ShutdownReapsUncooperativeServer)
{
    Runtime runtime;
    auto connected =
        syncWait(mcp::Client::connect(runtime, server("stubborn")));
    ASSERT_TRUE(connected);
    const auto start = std::chrono::steady_clock::now();
    EXPECT_TRUE(syncWait((*connected)->shutdown()));
    EXPECT_LT(std::chrono::steady_clock::now() - start, 2s);
    EXPECT_FALSE(
        syncWait((*connected)->callTool("add", {{"a", 1}, {"b", 2}}, 1s)));
}

TEST(McpJsonrpcTest, ReconnectsReadOnlyOperationsAtMostOnce)
{
    Runtime runtime;
    std::atomic<int> reconnects{0};
    auto connected = syncWait(mcp::Client::connect(
        runtime, server("resource_disconnect"),
        [&reconnects](const mcp::Event& event)
        {
            const auto* transition = std::get_if<mcp::LifecycleEvent>(&event);
            if(transition &&
               transition->kind == mcp::LifecycleEventKind::RECONNECTING)
            {
                ++reconnects;
            }
        }));
    ASSERT_TRUE(connected);
    EXPECT_FALSE(syncWait((*connected)->readResource("example://help")));
    EXPECT_EQ(reconnects.load(), 1);
}

TEST(McpJsonrpcTest, CanDisableReconnectAndCancelBeforeSending)
{
    Runtime runtime;
    auto config = server();
    config.reconnect_on_transport_error = false;
    auto connected = syncWait(mcp::Client::connect(runtime, std::move(config)));
    ASSERT_TRUE(connected);
    CancellationSource cancel;
    cancel.cancel();
    auto cancelled = syncWait(
        (*connected)
            ->callTool("add", {{"a", 1}, {"b", 2}}, 1s, cancel.token()));
    ASSERT_FALSE(cancelled);
    EXPECT_NE(cancelled.error().as<CancelledError>(), nullptr);
    EXPECT_FALSE(
        syncWait((*connected)->callTool("disconnect", Json::object(), 1s)));
    EXPECT_FALSE(
        syncWait((*connected)->callTool("add", {{"a", 1}, {"b", 2}}, 1s)));
}

TEST(McpJsonrpcTest, AgentUsesMcpSchemasAndRecordsToolResults)
{
    using ::testing::_;
    using ::testing::Invoke;
    Runtime runtime;
    mw::HTTPSessionAsyncMock http;
    LlmClient client(runtime, http,
                     {.endpoint = *mw::URL::fromStr("https://example.test"),
                      .model = "test",
                      .max_retries = 0});
    auto memory = std::make_shared<FullMemory>();
    Agent agent(runtime, client, memory, {.stream_default = false});
    std::atomic<int> connected_events{0};
    agent.hooks().on_mcp_event = [&connected_events](const mcp::Event& event)
    {
        if(const auto* transition = std::get_if<mcp::LifecycleEvent>(&event);
           transition && transition->kind == mcp::LifecycleEventKind::CONNECTED)
        {
            ++connected_events;
        }
    };
    ASSERT_TRUE(syncWait(agent.connectMcp(server())));
    EXPECT_EQ(connected_events.load(), 1);
    int requests = 0;
    EXPECT_CALL(http, post(_))
        .Times(2)
        .WillRepeatedly(Invoke(
            [&requests](mw::HTTPRequest request)
            {
                const auto body = Json::parse(request.request_data);
                EXPECT_EQ(body["tools"][0]["function"]["name"],
                          "mcp_test__add");
                EXPECT_EQ(body["tools"][1]["function"]["name"],
                          "mcp_test__read_resource");
                Json message{{"role", "assistant"}, {"content", "sum is 5"}};
                if(++requests == 1)
                {
                    message["tool_calls"] =
                        Json::array({{{"id", "add-call"},
                                      {"type", "function"},
                                      {"function",
                                       {{"name", "mcp_test__add"},
                                        {"arguments", R"({"a":2,"b":3})"}}}}});
                }
                else
                {
                    EXPECT_EQ(body["messages"].back()["role"], "tool");
                    EXPECT_EQ(body["messages"].back()["content"], "5");
                }
                return mw::HTTPSessionAsyncMock::complete(mw::HTTPResponse(
                    200,
                    Json{{"choices", Json::array({{{"message", message}}})}}
                        .dump()));
            }));
    auto result = syncWait(agent.send("add two and three"));
    ASSERT_TRUE(result) << result.error().msg();
    EXPECT_EQ(result->content, "sum is 5");
    EXPECT_EQ(memory->size(), 4U);
}

TEST(McpJsonrpcTest, AgentRejectsCollisionsWithoutPartialRegistration)
{
    Runtime runtime;
    LlmClient client(runtime,
                     {.endpoint = *mw::URL::fromStr("https://example.test"),
                      .model = "test"});
    auto connected = syncWait(mcp::Client::connect(runtime, server()));
    ASSERT_TRUE(connected);
    auto remote = std::shared_ptr<mcp::Client>(std::move(*connected));
    {
        Agent agent(runtime, client, std::make_shared<FullMemory>());
        EXPECT_FALSE(agent.addMcpClient(nullptr));
        ASSERT_TRUE(agent.addTool(makeTool(
            ToolSchema{"mcp_test__read_resource",
                       "occupied",
                       {{"type", "object"}, {"properties", Json::object()}}},
            [](Json, ToolContext) { return "local"; })));
        EXPECT_FALSE(agent.addMcpClient(remote));
        // A failed batch must not leave the earlier add adapter registered.
        EXPECT_TRUE(agent.addTool(makeTool(
            ToolSchema{"mcp_test__add",
                       "local add",
                       {{"type", "object"}, {"properties", Json::object()}}},
            [](Json, ToolContext) { return "local"; })));
    }
    {
        Agent agent(runtime, client, std::make_shared<FullMemory>());
        ASSERT_TRUE(agent.addMcpClient(remote));
        EXPECT_FALSE(agent.addMcpClient(remote));
    }
    // Hooks bound by the destroyed agent must not retain a dangling this.
    EXPECT_TRUE(syncWait(remote->shutdown()));
}

TEST(McpJsonrpcTest, CancelsBlockedWritesWithoutWaitingForRequestDeadline)
{
    Runtime runtime;
    auto connected =
        syncWait(mcp::Client::connect(runtime, server("stop_reading")));
    ASSERT_TRUE(connected);
    CancellationSource cancel;
    std::jthread thread(
        [cancel]() mutable
        {
            std::this_thread::sleep_for(30ms);
            cancel.cancel();
        });
    const auto start = std::chrono::steady_clock::now();
    auto result = syncWait(
        (*connected)
            ->callTool("add", {{"padding", std::string(1024 * 1024, 'x')}}, 2s,
                       cancel.token()));
    ASSERT_FALSE(result);
    EXPECT_NE(result.error().as<CancelledError>(), nullptr);
    EXPECT_LT(std::chrono::steady_clock::now() - start, 1s);
}

TEST(McpJsonrpcTest, AppliesWriteDeadlineWhenServerStopsReading)
{
    Runtime runtime;
    auto connected =
        syncWait(mcp::Client::connect(runtime, server("stop_reading")));
    ASSERT_TRUE(connected);
    const auto start = std::chrono::steady_clock::now();
    auto result = syncWait(
        (*connected)
            ->callTool("add", {{"padding", std::string(1024 * 1024, 'x')}},
                       40ms));
    ASSERT_FALSE(result);
    EXPECT_NE(result.error().as<TimeoutError>(), nullptr);
    EXPECT_LT(std::chrono::steady_clock::now() - start, 1s);
}

TEST(McpJsonrpcTest, RejectsInvalidConfigurationAndDeniedDirectCalls)
{
    Runtime runtime;
    auto config = server();
    config.alias = "bad alias";
    EXPECT_FALSE(syncWait(mcp::Client::connect(runtime, config)));
    config = server();
    config.request_timeout = 0ms;
    EXPECT_FALSE(syncWait(mcp::Client::connect(runtime, config)));
    config = server();
    config.transport =
        mcp::StdioServerConfig{{"python3", std::string("bad\0arg", 7)}};
    EXPECT_FALSE(syncWait(mcp::Client::connect(runtime, config)));
    config = server();
    config.deny_tools = {"add"};
    auto connected = syncWait(mcp::Client::connect(runtime, config));
    ASSERT_TRUE(connected);
    EXPECT_TRUE((*connected)->tools().empty());
    EXPECT_FALSE(
        syncWait((*connected)->callTool("add", {{"a", 1}, {"b", 2}}, 1s)));
}
} // namespace
} // namespace agent_smith