source/Protocol/MCP/Server.cpp - llvm-project/lldb - Git at Google

 //===----------------------------------------------------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//

 #include "lldb/Protocol/MCP/Server.h"
 #include "lldb/Host/File.h"
 #include "lldb/Host/FileSystem.h"
 #include "lldb/Host/HostInfo.h"
 #include "lldb/Host/JSONTransport.h"
 #include "lldb/Protocol/MCP/MCPError.h"
 #include "lldb/Protocol/MCP/Protocol.h"
 #include "llvm/ADT/SmallString.h"
 #include "llvm/Support/FileSystem.h"
 #include "llvm/Support/JSON.h"
 #include "llvm/Support/Signals.h"

 using namespace llvm;
 using namespace lldb_private;
 using namespace lldb_protocol::mcp;

 ServerInfoHandle::ServerInfoHandle(StringRef filename) : m_filename(filename) {
   if (!m_filename.empty())
     sys::RemoveFileOnSignal(m_filename);
 }

 ServerInfoHandle::~ServerInfoHandle() { Remove(); }

 ServerInfoHandle::ServerInfoHandle(ServerInfoHandle &&other) {
   *this = std::move(other);
 }

 ServerInfoHandle &
 ServerInfoHandle::operator=(ServerInfoHandle &&other) noexcept {
   m_filename = std::move(other.m_filename);
   return *this;
 }

 void ServerInfoHandle::Remove() {
   if (m_filename.empty())
     return;

   sys::fs::remove(m_filename);
   sys::DontRemoveFileOnSignal(m_filename);
   m_filename.clear();
 }

 json::Value lldb_protocol::mcp::toJSON(const ServerInfo &SM) {
   return json::Object{{"connection_uri", SM.connection_uri}};
 }

 bool lldb_protocol::mcp::fromJSON(const json::Value &V, ServerInfo &SM,
                                   json::Path P) {
   json::ObjectMapper O(V, P);
   return O && O.map("connection_uri", SM.connection_uri);
 }

 Expected<ServerInfoHandle> ServerInfo::Write(const ServerInfo &info) {
   std::string buf = formatv("{0}", toJSON(info)).str();
   size_t num_bytes = buf.size();

   FileSpec user_lldb_dir = HostInfo::GetUserLLDBDir();

   Status error(sys::fs::create_directory(user_lldb_dir.GetPath()));
   if (error.Fail())
     return error.takeError();

   FileSpec mcp_registry_entry_path = user_lldb_dir.CopyByAppendingPathComponent(
       formatv("lldb-mcp-{0}.json", getpid()).str());

   const File::OpenOptions flags = File::eOpenOptionWriteOnly |
                                   File::eOpenOptionCanCreate |
                                   File::eOpenOptionTruncate;
   Expected<lldb::FileUP> file =
       FileSystem::Instance().Open(mcp_registry_entry_path, flags);
   if (!file)
     return file.takeError();
   if (llvm::Error error = (*file)->Write(buf.data(), num_bytes).takeError())
     return error;
   return ServerInfoHandle{mcp_registry_entry_path.GetPath()};
 }

 Expected<std::vector<ServerInfo>> ServerInfo::Load() {
   namespace path = llvm::sys::path;
   FileSpec user_lldb_dir = HostInfo::GetUserLLDBDir();
   FileSystem &fs = FileSystem::Instance();
   std::error_code EC;
   vfs::directory_iterator it = fs.DirBegin(user_lldb_dir, EC);
   vfs::directory_iterator end;
   std::vector<ServerInfo> infos;
   for (; it != end && !EC; it.increment(EC)) {
     auto &entry = *it;
     auto path = entry.path();
     auto name = path::filename(path);
     if (!name.starts_with("lldb-mcp-") || !name.ends_with(".json"))
       continue;

     auto buffer = fs.CreateDataBuffer(path);
     auto info = json::parse<ServerInfo>(toStringRef(buffer->GetData()));
     if (!info)
       return info.takeError();

     infos.emplace_back(std::move(*info));
   }

   return infos;
 }

 Server::Server(std::string name, std::string version,
                std::unique_ptr<MCPTransport> transport_up,
                lldb_private::MainLoop &loop)
     : m_name(std::move(name)), m_version(std::move(version)),
       m_transport_up(std::move(transport_up)), m_loop(loop) {
   AddRequestHandlers();
 }

 void Server::AddRequestHandlers() {
   AddRequestHandler("initialize", std::bind(&Server::InitializeHandler, this,
                                             std::placeholders::_1));
   AddRequestHandler("tools/list", std::bind(&Server::ToolsListHandler, this,
                                             std::placeholders::_1));
   AddRequestHandler("tools/call", std::bind(&Server::ToolsCallHandler, this,
                                             std::placeholders::_1));
   AddRequestHandler("resources/list", std::bind(&Server::ResourcesListHandler,
                                                 this, std::placeholders::_1));
   AddRequestHandler("resources/read", std::bind(&Server::ResourcesReadHandler,
                                                 this, std::placeholders::_1));
 }

 llvm::Expected<Response> Server::Handle(const Request &request) {
   auto it = m_request_handlers.find(request.method);
   if (it != m_request_handlers.end()) {
     llvm::Expected<Response> response = it->second(request);
     if (!response)
       return response;
     response->id = request.id;
     return *response;
   }

   return llvm::make_error<MCPError>(
       llvm::formatv("no handler for request: {0}", request.method).str());
 }

 void Server::Handle(const Notification &notification) {
   auto it = m_notification_handlers.find(notification.method);
   if (it != m_notification_handlers.end()) {
     it->second(notification);
     return;
   }
 }

 void Server::AddTool(std::unique_ptr<Tool> tool) {
   if (!tool)
     return;
   m_tools[tool->GetName()] = std::move(tool);
 }

 void Server::AddResourceProvider(
     std::unique_ptr<ResourceProvider> resource_provider) {
   if (!resource_provider)
     return;
   m_resource_providers.push_back(std::move(resource_provider));
 }

 void Server::AddRequestHandler(llvm::StringRef method, RequestHandler handler) {
   m_request_handlers[method] = std::move(handler);
 }

 void Server::AddNotificationHandler(llvm::StringRef method,
                                     NotificationHandler handler) {
   m_notification_handlers[method] = std::move(handler);
 }

 llvm::Expected<Response> Server::InitializeHandler(const Request &request) {
   Response response;
   InitializeResult result;
   result.protocolVersion = mcp::kProtocolVersion;
   result.capabilities = GetCapabilities();
   result.serverInfo.name = m_name;
   result.serverInfo.version = m_version;
   response.result = std::move(result);
   return response;
 }

 llvm::Expected<Response> Server::ToolsListHandler(const Request &request) {
   Response response;

   ListToolsResult result;
   for (const auto &tool : m_tools)
     result.tools.emplace_back(tool.second->GetDefinition());

   response.result = std::move(result);

   return response;
 }

 llvm::Expected<Response> Server::ToolsCallHandler(const Request &request) {
   Response response;

   if (!request.params)
     return llvm::createStringError("no tool parameters");
   CallToolParams params;
   json::Path::Root root("params");
   if (!fromJSON(request.params, params, root))
     return root.getError();

   llvm::StringRef tool_name = params.name;
   if (tool_name.empty())
     return llvm::createStringError("no tool name");

   auto it = m_tools.find(tool_name);
   if (it == m_tools.end())
     return llvm::createStringError(llvm::formatv("no tool \"{0}\"", tool_name));

   ToolArguments tool_args;
   if (params.arguments)
     tool_args = *params.arguments;

   llvm::Expected<CallToolResult> text_result = it->second->Call(tool_args);
   if (!text_result)
     return text_result.takeError();

   response.result = toJSON(*text_result);

   return response;
 }

 llvm::Expected<Response> Server::ResourcesListHandler(const Request &request) {
   Response response;

   ListResourcesResult result;
   for (std::unique_ptr<ResourceProvider> &resource_provider_up :
        m_resource_providers)
     for (const Resource &resource : resource_provider_up->GetResources())
       result.resources.push_back(resource);

   response.result = std::move(result);

   return response;
 }

 llvm::Expected<Response> Server::ResourcesReadHandler(const Request &request) {
   Response response;

   if (!request.params)
     return llvm::createStringError("no resource parameters");

   ReadResourceParams params;
   json::Path::Root root("params");
   if (!fromJSON(request.params, params, root))
     return root.getError();

   llvm::StringRef uri_str = params.uri;
   if (uri_str.empty())
     return llvm::createStringError("no resource uri");

   for (std::unique_ptr<ResourceProvider> &resource_provider_up :
        m_resource_providers) {
     llvm::Expected<ReadResourceResult> result =
         resource_provider_up->ReadResource(uri_str);
     if (result.errorIsA<UnsupportedURI>()) {
       llvm::consumeError(result.takeError());
       continue;
     }
     if (!result)
       return result.takeError();

     Response response;
     response.result = std::move(*result);
     return response;
   }

   return make_error<MCPError>(
       llvm::formatv("no resource handler for uri: {0}", uri_str).str(),
       MCPError::kResourceNotFound);
 }

 ServerCapabilities Server::GetCapabilities() {
   lldb_protocol::mcp::ServerCapabilities capabilities;
   capabilities.supportsToolsList = true;
   // FIXME: Support sending notifications when a debugger/target are
   // added/removed.
   capabilities.supportsResourcesList = false;
   return capabilities;
 }

 llvm::Error Server::Run() {
   auto handle = m_transport_up->RegisterMessageHandler(m_loop, *this);
   if (!handle)
     return handle.takeError();

   lldb_private::Status status = m_loop.Run();
   if (status.Fail())
     return status.takeError();

   return llvm::Error::success();
 }

 void Server::Received(const Request &request) {
   auto SendResponse = [this](const Response &response) {
     if (llvm::Error error = m_transport_up->Send(response))
       m_transport_up->Log(llvm::toString(std::move(error)));
   };

   llvm::Expected<Response> response = Handle(request);
   if (response)
     return SendResponse(*response);

   lldb_protocol::mcp::Error protocol_error;
   llvm::handleAllErrors(
       response.takeError(),
       [&](const MCPError &err) { protocol_error = err.toProtocolError(); },
       [&](const llvm::ErrorInfoBase &err) {
         protocol_error.code = MCPError::kInternalError;
         protocol_error.message = err.message();
       });
   Response error_response;
   error_response.id = request.id;
   error_response.result = std::move(protocol_error);
   SendResponse(error_response);
 }

 void Server::Received(const Response &response) {
   m_transport_up->Log("unexpected MCP message: response");
 }

 void Server::Received(const Notification &notification) {
   Handle(notification);
 }

 void Server::OnError(llvm::Error error) {
   m_transport_up->Log(llvm::toString(std::move(error)));
   TerminateLoop();
 }

 void Server::OnClosed() {
   m_transport_up->Log("EOF");
   TerminateLoop();
 }

 void Server::TerminateLoop() {
   m_loop.AddPendingCallback(
       [](lldb_private::MainLoopBase &loop) { loop.RequestTermination(); });
 }
	//===----------------------------------------------------------------------===//
	//
	// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
	// See https://llvm.org/LICENSE.txt for license information.
	// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
	//
	//===----------------------------------------------------------------------===//

	#include "lldb/Protocol/MCP/Server.h"
	#include "lldb/Host/File.h"
	#include "lldb/Host/FileSystem.h"
	#include "lldb/Host/HostInfo.h"
	#include "lldb/Host/JSONTransport.h"
	#include "lldb/Protocol/MCP/MCPError.h"
	#include "lldb/Protocol/MCP/Protocol.h"
	#include "llvm/ADT/SmallString.h"
	#include "llvm/Support/FileSystem.h"
	#include "llvm/Support/JSON.h"
	#include "llvm/Support/Signals.h"

	using namespace llvm;
	using namespace lldb_private;
	using namespace lldb_protocol::mcp;

	ServerInfoHandle::ServerInfoHandle(StringRef filename) : m_filename(filename) {
	if (!m_filename.empty())
	sys::RemoveFileOnSignal(m_filename);
	}

	ServerInfoHandle::~ServerInfoHandle() { Remove(); }

	ServerInfoHandle::ServerInfoHandle(ServerInfoHandle &&other) {
	*this = std::move(other);
	}

	ServerInfoHandle &
	ServerInfoHandle::operator=(ServerInfoHandle &&other) noexcept {
	m_filename = std::move(other.m_filename);
	return *this;
	}

	void ServerInfoHandle::Remove() {
	if (m_filename.empty())
	return;

	sys::fs::remove(m_filename);
	sys::DontRemoveFileOnSignal(m_filename);
	m_filename.clear();
	}

	json::Value lldb_protocol::mcp::toJSON(const ServerInfo &SM) {
	return json::Object{{"connection_uri", SM.connection_uri}};
	}

	bool lldb_protocol::mcp::fromJSON(const json::Value &V, ServerInfo &SM,
	json::Path P) {
	json::ObjectMapper O(V, P);
	return O && O.map("connection_uri", SM.connection_uri);
	}

	Expected<ServerInfoHandle> ServerInfo::Write(const ServerInfo &info) {
	std::string buf = formatv("{0}", toJSON(info)).str();
	size_t num_bytes = buf.size();

	FileSpec user_lldb_dir = HostInfo::GetUserLLDBDir();

	Status error(sys::fs::create_directory(user_lldb_dir.GetPath()));
	if (error.Fail())
	return error.takeError();

	FileSpec mcp_registry_entry_path = user_lldb_dir.CopyByAppendingPathComponent(
	formatv("lldb-mcp-{0}.json", getpid()).str());

	const File::OpenOptions flags = File::eOpenOptionWriteOnly \|
	File::eOpenOptionCanCreate \|
	File::eOpenOptionTruncate;
	Expected<lldb::FileUP> file =
	FileSystem::Instance().Open(mcp_registry_entry_path, flags);
	if (!file)
	return file.takeError();
	if (llvm::Error error = (*file)->Write(buf.data(), num_bytes).takeError())
	return error;
	return ServerInfoHandle{mcp_registry_entry_path.GetPath()};
	}

	Expected<std::vector<ServerInfo>> ServerInfo::Load() {
	namespace path = llvm::sys::path;
	FileSpec user_lldb_dir = HostInfo::GetUserLLDBDir();
	FileSystem &fs = FileSystem::Instance();
	std::error_code EC;
	vfs::directory_iterator it = fs.DirBegin(user_lldb_dir, EC);
	vfs::directory_iterator end;
	std::vector<ServerInfo> infos;
	for (; it != end && !EC; it.increment(EC)) {
	auto &entry = *it;
	auto path = entry.path();
	auto name = path::filename(path);
	if (!name.starts_with("lldb-mcp-") \|\| !name.ends_with(".json"))
	continue;

	auto buffer = fs.CreateDataBuffer(path);
	auto info = json::parse<ServerInfo>(toStringRef(buffer->GetData()));
	if (!info)
	return info.takeError();

	infos.emplace_back(std::move(*info));
	}

	return infos;
	}

	Server::Server(std::string name, std::string version,
	std::unique_ptr<MCPTransport> transport_up,
	lldb_private::MainLoop &loop)
	: m_name(std::move(name)), m_version(std::move(version)),
	m_transport_up(std::move(transport_up)), m_loop(loop) {
	AddRequestHandlers();
	}

	void Server::AddRequestHandlers() {
	AddRequestHandler("initialize", std::bind(&Server::InitializeHandler, this,
	std::placeholders::_1));
	AddRequestHandler("tools/list", std::bind(&Server::ToolsListHandler, this,
	std::placeholders::_1));
	AddRequestHandler("tools/call", std::bind(&Server::ToolsCallHandler, this,
	std::placeholders::_1));
	AddRequestHandler("resources/list", std::bind(&Server::ResourcesListHandler,
	this, std::placeholders::_1));
	AddRequestHandler("resources/read", std::bind(&Server::ResourcesReadHandler,
	this, std::placeholders::_1));
	}

	llvm::Expected<Response> Server::Handle(const Request &request) {
	auto it = m_request_handlers.find(request.method);
	if (it != m_request_handlers.end()) {
	llvm::Expected<Response> response = it->second(request);
	if (!response)
	return response;
	response->id = request.id;
	return *response;
	}

	return llvm::make_error<MCPError>(
	llvm::formatv("no handler for request: {0}", request.method).str());
	}

	void Server::Handle(const Notification &notification) {
	auto it = m_notification_handlers.find(notification.method);
	if (it != m_notification_handlers.end()) {
	it->second(notification);
	return;
	}
	}

	void Server::AddTool(std::unique_ptr<Tool> tool) {
	if (!tool)
	return;
	m_tools[tool->GetName()] = std::move(tool);
	}

	void Server::AddResourceProvider(
	std::unique_ptr<ResourceProvider> resource_provider) {
	if (!resource_provider)
	return;
	m_resource_providers.push_back(std::move(resource_provider));
	}

	void Server::AddRequestHandler(llvm::StringRef method, RequestHandler handler) {
	m_request_handlers[method] = std::move(handler);
	}

	void Server::AddNotificationHandler(llvm::StringRef method,
	NotificationHandler handler) {
	m_notification_handlers[method] = std::move(handler);
	}

	llvm::Expected<Response> Server::InitializeHandler(const Request &request) {
	Response response;
	InitializeResult result;
	result.protocolVersion = mcp::kProtocolVersion;
	result.capabilities = GetCapabilities();
	result.serverInfo.name = m_name;
	result.serverInfo.version = m_version;
	response.result = std::move(result);
	return response;
	}

	llvm::Expected<Response> Server::ToolsListHandler(const Request &request) {
	Response response;

	ListToolsResult result;
	for (const auto &tool : m_tools)
	result.tools.emplace_back(tool.second->GetDefinition());

	response.result = std::move(result);

	return response;
	}

	llvm::Expected<Response> Server::ToolsCallHandler(const Request &request) {
	Response response;

	if (!request.params)
	return llvm::createStringError("no tool parameters");
	CallToolParams params;
	json::Path::Root root("params");
	if (!fromJSON(request.params, params, root))
	return root.getError();

	llvm::StringRef tool_name = params.name;
	if (tool_name.empty())
	return llvm::createStringError("no tool name");

	auto it = m_tools.find(tool_name);
	if (it == m_tools.end())
	return llvm::createStringError(llvm::formatv("no tool \"{0}\"", tool_name));

	ToolArguments tool_args;
	if (params.arguments)
	tool_args = *params.arguments;

	llvm::Expected<CallToolResult> text_result = it->second->Call(tool_args);
	if (!text_result)
	return text_result.takeError();

	response.result = toJSON(*text_result);

	return response;
	}

	llvm::Expected<Response> Server::ResourcesListHandler(const Request &request) {
	Response response;

	ListResourcesResult result;
	for (std::unique_ptr<ResourceProvider> &resource_provider_up :
	m_resource_providers)
	for (const Resource &resource : resource_provider_up->GetResources())
	result.resources.push_back(resource);

	response.result = std::move(result);

	return response;
	}

	llvm::Expected<Response> Server::ResourcesReadHandler(const Request &request) {
	Response response;

	if (!request.params)
	return llvm::createStringError("no resource parameters");

	ReadResourceParams params;
	json::Path::Root root("params");
	if (!fromJSON(request.params, params, root))
	return root.getError();

	llvm::StringRef uri_str = params.uri;
	if (uri_str.empty())
	return llvm::createStringError("no resource uri");

	for (std::unique_ptr<ResourceProvider> &resource_provider_up :
	m_resource_providers) {
	llvm::Expected<ReadResourceResult> result =
	resource_provider_up->ReadResource(uri_str);
	if (result.errorIsA<UnsupportedURI>()) {
	llvm::consumeError(result.takeError());
	continue;
	}
	if (!result)
	return result.takeError();

	Response response;
	response.result = std::move(*result);
	return response;
	}

	return make_error<MCPError>(
	llvm::formatv("no resource handler for uri: {0}", uri_str).str(),
	MCPError::kResourceNotFound);
	}

	ServerCapabilities Server::GetCapabilities() {
	lldb_protocol::mcp::ServerCapabilities capabilities;
	capabilities.supportsToolsList = true;
	// FIXME: Support sending notifications when a debugger/target are
	// added/removed.
	capabilities.supportsResourcesList = false;
	return capabilities;
	}

	llvm::Error Server::Run() {
	auto handle = m_transport_up->RegisterMessageHandler(m_loop, *this);
	if (!handle)
	return handle.takeError();

	lldb_private::Status status = m_loop.Run();
	if (status.Fail())
	return status.takeError();

	return llvm::Error::success();
	}

	void Server::Received(const Request &request) {
	auto SendResponse = [this](const Response &response) {
	if (llvm::Error error = m_transport_up->Send(response))
	m_transport_up->Log(llvm::toString(std::move(error)));
	};

	llvm::Expected<Response> response = Handle(request);
	if (response)
	return SendResponse(*response);

	lldb_protocol::mcp::Error protocol_error;
	llvm::handleAllErrors(
	response.takeError(),
	[&](const MCPError &err) { protocol_error = err.toProtocolError(); },
	[&](const llvm::ErrorInfoBase &err) {
	protocol_error.code = MCPError::kInternalError;
	protocol_error.message = err.message();
	});
	Response error_response;
	error_response.id = request.id;
	error_response.result = std::move(protocol_error);
	SendResponse(error_response);
	}

	void Server::Received(const Response &response) {
	m_transport_up->Log("unexpected MCP message: response");
	}

	void Server::Received(const Notification &notification) {
	Handle(notification);
	}

	void Server::OnError(llvm::Error error) {
	m_transport_up->Log(llvm::toString(std::move(error)));
	TerminateLoop();
	}

	void Server::OnClosed() {
	m_transport_up->Log("EOF");
	TerminateLoop();
	}

	void Server::TerminateLoop() {
	m_loop.AddPendingCallback(
	[](lldb_private::MainLoopBase &loop) { loop.RequestTermination(); });
	}