334 lines · cpp
1//===----------------------------------------------------------------------===//2//3// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.4// See https://llvm.org/LICENSE.txt for license information.5// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception6//7//===----------------------------------------------------------------------===//8 9#include "ProtocolMCPTestUtilities.h" // IWYU pragma: keep10#include "TestingSupport/Host/JSONTransportTestUtilities.h"11#include "TestingSupport/SubsystemRAII.h"12#include "lldb/Host/FileSystem.h"13#include "lldb/Host/HostInfo.h"14#include "lldb/Host/JSONTransport.h"15#include "lldb/Host/MainLoop.h"16#include "lldb/Host/MainLoopBase.h"17#include "lldb/Host/Socket.h"18#include "lldb/Protocol/MCP/MCPError.h"19#include "lldb/Protocol/MCP/Protocol.h"20#include "lldb/Protocol/MCP/Resource.h"21#include "lldb/Protocol/MCP/Server.h"22#include "lldb/Protocol/MCP/Tool.h"23#include "lldb/Protocol/MCP/Transport.h"24#include "llvm/ADT/StringRef.h"25#include "llvm/Support/Error.h"26#include "llvm/Support/JSON.h"27#include "llvm/Testing/Support/Error.h"28#include "gmock/gmock.h"29#include "gtest/gtest.h"30#include <future>31#include <memory>32#include <optional>33#include <system_error>34 35using namespace llvm;36using namespace lldb;37using namespace lldb_private;38using namespace lldb_private::transport;39using namespace lldb_protocol::mcp;40 41// Flakey, see https://github.com/llvm/llvm-project/issues/152677.42#ifndef _WIN3243 44namespace {45 46template <typename T> Response make_response(T &&result, Id id = 1) {47 return Response{id, std::forward<T>(result)};48}49 50/// Test tool that returns it argument as text.51class TestTool : public Tool {52public:53 using Tool::Tool;54 55 llvm::Expected<CallToolResult> Call(const ToolArguments &args) override {56 std::string argument;57 if (const json::Object *args_obj =58 std::get<json::Value>(args).getAsObject()) {59 if (const json::Value *s = args_obj->get("arguments")) {60 argument = s->getAsString().value_or("");61 }62 }63 64 CallToolResult text_result;65 text_result.content.emplace_back(TextContent{{argument}});66 return text_result;67 }68};69 70class TestResourceProvider : public ResourceProvider {71 using ResourceProvider::ResourceProvider;72 73 std::vector<Resource> GetResources() const override {74 std::vector<Resource> resources;75 76 Resource resource;77 resource.uri = "lldb://foo/bar";78 resource.name = "name";79 resource.description = "description";80 resource.mimeType = "application/json";81 82 resources.push_back(resource);83 return resources;84 }85 86 llvm::Expected<ReadResourceResult>87 ReadResource(llvm::StringRef uri) const override {88 if (uri != "lldb://foo/bar")89 return llvm::make_error<UnsupportedURI>(uri.str());90 91 TextResourceContents contents;92 contents.uri = "lldb://foo/bar";93 contents.mimeType = "application/json";94 contents.text = "foobar";95 96 ReadResourceResult result;97 result.contents.push_back(contents);98 return result;99 }100};101 102/// Test tool that returns an error.103class ErrorTool : public Tool {104public:105 using Tool::Tool;106 107 llvm::Expected<CallToolResult> Call(const ToolArguments &args) override {108 return llvm::createStringError(109 std::error_code(eErrorCodeInternalError, std::generic_category()),110 "error");111 }112};113 114/// Test tool that fails but doesn't return an error.115class FailTool : public Tool {116public:117 using Tool::Tool;118 119 llvm::Expected<CallToolResult> Call(const ToolArguments &args) override {120 CallToolResult text_result;121 text_result.content.emplace_back(TextContent{{"failed"}});122 text_result.isError = true;123 return text_result;124 }125};126 127class TestServer : public Server {128public:129 using Server::Bind;130 using Server::Server;131};132 133using Transport = TestTransport<lldb_protocol::mcp::ProtocolDescriptor>;134 135class ProtocolServerMCPTest : public testing::Test {136public:137 SubsystemRAII<FileSystem, HostInfo, Socket> subsystems;138 139 MainLoop loop;140 lldb_private::MainLoop::ReadHandleUP handles[2];141 142 std::unique_ptr<Transport> to_server;143 MCPBinderUP binder;144 std::unique_ptr<TestServer> server_up;145 146 std::unique_ptr<Transport> to_client;147 MockMessageHandler<lldb_protocol::mcp::ProtocolDescriptor> client;148 149 std::vector<std::string> logged_messages;150 151 /// Runs the MainLoop a single time, executing any pending callbacks.152 void Run() {153 bool addition_succeeded = loop.AddPendingCallback(154 [](MainLoopBase &loop) { loop.RequestTermination(); });155 EXPECT_TRUE(addition_succeeded);156 EXPECT_THAT_ERROR(loop.Run().takeError(), Succeeded());157 }158 159 void SetUp() override {160 std::tie(to_client, to_server) = Transport::createPair();161 162 server_up = std::make_unique<TestServer>(163 "lldb-mcp", "0.1.0",164 [this](StringRef msg) { logged_messages.push_back(msg.str()); });165 binder = server_up->Bind(*to_client);166 auto server_handle = to_server->RegisterMessageHandler(loop, *binder);167 EXPECT_THAT_EXPECTED(server_handle, Succeeded());168 binder->OnError([](llvm::Error error) {169 llvm::errs() << formatv("Server transport error: {0}", error);170 });171 handles[0] = std::move(*server_handle);172 173 auto client_handle = to_client->RegisterMessageHandler(loop, client);174 EXPECT_THAT_EXPECTED(client_handle, Succeeded());175 handles[1] = std::move(*client_handle);176 }177 178 template <typename Result, typename Params>179 Expected<json::Value> Call(StringRef method, const Params ¶ms) {180 std::promise<Response> promised_result;181 Request req =182 lldb_protocol::mcp::Request{/*id=*/1, method.str(), toJSON(params)};183 EXPECT_THAT_ERROR(to_server->Send(req), Succeeded());184 EXPECT_CALL(client, Received(testing::An<const Response &>()))185 .WillOnce(186 [&](const Response &resp) { promised_result.set_value(resp); });187 Run();188 Response resp = promised_result.get_future().get();189 return toJSON(resp);190 }191 192 template <typename Result>193 Expected<json::Value>194 Capture(llvm::unique_function<void(Reply<Result>)> &fn) {195 std::promise<llvm::Expected<Result>> promised_result;196 fn([&promised_result](llvm::Expected<Result> result) {197 promised_result.set_value(std::move(result));198 });199 Run();200 llvm::Expected<Result> result = promised_result.get_future().get();201 if (!result)202 return result.takeError();203 return toJSON(*result);204 }205 206 template <typename Result, typename Params>207 Expected<json::Value>208 Capture(llvm::unique_function<void(const Params &, Reply<Result>)> &fn,209 const Params ¶ms) {210 std::promise<llvm::Expected<Result>> promised_result;211 fn(params, [&promised_result](llvm::Expected<Result> result) {212 promised_result.set_value(std::move(result));213 });214 Run();215 llvm::Expected<Result> result = promised_result.get_future().get();216 if (!result)217 return result.takeError();218 return toJSON(*result);219 }220};221 222template <typename T>223inline testing::internal::EqMatcher<llvm::json::Value> HasJSON(T x) {224 return testing::internal::EqMatcher<llvm::json::Value>(toJSON(x));225}226 227} // namespace228 229TEST_F(ProtocolServerMCPTest, Initialization) {230 EXPECT_THAT_EXPECTED(231 (Call<InitializeResult, InitializeParams>(232 "initialize",233 InitializeParams{/*protocolVersion=*/"2024-11-05",234 /*capabilities=*/{},235 /*clientInfo=*/{"lldb-unit", "0.1.0"}})),236 HasValue(make_response(237 InitializeResult{/*protocolVersion=*/"2024-11-05",238 /*capabilities=*/239 {240 /*supportsToolsList=*/true,241 /*supportsResourcesList=*/true,242 },243 /*serverInfo=*/{"lldb-mcp", "0.1.0"}})));244}245 246TEST_F(ProtocolServerMCPTest, ToolsList) {247 server_up->AddTool(std::make_unique<TestTool>("test", "test tool"));248 249 ToolDefinition test_tool;250 test_tool.name = "test";251 test_tool.description = "test tool";252 test_tool.inputSchema = json::Object{{"type", "object"}};253 254 EXPECT_THAT_EXPECTED(Call<ListToolsResult>("tools/list", Void{}),255 HasValue(make_response(ListToolsResult{{test_tool}})));256}257 258TEST_F(ProtocolServerMCPTest, ResourcesList) {259 server_up->AddResourceProvider(std::make_unique<TestResourceProvider>());260 261 EXPECT_THAT_EXPECTED(Call<ListResourcesResult>("resources/list", Void{}),262 HasValue(make_response(ListResourcesResult{{263 {264 /*uri=*/"lldb://foo/bar",265 /*name=*/"name",266 /*description=*/"description",267 /*mimeType=*/"application/json",268 },269 }})));270}271 272TEST_F(ProtocolServerMCPTest, ToolsCall) {273 server_up->AddTool(std::make_unique<TestTool>("test", "test tool"));274 275 EXPECT_THAT_EXPECTED(276 (Call<CallToolResult, CallToolParams>("tools/call",277 CallToolParams{278 /*name=*/"test",279 /*arguments=*/280 json::Object{281 {"arguments", "foo"},282 {"debugger_id", 0},283 },284 })),285 HasValue(make_response(CallToolResult{{{/*text=*/"foo"}}})));286}287 288TEST_F(ProtocolServerMCPTest, ToolsCallError) {289 server_up->AddTool(std::make_unique<ErrorTool>("error", "error tool"));290 291 EXPECT_THAT_EXPECTED((Call<CallToolResult, CallToolParams>(292 "tools/call", CallToolParams{293 /*name=*/"error",294 /*arguments=*/295 json::Object{296 {"arguments", "foo"},297 {"debugger_id", 0},298 },299 })),300 HasValue(make_response(lldb_protocol::mcp::Error{301 eErrorCodeInternalError, "error"})));302}303 304TEST_F(ProtocolServerMCPTest, ToolsCallFail) {305 server_up->AddTool(std::make_unique<FailTool>("fail", "fail tool"));306 307 EXPECT_THAT_EXPECTED((Call<CallToolResult, CallToolParams>(308 "tools/call", CallToolParams{309 /*name=*/"fail",310 /*arguments=*/311 json::Object{312 {"arguments", "foo"},313 {"debugger_id", 0},314 },315 })),316 HasValue(make_response(CallToolResult{317 {{/*text=*/"failed"}},318 /*isError=*/true,319 })));320}321 322TEST_F(ProtocolServerMCPTest, NotificationInitialized) {323 EXPECT_THAT_ERROR(to_server->Send(lldb_protocol::mcp::Notification{324 "notifications/initialized",325 std::nullopt,326 }),327 Succeeded());328 Run();329 EXPECT_THAT(logged_messages,330 testing::Contains("MCP initialization complete"));331}332 333#endif334