[lldb] Adding A new Binding helper for JSONTransport.

This adds a new Binding helper class to allow mapping of incoming and outgoing requests / events to specific handlers.

This should make it easier to create new protocol implementations and allow us to create a relay in the lldb-mcp binary.

Author: ashgti

lldb/include/lldb/Host/JSONTransport.h

@@ -14,6 +14,7 @@
 #ifndef LLDB_PROTOCOL_MCP_PROTOCOL_H
 #define LLDB_PROTOCOL_MCP_PROTOCOL_H
 
+#include "llvm/ADT/StringRef.h"
 #include "llvm/Support/JSON.h"
 #include <optional>
 #include <string>
@@ -324,4 +325,11 @@ bool fromJSON(const llvm::json::Value &, CallToolResult &, llvm::json::Path);
 
 } // namespace lldb_protocol::mcp
 
+namespace llvm::json {
+inline Value toJSON(const lldb_protocol::mcp::Void &) { return Object(); }
+inline bool fromJSON(const Value &, lldb_protocol::mcp::Void &, Path) {
+  return true;
+}
+} // namespace llvm::json
+
 #endif

lldb/include/lldb/Protocol/MCP/Protocol.h

@@ -9,7 +9,6 @@
 #ifndef LLDB_PROTOCOL_MCP_SERVER_H
 #define LLDB_PROTOCOL_MCP_SERVER_H
 
-#include "lldb/Host/JSONTransport.h"
 #include "lldb/Host/MainLoop.h"
 #include "lldb/Protocol/MCP/Protocol.h"
 #include "lldb/Protocol/MCP/Resource.h"
@@ -19,75 +18,67 @@
 #include "llvm/ADT/StringMap.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/Support/Error.h"
+#include "llvm/Support/FormatVariadic.h"
 #include "llvm/Support/JSON.h"
 #include "llvm/Support/Signals.h"
