250 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 "lldb/Protocol/MCP/Server.h"10#include "lldb/Host/File.h"11#include "lldb/Host/FileSystem.h"12#include "lldb/Host/HostInfo.h"13#include "lldb/Protocol/MCP/MCPError.h"14#include "lldb/Protocol/MCP/Protocol.h"15#include "lldb/Protocol/MCP/Transport.h"16#include "llvm/ADT/SmallString.h"17#include "llvm/Support/FileSystem.h"18#include "llvm/Support/JSON.h"19#include "llvm/Support/Signals.h"20 21using namespace llvm;22using namespace lldb_private;23using namespace lldb_protocol::mcp;24 25ServerInfoHandle::ServerInfoHandle(StringRef filename) : m_filename(filename) {26 if (!m_filename.empty())27 sys::RemoveFileOnSignal(m_filename);28}29 30ServerInfoHandle::~ServerInfoHandle() { Remove(); }31 32ServerInfoHandle::ServerInfoHandle(ServerInfoHandle &&other) {33 *this = std::move(other);34}35 36ServerInfoHandle &37ServerInfoHandle::operator=(ServerInfoHandle &&other) noexcept {38 m_filename = std::move(other.m_filename);39 return *this;40}41 42void ServerInfoHandle::Remove() {43 if (m_filename.empty())44 return;45 46 sys::fs::remove(m_filename);47 sys::DontRemoveFileOnSignal(m_filename);48 m_filename.clear();49}50 51json::Value lldb_protocol::mcp::toJSON(const ServerInfo &SM) {52 return json::Object{{"connection_uri", SM.connection_uri}};53}54 55bool lldb_protocol::mcp::fromJSON(const json::Value &V, ServerInfo &SM,56 json::Path P) {57 json::ObjectMapper O(V, P);58 return O && O.map("connection_uri", SM.connection_uri);59}60 61Expected<ServerInfoHandle> ServerInfo::Write(const ServerInfo &info) {62 std::string buf = formatv("{0}", toJSON(info)).str();63 size_t num_bytes = buf.size();64 65 FileSpec user_lldb_dir = HostInfo::GetUserLLDBDir();66 67 Status error(sys::fs::create_directory(user_lldb_dir.GetPath()));68 if (error.Fail())69 return error.takeError();70 71 FileSpec mcp_registry_entry_path = user_lldb_dir.CopyByAppendingPathComponent(72 formatv("lldb-mcp-{0}.json", getpid()).str());73 74 const File::OpenOptions flags = File::eOpenOptionWriteOnly |75 File::eOpenOptionCanCreate |76 File::eOpenOptionTruncate;77 Expected<lldb::FileUP> file =78 FileSystem::Instance().Open(mcp_registry_entry_path, flags);79 if (!file)80 return file.takeError();81 if (llvm::Error error = (*file)->Write(buf.data(), num_bytes).takeError())82 return error;83 return ServerInfoHandle{mcp_registry_entry_path.GetPath()};84}85 86Expected<std::vector<ServerInfo>> ServerInfo::Load() {87 namespace path = llvm::sys::path;88 FileSpec user_lldb_dir = HostInfo::GetUserLLDBDir();89 FileSystem &fs = FileSystem::Instance();90 std::error_code EC;91 vfs::directory_iterator it = fs.DirBegin(user_lldb_dir, EC);92 vfs::directory_iterator end;93 std::vector<ServerInfo> infos;94 for (; it != end && !EC; it.increment(EC)) {95 auto &entry = *it;96 auto path = entry.path();97 auto name = path::filename(path);98 if (!name.starts_with("lldb-mcp-") || !name.ends_with(".json"))99 continue;100 101 auto buffer = fs.CreateDataBuffer(path);102 auto info = json::parse<ServerInfo>(toStringRef(buffer->GetData()));103 if (!info)104 return info.takeError();105 106 infos.emplace_back(std::move(*info));107 }108 109 return infos;110}111 112Server::Server(std::string name, std::string version, LogCallback log_callback)113 : m_name(std::move(name)), m_version(std::move(version)),114 m_log_callback(std::move(log_callback)) {}115 116void Server::AddTool(std::unique_ptr<Tool> tool) {117 if (!tool)118 return;119 m_tools[tool->GetName()] = std::move(tool);120}121 122void Server::AddResourceProvider(123 std::unique_ptr<ResourceProvider> resource_provider) {124 if (!resource_provider)125 return;126 m_resource_providers.push_back(std::move(resource_provider));127}128 129MCPBinderUP Server::Bind(MCPTransport &transport) {130 MCPBinderUP binder_up = std::make_unique<MCPBinder>(transport);131 binder_up->Bind<InitializeResult, InitializeParams>(132 "initialize", &Server::InitializeHandler, this);133 binder_up->Bind<ListToolsResult, void>("tools/list",134 &Server::ToolsListHandler, this);135 binder_up->Bind<CallToolResult, CallToolParams>(136 "tools/call", &Server::ToolsCallHandler, this);137 binder_up->Bind<ListResourcesResult, void>(138 "resources/list", &Server::ResourcesListHandler, this);139 binder_up->Bind<ReadResourceResult, ReadResourceParams>(140 "resources/read", &Server::ResourcesReadHandler, this);141 binder_up->Bind<void>("notifications/initialized",142 [this]() { Log("MCP initialization complete"); });143 return binder_up;144}145 146llvm::Error Server::Accept(MainLoop &loop, MCPTransportUP transport) {147 MCPBinderUP binder = Bind(*transport);148 MCPTransport *transport_ptr = transport.get();149 binder->OnDisconnect([this, transport_ptr]() {150 assert(m_instances.find(transport_ptr) != m_instances.end() &&151 "Client not found in m_instances");152 m_instances.erase(transport_ptr);153 });154 binder->OnError([this](llvm::Error err) {155 Logv("Transport error: {0}", llvm::toString(std::move(err)));156 });157 158 auto handle = transport->RegisterMessageHandler(loop, *binder);159 if (!handle)160 return handle.takeError();161 162 m_instances[transport_ptr] =163 Client{std::move(*handle), std::move(transport), std::move(binder)};164 return llvm::Error::success();165}166 167Expected<InitializeResult>168Server::InitializeHandler(const InitializeParams &request) {169 InitializeResult result;170 result.protocolVersion = mcp::kProtocolVersion;171 result.capabilities = GetCapabilities();172 result.serverInfo.name = m_name;173 result.serverInfo.version = m_version;174 return result;175}176 177llvm::Expected<ListToolsResult> Server::ToolsListHandler() {178 ListToolsResult result;179 for (const auto &tool : m_tools)180 result.tools.emplace_back(tool.second->GetDefinition());181 182 return result;183}184 185llvm::Expected<CallToolResult>186Server::ToolsCallHandler(const CallToolParams ¶ms) {187 llvm::StringRef tool_name = params.name;188 if (tool_name.empty())189 return llvm::createStringError("no tool name");190 191 auto it = m_tools.find(tool_name);192 if (it == m_tools.end())193 return llvm::createStringError(llvm::formatv("no tool \"{0}\"", tool_name));194 195 ToolArguments tool_args;196 if (params.arguments)197 tool_args = *params.arguments;198 199 llvm::Expected<CallToolResult> text_result = it->second->Call(tool_args);200 if (!text_result)201 return text_result.takeError();202 203 return text_result;204}205 206llvm::Expected<ListResourcesResult> Server::ResourcesListHandler() {207 ListResourcesResult result;208 for (std::unique_ptr<ResourceProvider> &resource_provider_up :209 m_resource_providers)210 for (const Resource &resource : resource_provider_up->GetResources())211 result.resources.push_back(resource);212 213 return result;214}215 216Expected<ReadResourceResult>217Server::ResourcesReadHandler(const ReadResourceParams ¶ms) {218 StringRef uri_str = params.uri;219 if (uri_str.empty())220 return createStringError("no resource uri");221 222 for (std::unique_ptr<ResourceProvider> &resource_provider_up :223 m_resource_providers) {224 Expected<ReadResourceResult> result =225 resource_provider_up->ReadResource(uri_str);226 if (result.errorIsA<UnsupportedURI>()) {227 consumeError(result.takeError());228 continue;229 }230 if (!result)231 return result.takeError();232 233 return *result;234 }235 236 return make_error<MCPError>(237 formatv("no resource handler for uri: {0}", uri_str).str(),238 MCPError::kResourceNotFound);239}240 241ServerCapabilities Server::GetCapabilities() {242 lldb_protocol::mcp::ServerCapabilities capabilities;243 capabilities.supportsToolsList = true;244 capabilities.supportsResourcesList = true;245 // FIXME: Support sending notifications when a debugger/target are246 // added/removed.247 capabilities.supportsResourcesSubscribe = false;248 return capabilities;249}250