cpp-mcp/src/mcp_stdio_client.cpp

/**
 * @file mcp_stdio_client.cpp
 * @brief Implementation of the MCP stdio client
 * 
 * This file implements the client-side functionality for the Model Context Protocol
 * using standard input/output (stdio) as the transport mechanism.
 * Follows the 2024-11-05 protocol specification.
 */

#include "mcp_stdio_client.h"

#include <unistd.h>
#include <sys/types.h>
#include <sys/wait.h>
#include <fcntl.h>
#include <signal.h>
#include <cstring>
#include <sstream>
#include <iostream>
#include <chrono>

namespace mcp {

stdio_client::stdio_client(const std::string& command, const json& capabilities)
    : command_(command), capabilities_(capabilities) {
    
    LOG_INFO("Creating MCP stdio client for command: ", command);
}

stdio_client::~stdio_client() {
    stop_server_process();
}

bool stdio_client::initialize(const std::string& client_name, const std::string& client_version) {
    LOG_INFO("Initializing MCP stdio client...");
    
    if (!start_server_process()) {
        LOG_ERROR("Failed to start server process");
        return false;
    }
    
    request req = request::create("initialize", {
        {"protocolVersion", MCP_VERSION},
        {"capabilities", capabilities_},
        {"clientInfo", {
            {"name", client_name},
            {"version", client_version}
        }}
    });
    
    try {
        json result = send_jsonrpc(req);
        
        server_capabilities_ = result["capabilities"];
        
        request notification = request::create_notification("initialized");
        send_jsonrpc(notification);
        
        initialized_ = true;
        init_cv_.notify_all();
        
        return true;
    } catch (const std::exception& e) {
        LOG_ERROR("Initialization failed: ", e.what());
        stop_server_process();
        return false;
    }
}

bool stdio_client::ping() {
    if (!running_) {
        return false;
    }
    
    request req = request::create("ping", {});
    
    try {
        json result = send_jsonrpc(req);
        return result.empty();
    } catch (const std::exception& e) {
        return false;
    }
}

void stdio_client::set_capabilities(const json& capabilities) {
    std::lock_guard<std::mutex> lock(mutex_);
    capabilities_ = capabilities;
}

response stdio_client::send_request(const std::string& method, const json& params) {
    if (!running_) {
        throw mcp_exception(error_code::internal_error, "Server process not running");
    }
    
    request req = request::create(method, params);
    json result = send_jsonrpc(req);
    
    response res;
    res.jsonrpc = "2.0";
    res.id = req.id;
    res.result = result;
    
    return res;
}

void stdio_client::send_notification(const std::string& method, const json& params) {
    if (!running_) {
        throw mcp_exception(error_code::internal_error, "Server process not running");
    }
    
    request req = request::create_notification(method, params);
    send_jsonrpc(req);
}

json stdio_client::get_server_capabilities() {
    return server_capabilities_;
}

json stdio_client::call_tool(const std::string& tool_name, const json& arguments) {
    return send_request("tools/call", {
        {"name", tool_name},
        {"arguments", arguments}
    }).result;
}

std::vector<tool> stdio_client::get_tools() {
    json response_json = send_request("tools/list", {}).result;
    std::vector<tool> tools;
    
    json tools_json;
    if (response_json.contains("tools") && response_json["tools"].is_array()) {
        tools_json = response_json["tools"];
    } else if (response_json.is_array()) {
        tools_json = response_json;
    } else {
        return tools;
    }
    
    for (const auto& tool_json : tools_json) {
        tool t;
        t.name = tool_json["name"];
        t.description = tool_json["description"];
        
        if (tool_json.contains("inputSchema")) {
            t.parameters_schema = tool_json["inputSchema"];
        }
        
        tools.push_back(t);
    }
    
    return tools;
}

json stdio_client::get_capabilities() {
    return capabilities_;
}

json stdio_client::list_resources(const std::string& cursor) {
    json params = json::object();
    if (!cursor.empty()) {
        params["cursor"] = cursor;
    }
    return send_request("resources/list", params).result;
}

json stdio_client::read_resource(const std::string& resource_uri) {
    return send_request("resources/read", {
        {"uri", resource_uri}
    }).result;
}

json stdio_client::subscribe_to_resource(const std::string& resource_uri) {
    return send_request("resources/subscribe", {
        {"uri", resource_uri}
    }).result;
}

json stdio_client::list_resource_templates() {
    return send_request("resources/templates/list").result;
}

bool stdio_client::is_running() const {
    return running_;
}

bool stdio_client::start_server_process() {
    if (running_) {
        LOG_INFO("Server process already running");
        return true;
    }
    
    LOG_INFO("Starting server process: ", command_);
    
    // 创建管道
    if (pipe(stdin_pipe_) == -1) {
        LOG_ERROR("Failed to create stdin pipe: ", strerror(errno));
        return false;
    }
    
    if (pipe(stdout_pipe_) == -1) {
        LOG_ERROR("Failed to create stdout pipe: ", strerror(errno));
        close(stdin_pipe_[0]);
        close(stdin_pipe_[1]);
        return false;
    }
    
    // 创建子进程
    process_id_ = fork();
    
    if (process_id_ == -1) {
        LOG_ERROR("Failed to fork process: ", strerror(errno));
        close(stdin_pipe_[0]);
        close(stdin_pipe_[1]);
        close(stdout_pipe_[0]);
        close(stdout_pipe_[1]);
        return false;
    }
    
    if (process_id_ == 0) {
        // 子进程
        
        // 关闭不需要的管道端
        close(stdin_pipe_[1]);  // 关闭写入端
        close(stdout_pipe_[0]); // 关闭读取端
        
        // 重定向标准输入/输出
        if (dup2(stdin_pipe_[0], STDIN_FILENO) == -1) {
            LOG_ERROR("Failed to redirect stdin: ", strerror(errno));
            exit(EXIT_FAILURE);
        }
        
        if (dup2(stdout_pipe_[1], STDOUT_FILENO) == -1) {
            LOG_ERROR("Failed to redirect stdout: ", strerror(errno));
            exit(EXIT_FAILURE);
        }
        
        // 关闭已重定向的文件描述符
        close(stdin_pipe_[0]);
        close(stdout_pipe_[1]);
        
        // 设置非阻塞模式
        int flags = fcntl(STDIN_FILENO, F_GETFL, 0);
        fcntl(STDIN_FILENO, F_SETFL, flags | O_NONBLOCK);
        
        // 执行命令
        std::vector<std::string> args;
        std::istringstream iss(command_);
        std::string arg;
        
        while (iss >> arg) {
            args.push_back(arg);
        }
        
        std::vector<char*> c_args;
        for (auto& a : args) {
            c_args.push_back(const_cast<char*>(a.c_str()));
        }
        c_args.push_back(nullptr);
        
        execvp(c_args[0], c_args.data());
        
        // 如果execvp返回，则表示出错
        LOG_ERROR("Failed to execute command: ", strerror(errno));
        exit(EXIT_FAILURE);
    }
    
    // 父进程
    
    // 关闭不需要的管道端
    close(stdin_pipe_[0]);  // 关闭读取端
    close(stdout_pipe_[1]); // 关闭写入端
    
    // 设置非阻塞模式
    int flags = fcntl(stdout_pipe_[0], F_GETFL, 0);
    fcntl(stdout_pipe_[0], F_SETFL, flags | O_NONBLOCK);
    
    running_ = true;
    
    // 启动读取线程
    read_thread_ = std::make_unique<std::thread>(&stdio_client::read_thread_func, this);
    
    // 等待一段时间，确保进程启动
    std::this_thread::sleep_for(std::chrono::milliseconds(500));
    
    // 检查进程是否仍在运行
    int status;
    pid_t result = waitpid(process_id_, &status, WNOHANG);
    
    if (result == process_id_) {
        LOG_ERROR("Server process exited immediately with status: ", WEXITSTATUS(status));
        running_ = false;
        
        if (read_thread_ && read_thread_->joinable()) {
            read_thread_->join();
        }
        
        close(stdin_pipe_[1]);
        close(stdout_pipe_[0]);
        
        return false;
    } else if (result == -1) {
        LOG_ERROR("Failed to check process status: ", strerror(errno));
        running_ = false;
        
        if (read_thread_ && read_thread_->joinable()) {
            read_thread_->join();
        }
        
        close(stdin_pipe_[1]);
        close(stdout_pipe_[0]);
        
        return false;
    }
    
    LOG_INFO("Server process started successfully, PID: ", process_id_);
    return true;
}

void stdio_client::stop_server_process() {
    if (!running_) {
        return;
    }
    
    LOG_INFO("Stopping server process...");
    
    running_ = false;
    
    // 关闭管道
    if (stdin_pipe_[1] != -1) {
        close(stdin_pipe_[1]);
        stdin_pipe_[1] = -1;
    }
    
    if (stdout_pipe_[0] != -1) {
        close(stdout_pipe_[0]);
        stdout_pipe_[0] = -1;
    }
    
    // 等待读取线程结束
    if (read_thread_ && read_thread_->joinable()) {
        read_thread_->join();
    }
    
    // 终止进程
    if (process_id_ > 0) {
        LOG_INFO("Sending SIGTERM to process: ", process_id_);
        kill(process_id_, SIGTERM);
        
        // 等待进程结束
        int status;
        pid_t result = waitpid(process_id_, &status, WNOHANG);
        
        if (result == 0) {
            // 进程仍在运行，等待一段时间
            std::this_thread::sleep_for(std::chrono::seconds(2));
            
            result = waitpid(process_id_, &status, WNOHANG);
            
            if (result == 0) {
                // 进程仍在运行，强制终止
                LOG_WARNING("Process did not terminate, sending SIGKILL");
                kill(process_id_, SIGKILL);
                waitpid(process_id_, &status, 0);
            }
        }
        
        process_id_ = -1;
    }
    
    LOG_INFO("Server process stopped");
}

void stdio_client::read_thread_func() {
    LOG_INFO("Read thread started");
    
    const int buffer_size = 4096;
    char buffer[buffer_size];
    std::string data_buffer;
    
    while (running_) {
        // 读取数据
        ssize_t bytes_read = read(stdout_pipe_[0], buffer, buffer_size - 1);
        
        if (bytes_read > 0) {
            buffer[bytes_read] = '\0';
            data_buffer.append(buffer, bytes_read);
            
            // 处理完整的JSON-RPC消息
            size_t pos = 0;
            while ((pos = data_buffer.find('\n')) != std::string::npos) {
                std::string line = data_buffer.substr(0, pos);
                data_buffer.erase(0, pos + 1);
                
                if (!line.empty()) {
                    try {
                        json message = json::parse(line);
                        
                        if (message.contains("jsonrpc") && message["jsonrpc"] == "2.0") {
                            if (message.contains("id") && !message["id"].is_null()) {
                                // 这是一个响应
                                json id = message["id"];
                                
                                std::lock_guard<std::mutex> lock(response_mutex_);
                                auto it = pending_requests_.find(id);
                                
                                if (it != pending_requests_.end()) {
                                    if (message.contains("result")) {
                                        it->second.set_value(message["result"]);
                                    } else if (message.contains("error")) {
                                        json error_result = {
                                            {"isError", true},
                                            {"error", message["error"]}
                                        };
                                        it->second.set_value(error_result);
                                    } else {
                                        it->second.set_value(json::object());
                                    }
                                    
                                    pending_requests_.erase(it);
                                } else {
                                    LOG_WARNING("Received response for unknown request ID: ", id);
                                }
                            } else if (message.contains("method")) {
                                // 这是一个请求或通知
                                LOG_INFO("Received request/notification: ", message["method"]);
                                // 目前不处理服务器发来的请求
                            }
                        }
                    } catch (const json::exception& e) {
                        LOG_ERROR("Failed to parse JSON-RPC message: ", e.what(), ", message: ", line);
                    }
                }
            }
        } else if (bytes_read == 0) {
            // 管道已关闭
            LOG_WARNING("Pipe closed by server");
            break;
        } else if (bytes_read == -1) {
            if (errno == EAGAIN || errno == EWOULDBLOCK) {
                // 非阻塞模式下没有数据可读
                std::this_thread::sleep_for(std::chrono::milliseconds(10));
            } else {
                LOG_ERROR("Error reading from pipe: ", strerror(errno));
                break;
            }
        }
    }
    
    LOG_INFO("Read thread stopped");
}

json stdio_client::send_jsonrpc(const request& req) {
    if (!running_) {
        throw mcp_exception(error_code::internal_error, "Server process not running");
    }
    
    json req_json = req.to_json();
    std::string req_str = req_json.dump() + "\n";
    
    // 发送请求
    ssize_t bytes_written = write(stdin_pipe_[1], req_str.c_str(), req_str.size());
    
    if (bytes_written != static_cast<ssize_t>(req_str.size())) {
        LOG_ERROR("Failed to write complete request: ", strerror(errno));
        throw mcp_exception(error_code::internal_error, "Failed to write to pipe");
    }
    
    // 如果是通知，不需要等待响应
    if (req.is_notification()) {
        return json::object();
    }
    
    // 创建Promise和Future
    std::promise<json> response_promise;
    std::future<json> response_future = response_promise.get_future();
    
    {
        std::lock_guard<std::mutex> lock(response_mutex_);
        pending_requests_[req.id] = std::move(response_promise);
    }
    
    // 等待响应，设置超时
    const auto timeout = std::chrono::seconds(30);
    auto status = response_future.wait_for(timeout);
    
    if (status == std::future_status::ready) {
        json response = response_future.get();
        
        if (response.contains("isError") && response["isError"].get<bool>()) {
            int code = response["error"]["code"];
            std::string message = response["error"]["message"];
            
            throw mcp_exception(static_cast<error_code>(code), message);
        }
        
        return response;
    } else {
        {
            std::lock_guard<std::mutex> lock(response_mutex_);
            pending_requests_.erase(req.id);
        }
        
        throw mcp_exception(error_code::internal_error, "Timeout waiting for response");
    }
}

} // namespace mcp