-#include <functional>
 #include <memory>
 #include <string>
 #include <vector>
 
 namespace lldb_protocol::mcp {
 
-class Server : public MCPTransport::MessageHandler {
-  using ClosedCallback = llvm::unique_function<void()>;
+
+class Server {
+
+  using MCPTransportUP = std::unique_ptr<lldb_protocol::mcp::MCPTransport>;
+
+  using ReadHandleUP = lldb_private::MainLoop::ReadHandleUP;
 
 public:
-  Server(std::string name, std::string version, MCPTransport &client,
-         LogCallback log_callback = {}, ClosedCallback closed_callback = {});
+  Server(std::string name, std::string version, LogCallback log_callback = {});
   ~Server() = default;
 
-  using NotificationHandler = std::function<void(const Notification &)>;
-
   void AddTool(std::unique_ptr<Tool> tool);
   void AddResourceProvider(std::unique_ptr<ResourceProvider> resource_provider);
-  void AddNotificationHandler(llvm::StringRef method,
-                              NotificationHandler handler);
-
-protected:
-  ServerCapabilities GetCapabilities();
 
-  using RequestHandler =
-      std::function<llvm::Expected<Response>(const Request &)>;
+  llvm::Error Accept(lldb_private::MainLoop &, MCPTransportUP);
 
-  void AddRequestHandlers();
-
-  void AddRequestHandler(llvm::StringRef method, RequestHandler handler);
-
-  llvm::Expected<std::optional<Message>> HandleData(llvm::StringRef data);
-
-  llvm::Expected<Response> Handle(const Request &request);
-  void Handle(const Notification &notification);
+protected:
+  MCPTransport::BinderUP Bind(MCPTransport &);
 
-  llvm::Expected<Response> InitializeHandler(const Request &);
+  ServerCapabilities GetCapabilities();
 
-  llvm::Expected<Response> ToolsListHandler(const Request &);
-  llvm::Expected<Response> ToolsCallHandler(const Request &);
+  llvm::Expected<InitializeResult> InitializeHandler(const InitializeParams &);
 
-  llvm::Expected<Response> ResourcesListHandler(const Request &);
-  llvm::Expected<Response> ResourcesReadHandler(const Request &);
+  llvm::Expected<ListToolsResult> ToolsListHandler();
+  llvm::Expected<CallToolResult> ToolsCallHandler(const CallToolParams &);
 
-  void Received(const Request &) override;
-  void Received(const Response &) override;
-  void Received(const Notification &) override;
-  void OnError(llvm::Error) override;
-  void OnClosed() override;
+  llvm::Expected<ListResourcesResult> ResourcesListHandler();
+  llvm::Expected<ReadResourceResult>
+  ResourcesReadHandler(const ReadResourceParams &);
 
-protected:
-  void Log(llvm::StringRef);
+  template <typename... Ts> inline auto Logv(const char *Fmt, Ts &&...Vals) {
+    Log(llvm::formatv(Fmt, std::forward<Ts>(Vals)...).str());
+  }
+  void Log(llvm::StringRef message) {
+    if (m_log_callback)
+      m_log_callback(message);
+  }
 
 private:
   const std::string m_name;
   const std::string m_version;
 
-  MCPTransport &m_client;
   LogCallback m_log_callback;
-  ClosedCallback m_closed_callback;
+  struct Client {
+    ReadHandleUP handle;
+    MCPTransportUP transport;
+    MCPTransport::BinderUP binder;
+  };
+  std::map<MCPTransport *, Client> m_instances;
 
   llvm::StringMap<std::unique_ptr<Tool>> m_tools;
   std::vector<std::unique_ptr<ResourceProvider>> m_resource_providers;
-
-  llvm::StringMap<RequestHandler> m_request_handlers;
-  llvm::StringMap<NotificationHandler> m_notification_handlers;
 };
 
 class ServerInfoHandle;
@@ -121,7 +112,7 @@ class ServerInfoHandle {
   ServerInfoHandle &operator=(const ServerInfoHandle &) = delete;
   /// @}
 
-  /// Remove the file.
+  /// Remove the file on disk, if one is tracked.
   void Remove();
 
 private:

lldb/include/lldb/Protocol/MCP/Server.h

@@ -10,22 +10,95 @@
 #define LLDB_PROTOCOL_MCP_TRANSPORT_H
 
 #include "lldb/Host/JSONTransport.h"
+#include "lldb/Protocol/MCP/MCPError.h"
 #include "lldb/Protocol/MCP/Protocol.h"
 #include "lldb/lldb-forward.h"
 #include "llvm/ADT/FunctionExtras.h"
 #include "llvm/ADT/StringRef.h"
+#include "llvm/Support/Error.h"
+
+namespace lldb_private {
+/// Specializations of the JSONTransport protocol functions for MCP.
+/// @{
+template <>
+inline lldb_protocol::mcp::Request
+make_request(int64_t id, llvm::StringRef method,
+             std::optional<llvm::json::Value> params) {
+  return lldb_protocol::mcp::Request{id, method.str(), params};
+}
+template <>
+inline lldb_protocol::mcp::Response
+make_response(const lldb_protocol::mcp::Request &req, llvm::Error error) {
+  lldb_protocol::mcp::Error protocol_error;
+  llvm::handleAllErrors(
+      std::move(error),
+      [&](const lldb_protocol::mcp::MCPError &err) {
+        protocol_error = err.toProtocolError();
+      },
+      [&](const llvm::ErrorInfoBase &err) {
+        protocol_error.code = lldb_protocol::mcp::MCPError::kInternalError;
+        protocol_error.message = err.message();
+      });
+
+  return lldb_protocol::mcp::Response{req.id, std::move(protocol_error)};
+}
+template <>
+inline lldb_protocol::mcp::Response
+make_response(const lldb_protocol::mcp::Request &req,
+              llvm::json::Value result) {
+  return lldb_protocol::mcp::Response{req.id, std::move(result)};
+}
+template <>
+inline lldb_protocol::mcp::Notification
+make_event(llvm::StringRef method, std::optional<llvm::json::Value> params) {
+  return lldb_protocol::mcp::Notification{method.str(), params};
+}
+template <>
+inline llvm::Expected<llvm::json::Value>
+get_result(const lldb_protocol::mcp::Response &resp) {
+  if (const lldb_protocol::mcp::Error *error =
+          std::get_if<lldb_protocol::mcp::Error>(&resp.result))
+    return llvm::make_error<lldb_protocol::mcp::MCPError>(error->message,
+                                                          error->code);
+  return std::get<llvm::json::Value>(resp.result);
+}
+template <> inline int64_t get_id(const lldb_protocol::mcp::Response &resp) {
+  return std::get<int64_t>(resp.id);
+}
+template <>
+inline llvm::StringRef get_method(const lldb_protocol::mcp::Request &req) {
+  return req.method;
+}
+template <>
+inline llvm::StringRef get_method(const lldb_protocol::mcp::Notification &evt) {
+  return evt.method;
+}
+template <>
+inline llvm::json::Value get_params(const lldb_protocol::mcp::Request &req) {
+  return req.params;
+}
+template <>
+inline llvm::json::Value
+get_params(const lldb_protocol::mcp::Notification &evt) {
+  return evt.params;
+}
+/// @}
+
+} // end namespace lldb_private
 
 namespace lldb_protocol::mcp {
 
 /// Generic transport that uses the MCP protocol.
-using MCPTransport = lldb_private::Transport<Request, Response, Notification>;
+using MCPTransport =
+    lldb_private::JSONTransport<int64_t, Request, Response, Notification>;
 
 /// Generic logging callback, to allow the MCP server / client / transport layer
 /// to be independent of the lldb log implementation.
 using LogCallback = llvm::unique_function<void(llvm::StringRef message)>;
 
 class Transport final
-    : public lldb_private::JSONRPCTransport<Request, Response, Notification> {
+    : public lldb_private::JSONRPCTransport<int64_t, Request, Response,
+                                            Notification> {
 public:
   Transport(lldb::IOObjectSP in, lldb::IOObjectSP out,
             LogCallback log_callback = {});

lldb/include/lldb/Protocol/MCP/Transport.h

@@ -30,3 +30,13 @@ void TransportUnhandledContentsError::log(llvm::raw_ostream &OS) const {
 std::error_code TransportUnhandledContentsError::convertToErrorCode() const {
   return std::make_error_code(std::errc::bad_message);
 }
+
+char InvalidParams::ID;
+
+void InvalidParams::log(llvm::raw_ostream &OS) const {
+  OS << "invalid parameters for method '" << m_method << "': '" << m_context
+     << "'";
+}
+std::error_code InvalidParams::convertToErrorCode() const {
+  return std::make_error_code(std::errc::invalid_argument);
+}

lldb/source/Host/common/JSONTransport.cpp

@@ -52,11 +52,6 @@ llvm::StringRef ProtocolServerMCP::GetPluginDescriptionStatic() {
 }
 
 void ProtocolServerMCP::Extend(lldb_protocol::mcp::Server &server) const {
-  server.AddNotificationHandler("notifications/initialized",
-                                [](const lldb_protocol::mcp::Notification &) {
-                                  LLDB_LOG(GetLog(LLDBLog::Host),
-                                           "MCP initialization complete");
-                                });
   server.AddTool(
       std::make_unique<CommandTool>("command", "Run an lldb command."));
   server.AddTool(std::make_unique<DebuggerListTool>(
@@ -74,26 +69,9 @@ void ProtocolServerMCP::AcceptCallback(std::unique_ptr<Socket> socket) {
       io_sp, io_sp, [client_name](llvm::StringRef message) {
         LLDB_LOG(GetLog(LLDBLog::Host), "{0}: {1}", client_name, message);
       });
-  MCPTransport *transport_ptr = transport_up.get();
-  auto instance_up = std::make_unique<lldb_protocol::mcp::Server>(
-      std::string(kName), std::string(kVersion), *transport_up,
-      /*log_callback=*/
-      [client_name](llvm::StringRef message) {
-        LLDB_LOG(GetLog(LLDBLog::Host), "{0} Server: {1}", client_name,
-                 message);
-      },
-      /*closed_callback=*/
-      [this, transport_ptr]() { m_instances.erase(transport_ptr); });
-  Extend(*instance_up);
-  llvm::Expected<MainLoop::ReadHandleUP> handle =
-      transport_up->RegisterMessageHandler(m_loop, *instance_up);
-  if (!handle) {
-    LLDB_LOG_ERROR(log, handle.takeError(), "Failed to run MCP server: {0}");
-    return;
-  }
-  m_instances[transport_ptr] =
-      std::make_tuple<ServerUP, ReadHandleUP, TransportUP>(
-          std::move(instance_up), std::move(*handle), std::move(transport_up));
+
+  if (auto error = m_server->Accept(m_loop, std::move(transport_up)))
+    LLDB_LOG_ERROR(log, std::move(error), "{0}:");
 }
 
 llvm::Error ProtocolServerMCP::Start(ProtocolServer::Connection connection) {
@@ -124,13 +102,20 @@ llvm::Error ProtocolServerMCP::Start(ProtocolServer::Connection connection) {
       llvm::join(m_listener->GetListeningConnectionURI(), ", ");
 
   ServerInfo info{listening_uris[0]};
-  llvm::Expected<ServerInfoHandle> handle = ServerInfo::Write(info);
-  if (!handle)
-    return handle.takeError();
+  llvm::Expected<ServerInfoHandle> server_info_handle = ServerInfo::Write(info);
+  if (!server_info_handle)
+    return server_info_handle.takeError();
+
+  m_client_count = 0;
+  m_server = std::make_unique<lldb_protocol::mcp::Server>(
+      std::string(kName), std::string(kVersion), [](StringRef message) {
+        LLDB_LOG(GetLog(LLDBLog::Host), "MCP Server: {0}", message);
+      });
+  Extend(*m_server);
 
   m_running = true;
-  m_server_info_handle = std::move(*handle);
-  m_listen_handlers = std::move(*handles);
+  m_server_info_handle = std::move(*server_info_handle);
+  m_accept_handles = std::move(*handles);
   m_loop_thread = std::thread([=] {
     llvm::set_thread_name("protocol-server.mcp");
     m_loop.Run();
@@ -155,9 +140,10 @@ llvm::Error ProtocolServerMCP::Stop() {
   if (m_loop_thread.joinable())
     m_loop_thread.join();
 
+  m_accept_handles.clear();
+
+  m_server.reset(nullptr);
   m_server_info_handle.Remove();
-  m_listen_handlers.clear();
-  m_instances.clear();
 
   return llvm::Error::success();
 }

lldb/source/Plugins/Protocol/MCP/ProtocolServerMCP.cpp

@@ -23,16 +23,17 @@
 namespace lldb_private::mcp {
 
 class ProtocolServerMCP : public ProtocolServer {
-  using ReadHandleUP = MainLoopBase::ReadHandleUP;
-  using TransportUP = std::unique_ptr<lldb_protocol::mcp::MCPTransport>;
+
   using ServerUP = std::unique_ptr<lldb_protocol::mcp::Server>;
 
+  using ReadHandleUP = MainLoop::ReadHandleUP;
+
 public:
   ProtocolServerMCP();
-  virtual ~ProtocolServerMCP() override;
+  ~ProtocolServerMCP() override;
 
-  virtual llvm::Error Start(ProtocolServer::Connection connection) override;
-  virtual llvm::Error Stop() override;
+  llvm::Error Start(ProtocolServer::Connection connection) override;
+  llvm::Error Stop() override;
 
   static void Initialize();
   static void Terminate();
@@ -56,19 +57,18 @@ class ProtocolServerMCP : public ProtocolServer {
 
   bool m_running = false;
 
-  lldb_protocol::mcp::ServerInfoHandle m_server_info_handle;
   lldb_private::MainLoop m_loop;
   std::thread m_loop_thread;
   std::mutex m_mutex;
   size_t m_client_count = 0;
 
   std::unique_ptr<Socket> m_listener;
+  std::vector<ReadHandleUP> m_accept_handles;
 
-  std::vector<ReadHandleUP> m_listen_handlers;
-  std::map<lldb_protocol::mcp::MCPTransport *,
-           std::tuple<ServerUP, ReadHandleUP, TransportUP>>
-      m_instances;
+  ServerUP m_server;
+  lldb_protocol::mcp::ServerInfoHandle m_server_info_handle;
 };
+
 } // namespace lldb_private::mcp
 
 #endif