Don't use JSONTransport in ProtocolServerMCP

Author: JDevlieghere

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

@@ -22,6 +22,8 @@ using namespace llvm;
 
 LLDB_PLUGIN_DEFINE(ProtocolServerMCP)
 
+static constexpr size_t kChunkSize = 1024;
+
 ProtocolServerMCP::ProtocolServerMCP(Debugger &debugger)
     : ProtocolServer(), m_debugger(debugger) {
   AddRequestHandler("initialize",
@@ -91,21 +93,52 @@ void ProtocolServerMCP::AcceptCallback(std::unique_ptr<Socket> socket) {
            m_clients.size() + 1);
 
   lldb::IOObjectSP io_sp = std::move(socket);
-  auto transport_sp = std::make_shared<JSONRPCTransport>(io_sp, io_sp);
+  auto client_up = std::make_unique<Client>();
+  client_up->io_sp = io_sp;
+  Client *client = client_up.get();
 
   Status status;
   auto read_handle_up = m_loop.RegisterReadObject(
       io_sp,
-      [=](MainLoopBase &) {
-        if (llvm::Error err = HandleData(*transport_sp)) {
-          LLDB_LOG_ERROR(GetLog(LLDBLog::Host), std::move(err), "{0}");
+      [this, client](MainLoopBase &loop) {
+        if (Error error = ReadCallback(*client)) {
+          LLDB_LOG_ERROR(GetLog(LLDBLog::Host), std::move(error), "{0}");
+          client->read_handle_up.reset();
         }
       },
       status);
   if (status.Fail())
     return;
 
-  m_clients.emplace_back(io_sp, std::move(read_handle_up));
+  client_up->read_handle_up = std::move(read_handle_up);
+  m_clients.emplace_back(std::move(client_up));
+}
+
+llvm::Error ProtocolServerMCP::ReadCallback(Client &client) {
+  char chunk[kChunkSize];
+  size_t bytes_read = sizeof(chunk);
+  if (Status status = client.io_sp->Read(chunk, bytes_read); status.Fail())
+    return status.takeError();
+  client.buffer.append(chunk, bytes_read);
+
+  for (std::string::size_type pos;
+       (pos = client.buffer.find('\n')) != std::string::npos;) {
+    llvm::Expected<std::optional<protocol::Message>> message =
+        HandleData(StringRef(client.buffer.data(), pos));
+    client.buffer = client.buffer.erase(0, pos + 1);
+    if (!message)
+      return message.takeError();
+
+    if (*message) {
+      std::string Output;
+      llvm::raw_string_ostream OS(Output);
+      OS << llvm::formatv("{0}", toJSON(**message)) << '\n';
+      size_t num_bytes = Output.size();
+      return client.io_sp->Write(Output.data(), num_bytes).takeError();
+    }
+  }
+
+  return llvm::Error::success();
 }
 
 llvm::Error ProtocolServerMCP::Start(ProtocolServer::Connection connection) {
@@ -165,58 +198,43 @@ llvm::Error ProtocolServerMCP::Stop() {
   return llvm::Error::success();
 }
 
-llvm::Error ProtocolServerMCP::HandleData(JSONTransport &transport) {
-  llvm::Expected<protocol::Message> maybe_message =
-      transport.Read<protocol::Message>(std::chrono::seconds(1));
-  if (maybe_message.errorIsA<TransportEOFError>() ||
-      maybe_message.errorIsA<TransportInvalidError>() ||
-      maybe_message.errorIsA<TransportTimeoutError>()) {
-    consumeError(maybe_message.takeError());
-    return llvm::Error::success();
-  }
-
-  if (llvm::Error err = maybe_message.takeError())
-    return err;
+llvm::Expected<std::optional<protocol::Message>>
+ProtocolServerMCP::HandleData(llvm::StringRef data) {
+  auto message = llvm::json::parse<protocol::Message>(/*JSON=*/data);
+  if (!message)
+    return message.takeError();
 
-  protocol::Message &message = *maybe_message;
   if (const protocol::Request *request =
-          std::get_if<protocol::Request>(&message)) {
-    llvm::Expected<protocol::Response> maybe_response = Handle(*request);
+          std::get_if<protocol::Request>(&(*message))) {
+    llvm::Expected<protocol::Response> response = Handle(*request);
 
-    // Handle failures.
-    if (!maybe_response) {
+    // Handle failures by converting them into an Error message.
+    if (!response) {
       protocol::Error protocol_error;
       llvm::handleAllErrors(
-          maybe_response.takeError(),
+          response.takeError(),
           [&](const MCPError &err) { protocol_error = err.toProtcolError(); },
           [&](const llvm::ErrorInfoBase &err) {
             protocol_error.error.code = -1;
             protocol_error.error.message = err.message();
           });
       protocol_error.id = request->id;
-      if (llvm::Error err = transport.Write(protocol_error))
-        return err;
-
-      return llvm::Error::success();
+      return protocol_error;
     }
 
-    // Handle success.
-    if (llvm::Error err = transport.Write(*maybe_response))
-      return err;
-
-    return llvm::Error::success();
+    return *response;
   }
 
   if (const protocol::Notification *notification =
-          std::get_if<protocol::Notification>(&message)) {
+          std::get_if<protocol::Notification>(&(*message))) {
     Handle(*notification);
-    return llvm::Error::success();
+    return std::nullopt;
   }
 
-  if (std::get_if<protocol::Error>(&message))
+  if (std::get_if<protocol::Error>(&(*message)))
     return llvm::createStringError("unexpected MCP message: error");
 
-  if (std::get_if<protocol::Response>(&message))
+  if (std::get_if<protocol::Response>(&(*message)))
     return llvm::createStringError("unexpected MCP message: response");
 
   llvm_unreachable("all message types handled");

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

@@ -12,7 +12,6 @@
 #include "Protocol.h"
 #include "Tool.h"
 #include "lldb/Core/ProtocolServer.h"
-#include "lldb/Host/JSONTransport.h"
 #include "lldb/Host/MainLoop.h"
 #include "lldb/Host/Socket.h"
 #include "llvm/ADT/StringMap.h"
@@ -54,7 +53,8 @@ class ProtocolServerMCP : public ProtocolServer {
 private:
   void AcceptCallback(std::unique_ptr<Socket> socket);
 
-  llvm::Error HandleData(JSONTransport &transport);
+  llvm::Expected<std::optional<protocol::Message>>
+  HandleData(llvm::StringRef data);
 
   llvm::Expected<protocol::Response> Handle(protocol::Request request);
   void Handle(protocol::Notification notification);
@@ -80,8 +80,14 @@ class ProtocolServerMCP : public ProtocolServer {
 
   std::unique_ptr<Socket> m_listener;
   std::vector<MainLoopBase::ReadHandleUP> m_listen_handlers;
-  std::vector<std::pair<lldb::IOObjectSP, MainLoopBase::ReadHandleUP>>
-      m_clients;
+
+  struct Client {
+    lldb::IOObjectSP io_sp;
+    MainLoopBase::ReadHandleUP read_handle_up;
+    std::string buffer;
+  };
+  llvm::Error ReadCallback(Client &client);
+  std::vector<std::unique_ptr<Client>> m_clients;
 
   std::mutex m_server_mutex;
   llvm::StringMap<std::unique_ptr<Tool>> m_tools;

lldb/unittests/Protocol/ProtocolMCPServerTest.cpp

@@ -13,6 +13,7 @@
 #include "lldb/Core/ProtocolServer.h"
 #include "lldb/Host/FileSystem.h"
 #include "lldb/Host/HostInfo.h"
+#include "lldb/Host/JSONTransport.h"
 #include "lldb/Host/Socket.h"
 #include "llvm/Testing/Support/Error.h"
 #include "gtest/gtest.h"