152 lines · cpp
1//===- ProtocolServerMCP.cpp ----------------------------------------------===//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 "ProtocolServerMCP.h"10#include "Resource.h"11#include "Tool.h"12#include "lldb/Core/PluginManager.h"13#include "lldb/Protocol/MCP/Server.h"14#include "lldb/Utility/LLDBLog.h"15#include "lldb/Utility/Log.h"16#include "llvm/ADT/StringExtras.h"17#include "llvm/Support/Error.h"18#include "llvm/Support/Threading.h"19#include <thread>20 21using namespace lldb_private;22using namespace lldb_private::mcp;23using namespace lldb_protocol::mcp;24using namespace llvm;25 26LLDB_PLUGIN_DEFINE(ProtocolServerMCP)27 28static constexpr llvm::StringLiteral kName = "lldb-mcp";29static constexpr llvm::StringLiteral kVersion = "0.1.0";30 31ProtocolServerMCP::ProtocolServerMCP() : ProtocolServer() {}32 33ProtocolServerMCP::~ProtocolServerMCP() { llvm::consumeError(Stop()); }34 35void ProtocolServerMCP::Initialize() {36 PluginManager::RegisterPlugin(GetPluginNameStatic(),37 GetPluginDescriptionStatic(), CreateInstance);38}39 40void ProtocolServerMCP::Terminate() {41 if (llvm::Error error = ProtocolServer::Terminate())42 LLDB_LOG_ERROR(GetLog(LLDBLog::Host), std::move(error), "{0}");43 PluginManager::UnregisterPlugin(CreateInstance);44}45 46lldb::ProtocolServerUP ProtocolServerMCP::CreateInstance() {47 return std::make_unique<ProtocolServerMCP>();48}49 50llvm::StringRef ProtocolServerMCP::GetPluginDescriptionStatic() {51 return "MCP Server.";52}53 54void ProtocolServerMCP::Extend(lldb_protocol::mcp::Server &server) const {55 server.AddTool(56 std::make_unique<CommandTool>("command", "Run an lldb command."));57 server.AddTool(std::make_unique<DebuggerListTool>(58 "debugger_list", "List debugger instances with their debugger_id."));59 server.AddResourceProvider(std::make_unique<DebuggerResourceProvider>());60}61 62void ProtocolServerMCP::AcceptCallback(std::unique_ptr<Socket> socket) {63 Log *log = GetLog(LLDBLog::Host);64 std::string client_name = llvm::formatv("client_{0}", ++m_client_count);65 LLDB_LOG(log, "New MCP client connected: {0}", client_name);66 67 lldb::IOObjectSP io_sp = std::move(socket);68 auto transport_up = std::make_unique<lldb_protocol::mcp::Transport>(69 io_sp, io_sp, [client_name](llvm::StringRef message) {70 LLDB_LOG(GetLog(LLDBLog::Host), "{0}: {1}", client_name, message);71 });72 73 if (auto error = m_server->Accept(m_loop, std::move(transport_up)))74 LLDB_LOG_ERROR(log, std::move(error), "{0}:");75}76 77llvm::Error ProtocolServerMCP::Start(ProtocolServer::Connection connection) {78 std::lock_guard<std::mutex> guard(m_mutex);79 80 if (m_running)81 return llvm::createStringError("the MCP server is already running");82 83 Status status;84 m_listener = Socket::Create(connection.protocol, status);85 if (status.Fail())86 return status.takeError();87 88 status = m_listener->Listen(connection.name, /*backlog=*/5);89 if (status.Fail())90 return status.takeError();91 92 auto handles =93 m_listener->Accept(m_loop, std::bind(&ProtocolServerMCP::AcceptCallback,94 this, std::placeholders::_1));95 if (llvm::Error error = handles.takeError())96 return error;97 98 auto listening_uris = m_listener->GetListeningConnectionURI();99 if (listening_uris.empty())100 return createStringError("failed to get listening connections");101 std::string address =102 llvm::join(m_listener->GetListeningConnectionURI(), ", ");103 104 ServerInfo info{listening_uris[0]};105 llvm::Expected<ServerInfoHandle> server_info_handle = ServerInfo::Write(info);106 if (!server_info_handle)107 return server_info_handle.takeError();108 109 m_client_count = 0;110 m_server = std::make_unique<lldb_protocol::mcp::Server>(111 std::string(kName), std::string(kVersion), [](StringRef message) {112 LLDB_LOG(GetLog(LLDBLog::Host), "MCP Server: {0}", message);113 });114 Extend(*m_server);115 116 m_running = true;117 m_server_info_handle = std::move(*server_info_handle);118 m_accept_handles = std::move(*handles);119 m_loop_thread = std::thread([this] {120 llvm::set_thread_name("protocol-server.mcp");121 m_loop.Run();122 });123 124 return llvm::Error::success();125}126 127llvm::Error ProtocolServerMCP::Stop() {128 {129 std::lock_guard<std::mutex> guard(m_mutex);130 if (!m_running)131 return createStringError("the MCP sever is not running");132 m_running = false;133 }134 135 // Stop the main loop.136 bool addition_succeeded = m_loop.AddPendingCallback(137 [](lldb_private::MainLoopBase &loop) { loop.RequestTermination(); });138 139 // Wait for the main loop to exit, but not if we didn't succeed in inserting140 // our pending callback or we'll wait forever.141 if (addition_succeeded && m_loop_thread.joinable())142 m_loop_thread.join();143 144 m_accept_handles.clear();145 146 m_server.reset(nullptr);147 m_server_info_handle.Remove();148 m_listener.reset();149 150 return llvm::Error::success();151}152