| 1 | //===- ProtocolServerMCP.h ------------------------------------------------===// |
|---|---|
| 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-exception |
| 6 | // |
| 7 | //===----------------------------------------------------------------------===// |
| 8 | |
| 9 | #ifndef LLDB_PLUGINS_PROTOCOL_MCP_PROTOCOLSERVERMCP_H |
| 10 | #define LLDB_PLUGINS_PROTOCOL_MCP_PROTOCOLSERVERMCP_H |
| 11 | |
| 12 | #include "Protocol.h" |
| 13 | #include "Resource.h" |
| 14 | #include "Tool.h" |
| 15 | #include "lldb/Core/ProtocolServer.h" |
| 16 | #include "lldb/Host/MainLoop.h" |
| 17 | #include "lldb/Host/Socket.h" |
| 18 | #include "llvm/ADT/StringMap.h" |
| 19 | #include <thread> |
| 20 | |
| 21 | namespace lldb_private::mcp { |
| 22 | |
| 23 | class ProtocolServerMCP : public ProtocolServer { |
| 24 | public: |
| 25 | ProtocolServerMCP(); |
| 26 | virtual ~ProtocolServerMCP() override; |
| 27 | |
| 28 | virtual llvm::Error Start(ProtocolServer::Connection connection) override; |
| 29 | virtual llvm::Error Stop() override; |
| 30 | |
| 31 | static void Initialize(); |
| 32 | static void Terminate(); |
| 33 | |
| 34 | static llvm::StringRef GetPluginNameStatic() { return "MCP"; } |
| 35 | static llvm::StringRef GetPluginDescriptionStatic(); |
| 36 | |
| 37 | static lldb::ProtocolServerUP CreateInstance(); |
| 38 | |
| 39 | llvm::StringRef GetPluginName() override { return GetPluginNameStatic(); } |
| 40 | |
| 41 | Socket *GetSocket() const override { return m_listener.get(); } |
| 42 | |
| 43 | protected: |
| 44 | using RequestHandler = std::function<llvm::Expected<protocol::Response>( |
| 45 | const protocol::Request &)>; |
| 46 | using NotificationHandler = |
| 47 | std::function<void(const protocol::Notification &)>; |
| 48 | |
| 49 | void AddTool(std::unique_ptr<Tool> tool); |
| 50 | void AddResourceProvider(std::unique_ptr<ResourceProvider> resource_provider); |
| 51 | |
| 52 | void AddRequestHandler(llvm::StringRef method, RequestHandler handler); |
| 53 | void AddNotificationHandler(llvm::StringRef method, |
| 54 | NotificationHandler handler); |
| 55 | |
| 56 | private: |
| 57 | void AcceptCallback(std::unique_ptr<Socket> socket); |
| 58 | |
| 59 | llvm::Expected<std::optional<protocol::Message>> |
| 60 | HandleData(llvm::StringRef data); |
| 61 | |
| 62 | llvm::Expected<protocol::Response> Handle(protocol::Request request); |
| 63 | void Handle(protocol::Notification notification); |
| 64 | |
| 65 | llvm::Expected<protocol::Response> |
| 66 | InitializeHandler(const protocol::Request &); |
| 67 | |
| 68 | llvm::Expected<protocol::Response> |
| 69 | ToolsListHandler(const protocol::Request &); |
| 70 | llvm::Expected<protocol::Response> |
| 71 | ToolsCallHandler(const protocol::Request &); |
| 72 | |
| 73 | llvm::Expected<protocol::Response> |
| 74 | ResourcesListHandler(const protocol::Request &); |
| 75 | llvm::Expected<protocol::Response> |
| 76 | ResourcesReadHandler(const protocol::Request &); |
| 77 | |
| 78 | protocol::Capabilities GetCapabilities(); |
| 79 | |
| 80 | llvm::StringLiteral kName = "lldb-mcp"; |
| 81 | llvm::StringLiteral kVersion = "0.1.0"; |
| 82 | |
| 83 | bool m_running = false; |
| 84 | |
| 85 | MainLoop m_loop; |
| 86 | std::thread m_loop_thread; |
| 87 | |
| 88 | std::unique_ptr<Socket> m_listener; |
| 89 | std::vector<MainLoopBase::ReadHandleUP> m_listen_handlers; |
| 90 | |
| 91 | struct Client { |
| 92 | lldb::IOObjectSP io_sp; |
| 93 | MainLoopBase::ReadHandleUP read_handle_up; |
| 94 | std::string buffer; |
| 95 | }; |
| 96 | llvm::Error ReadCallback(Client &client); |
| 97 | std::vector<std::unique_ptr<Client>> m_clients; |
| 98 | |
| 99 | std::mutex m_server_mutex; |
| 100 | llvm::StringMap<std::unique_ptr<Tool>> m_tools; |
| 101 | std::vector<std::unique_ptr<ResourceProvider>> m_resource_providers; |
| 102 | |
| 103 | llvm::StringMap<RequestHandler> m_request_handlers; |
| 104 | llvm::StringMap<NotificationHandler> m_notification_handlers; |
| 105 | }; |
| 106 | } // namespace lldb_private::mcp |
| 107 | |
| 108 | #endif |
| 109 |
