transport_stdio.cpp

Go to the documentation of this file.

// SPDX-License-Identifier: Apache-2.0
#include <entropic/mcp/transport_stdio.h>
#include <entropic/types/logging.h>
 
#include <cerrno>
#include <cstring>
#include <chrono>
 
#include <fcntl.h>
#include <poll.h>
#include <signal.h>
#include <spawn.h>
#include <sys/wait.h>
#include <unistd.h>
 
namespace {
struct IgnoreSigpipe {
    IgnoreSigpipe() { signal(SIGPIPE, SIG_IGN); }
};
static IgnoreSigpipe ignore_sigpipe;
} // namespace
 
extern char** environ;
 
static auto logger = entropic::log::get("mcp.transport.stdio");
 
namespace entropic {
 
static std::string sanitize_display_name(const std::string& raw) {
    std::string out;
    out.reserve(raw.size());
    for (char c : raw) {
        if (out.empty() && c == '/') { continue; }
        if (c == '[' || c == ']') { continue; }
        if (static_cast<unsigned char>(c) < 0x20 || c == 0x7f) { continue; }
        out.push_back(c);
    }
    return out.empty() ? std::string{"server"} : out;
}
 
StdioTransport::StdioTransport(
    std::string command,
    std::vector<std::string> args,
    std::map<std::string, std::string> env,
    uint32_t default_timeout_ms)
    : display_name_(sanitize_display_name(command)),
      command_(std::move(command)),
      args_(std::move(args)),
      env_(std::move(env)),
      default_timeout_ms_(default_timeout_ms) {}
 
StdioTransport::StdioTransport(
    std::string display_name,
    std::string command,
    std::vector<std::string> args,
    std::map<std::string, std::string> env,
    uint32_t default_timeout_ms)
    : display_name_(sanitize_display_name(
          display_name.empty() ? command : std::move(display_name))),
      command_(std::move(command)),
      args_(std::move(args)),
      env_(std::move(env)),
      default_timeout_ms_(default_timeout_ms) {}
 
StdioTransport::~StdioTransport() {
    close();
}
 
bool StdioTransport::create_all_pipes(int (&fds)[6]) {
    bool ok = create_pipe(fds[0], fds[1]) &&
              create_pipe(fds[2], fds[3]) &&
              create_pipe(fds[4], fds[5]);
    if (!ok) {
        logger->error("Failed to create pipes: {}", strerror(errno));
        for (auto& fd : fds) { close_fd(fd); }
    }
    return ok;
}
 
bool StdioTransport::open() {
    if (connected_) {
        return true;
    }
 
    if (!open_child_process()) {
        return false;
    }
 
    connected_ = true;
    stderr_thread_ = std::thread(
        &StdioTransport::stderr_reader_loop, this);
 
    logger->info("Spawned child process PID {} for '{}' (cmd: {})",
                 child_pid_, display_name_, command_);
    return true;
}
 
bool StdioTransport::open_child_process() {
    int fds[6] = {-1, -1, -1, -1, -1, -1};
    if (!create_all_pipes(fds)) {
        return false;
    }
 
    auto env_strs = build_env();
    bool ok = spawn_child(fds[0], fds[3], fds[5], env_strs);
 
    // Close child-side pipe ends in parent
    close_fd(fds[0]);
    close_fd(fds[3]);
    close_fd(fds[5]);
 
    if (!ok) {
        close_fd(fds[1]);
        close_fd(fds[2]);
        close_fd(fds[4]);
        return false;
    }
 
    stdin_fd_ = fds[1];
    stdout_fd_ = fds[2];
    stderr_fd_ = fds[4];
    return true;
}
 
void StdioTransport::close() {
    connected_ = false;
 
    terminate_child();
    close_fd(stdin_fd_);
    close_fd(stdout_fd_);
    close_fd(stderr_fd_);
 
    if (stderr_thread_.joinable()) {
        stderr_thread_.join();
    }
 
    logger->info("Closed stdio transport for '{}'", display_name_);
}
 
std::string StdioTransport::send_request(
    const std::string& request_json,
    uint32_t timeout_ms) {
 
    if (!connected_ || cancel_flag_.load(std::memory_order_acquire)) {
        return "";
    }
 
    uint32_t actual_timeout = timeout_ms > 0
        ? timeout_ms : default_timeout_ms_;
 
    std::lock_guard<std::mutex> lock(io_mutex_);
 
    std::string msg = request_json + "\n";
    ssize_t written = ::write(stdin_fd_, msg.data(), msg.size());
    if (written < 0 || static_cast<size_t>(written) != msg.size()) {
        logger->error("Write to child stdin failed: {}",
                      strerror(errno));
        connected_ = false;
        return "";
    }
 
    return read_line(stdout_fd_, actual_timeout);
}
 
bool StdioTransport::is_connected() const {
    if (!connected_) {
        return false;
    }
    if (child_pid_ > 0 && kill(child_pid_, 0) == 0) {
        return true;
    }
    return false;
}
 
void StdioTransport::interrupt() {
    cancel_flag_.store(true, std::memory_order_release);
}
 
static std::vector<char*> build_argv(const std::string& command,
                                     const std::vector<std::string>& args) {
    std::vector<char*> argv;
    argv.push_back(const_cast<char*>(command.c_str()));
    for (const auto& arg : args) {
        argv.push_back(const_cast<char*>(arg.c_str()));
    }
    argv.push_back(nullptr);
    return argv;
}
 
static std::vector<char*> build_envp(
    const std::vector<std::string>& env_strs) {
    std::vector<char*> envp;
    for (const auto& e : env_strs) {
        envp.push_back(const_cast<char*>(e.c_str()));
    }
    envp.push_back(nullptr);
    return envp;
}
 
bool StdioTransport::spawn_child(
    int stdin_r, int stdout_w, int stderr_w,
    const std::vector<std::string>& env_strs) {
 
    posix_spawn_file_actions_t actions;
    posix_spawn_file_actions_init(&actions);
    posix_spawn_file_actions_adddup2(&actions, stdin_r, STDIN_FILENO);
    posix_spawn_file_actions_adddup2(&actions, stdout_w, STDOUT_FILENO);
    posix_spawn_file_actions_adddup2(&actions, stderr_w, STDERR_FILENO);
 
    auto argv = build_argv(command_, args_);
    auto envp = build_envp(env_strs);
 
    int err = posix_spawnp(
        &child_pid_, command_.c_str(), &actions,
        nullptr, argv.data(), envp.data());
 
    posix_spawn_file_actions_destroy(&actions);
 
    if (err != 0) {
        logger->error("posix_spawnp failed for '{}' (cmd: {}): {}",
                      display_name_, command_, strerror(err));
        child_pid_ = -1;
        return false;
    }
    return true;
}
 
std::vector<std::string> StdioTransport::build_env() const {
    std::map<std::string, std::string> merged;
 
    // Copy parent environment
    for (char** ep = environ; ep && *ep; ++ep) {
        std::string entry(*ep);
        auto eq = entry.find('=');
        if (eq != std::string::npos) {
            merged[entry.substr(0, eq)] = entry.substr(eq + 1);
        }
    }
 
    // Apply overrides
    for (const auto& [key, val] : env_) {
        merged[key] = val;
    }
 
    std::vector<std::string> result;
    result.reserve(merged.size());
    for (const auto& [key, val] : merged) {
        result.push_back(key + "=" + val);
    }
    return result;
}
 
bool StdioTransport::create_pipe(int& read_fd, int& write_fd) {
    int fds[2];
    if (::pipe(fds) != 0) {
        return false;
    }
    ::fcntl(fds[0], F_SETFD, FD_CLOEXEC);
    ::fcntl(fds[1], F_SETFD, FD_CLOEXEC);
    read_fd = fds[0];
    write_fd = fds[1];
    return true;
}
 
int StdioTransport::poll_until_ready(
    int fd,
    std::chrono::steady_clock::time_point deadline) {
 
    auto remaining = std::chrono::duration_cast<
        std::chrono::milliseconds>(
        deadline - std::chrono::steady_clock::now());
 
    if (remaining.count() <= 0) {
        return 0;
    }
 
    // P1-10 (2.0.6-rc16): cap the single poll slice at 100ms so the
    // caller's cancel_flag_ is re-checked at least ten times/second.
    constexpr int kSliceMs = 100;
    int slice = static_cast<int>(remaining.count());
    if (slice > kSliceMs) { slice = kSliceMs; }
 
    struct pollfd pfd{fd, POLLIN, 0};
    int rc = ::poll(&pfd, 1, slice);
    if (rc == 0) {
        // Slice expired; report 0 only if the full deadline passed.
        if (std::chrono::steady_clock::now() < deadline) {
            return -2;  // caller retries after checking cancel_flag
        }
    }
    return rc;
}
 
std::string StdioTransport::read_line(int fd, uint32_t timeout_ms) {
    std::string line;
    auto deadline = std::chrono::steady_clock::now() +
                    std::chrono::milliseconds(timeout_ms);
 
    while (true) {
        // P1-10: short-circuit if the engine interrupted this request.
        if (cancel_flag_.load(std::memory_order_acquire)) {
            logger->info("Transport read cancelled by interrupt");
            break;
        }
        int ready = poll_until_ready(fd, deadline);
        if (ready == -2) { continue; }  // slice expired, re-check cancel
        if (ready <= 0) {
            if (ready == 0) { logger->warn("Read timeout after {}ms", timeout_ms); }
            break;
        }
        char ch = 0;
        if (::read(fd, &ch, 1) <= 0) { break; }
        if (ch == '\n') { return line; }
        line += ch;
    }
    return "";
}
 
void StdioTransport::stderr_reader_loop() {
    constexpr int poll_timeout_ms = 500;
    while (connected_) {
        struct pollfd pfd{stderr_fd_, POLLIN, 0};
        int ret = ::poll(&pfd, 1, poll_timeout_ms);
        if (ret <= 0) {
            continue;
        }
 
        char buf[1024];
        ssize_t n = ::read(stderr_fd_, buf, sizeof(buf) - 1);
        if (n <= 0) {
            break;
        }
        buf[n] = '\0';
        logger->warn("[{}] {}", display_name_, buf);
    }
}
 
void StdioTransport::terminate_child() {
    if (child_pid_ <= 0) {
        return;
    }
 
    ::kill(child_pid_, SIGTERM);
 
    // Wait up to 3 seconds for graceful exit
    constexpr int max_wait_ms = 3000;
    constexpr int poll_interval_ms = 50;
    int waited = 0;
    while (waited < max_wait_ms) {
        int status = 0;
        pid_t result = ::waitpid(child_pid_, &status, WNOHANG);
        if (result == child_pid_) {
            child_pid_ = -1;
            return;
        }
        std::this_thread::sleep_for(
            std::chrono::milliseconds(poll_interval_ms));
        waited += poll_interval_ms;
    }
 
    // Force kill
    ::kill(child_pid_, SIGKILL);
    ::waitpid(child_pid_, nullptr, 0);
    logger->warn("Force-killed child PID {}", child_pid_);
    child_pid_ = -1;
}
 
void StdioTransport::close_fd(int& fd) {
    if (fd >= 0) {
        ::close(fd);
        fd = -1;
    }
}
 
} // namespace entropic