yams/tools/yams-mcp/main.cpp

#include <spdlog/sinks/rotating_file_sink.h>
#include <spdlog/sinks/stdout_color_sinks.h>
#include <spdlog/spdlog.h>

#include <atomic>
#include <csignal>
#include <cstdio>
#include <cstdlib>
#include <iostream>
#include <memory>
#include <thread>

#include <CLI/CLI.hpp>

#include <yams/config/config_helpers.h>
#include <yams/mcp/mcp_server.h>

#ifdef _WIN32
#include <fcntl.h>
#include <io.h>

#include <yams/platform/windows_init.h>
#endif

namespace {

std::atomic<bool> g_running{true};

void configureMcpStandardStreams() {
#ifdef _WIN32
    // Keep C and C++ streams synchronized on Windows so std::getline and pipe state agree.
    std::ios::sync_with_stdio(true);
    _setmode(_fileno(stdin), _O_BINARY);
    _setmode(_fileno(stdout), _O_BINARY);
    setvbuf(stdout, nullptr, _IONBF, 0);
    setvbuf(stdin, nullptr, _IONBF, 0);
#else
    std::ios::sync_with_stdio(false);
#endif
    std::cin.tie(nullptr);
    std::cout << std::unitbuf;
    std::cerr << std::unitbuf;
}

void signalHandler(int signal) {
    spdlog::info("Received signal {}, shutting down...", signal);
    g_running = false;
}

} // namespace

int main(int argc, char* argv[]) {
    std::set_terminate([]() noexcept {
        try {
            if (auto ep = std::current_exception()) {
                try {
                    std::rethrow_exception(ep);
                } catch (const std::exception& e) {
                    spdlog::critical("std::terminate: uncaught exception: {}", e.what());
                } catch (...) {
                    spdlog::critical("std::terminate: unknown uncaught exception");
                }
            } else {
                spdlog::critical("std::terminate called without active exception");
            }
        } catch (const std::exception& e) {
            std::cerr << "std::terminate: failed while logging exception: " << e.what()
                      << std::endl;
        } catch (...) {
            std::cerr << "std::terminate: failed while logging unknown exception" << std::endl;
        }
        std::_Exit(1);
    });
    CLI::App app{"YAMS MCP Server - Model Context Protocol server for YAMS"};

    std::string log_level = "info";
    std::string log_file;
    bool daemon_mode = false;
    std::string list_mode = "auto";
    std::string grep_mode = "auto";
    std::string retrieval_mode = "auto";
    std::string daemon_socket;

    app.add_option("-l,--log-level", log_level, "Log level (trace, debug, info, warn, error)")
        ->default_val("info");
    app.add_option("--log-file", log_file, "Log file path (optional)");
    app.add_flag("-d,--daemon", daemon_mode, "Run as daemon");
    app.add_option("--list-mode", list_mode, "Hot/cold mode for list: hot_only|cold_only|auto")
        ->check(CLI::IsMember({"hot_only", "hot", "cold_only", "cold", "auto"}))
        ->default_val("auto");
    app.add_option("--grep-mode", grep_mode, "Hot/cold mode for grep: hot_only|cold_only|auto")
        ->check(CLI::IsMember({"hot_only", "hot", "cold_only", "cold", "auto"}))
        ->default_val("auto");
    app.add_option("--retrieval-mode", retrieval_mode,
                   "Hot/cold mode for retrieval: hot_only|cold_only|auto")
        ->check(CLI::IsMember({"hot_only", "hot", "cold_only", "cold", "auto"}))
        ->default_val("auto");
    app.add_option("--daemon-socket", daemon_socket, "Override daemon socket path")
        ->envname("YAMS_DAEMON_SOCKET");
    CLI11_PARSE(app, argc, argv);
    configureMcpStandardStreams();

    try {
        std::vector<spdlog::sink_ptr> sinks;
        if (!daemon_mode) {
            auto console_sink = std::make_shared<spdlog::sinks::stderr_color_sink_mt>();
            sinks.push_back(console_sink);
        }
        if (!log_file.empty()) {
            auto file_sink = std::make_shared<spdlog::sinks::rotating_file_sink_mt>(
                log_file, 10 * 1024 * 1024, 3);
            sinks.push_back(file_sink);
        }

        auto logger = std::make_shared<spdlog::logger>("yams-mcp", sinks.begin(), sinks.end());
        spdlog::set_default_logger(logger);

        if (log_level == "trace")
            spdlog::set_level(spdlog::level::trace);
        else if (log_level == "debug")
            spdlog::set_level(spdlog::level::debug);
        else if (log_level == "info")
            spdlog::set_level(spdlog::level::info);
        else if (log_level == "warn")
            spdlog::set_level(spdlog::level::warn);
        else if (log_level == "error")
            spdlog::set_level(spdlog::level::err);

        spdlog::set_pattern("[%Y-%m-%d %H:%M:%S.%e] [%l] [%n] %v");
        // These CLI compatibility modes are installed before server threads start and use the
        // shared process-environment boundary. New MCP tuning belongs in typed server options.
        for (const auto& [name, value] :
             {std::pair{"YAMS_LIST_MODE", list_mode}, std::pair{"YAMS_GREP_MODE", grep_mode},
              std::pair{"YAMS_RETRIEVAL_MODE", retrieval_mode}}) {
            if (!value.empty() && !yams::config::set_environment(name, value.c_str())) {
                throw std::runtime_error(std::string{"Failed to apply MCP compatibility mode "} +
                                         name);
            }
        }
    } catch (const std::exception& e) {
        std::cerr << "Failed to setup logging: " << e.what() << std::endl;
        return 1;
    }

    spdlog::info("YAMS MCP Server v{}.{}.{}", 0, 0, 2);
    spdlog::info("Transport: STDIO");

    std::signal(SIGINT, signalHandler);
    std::signal(SIGTERM, signalHandler);
#ifndef _WIN32
    // Prevent abrupt termination when the client closes the pipe early
    std::signal(SIGPIPE, SIG_IGN);
#endif

    try {
        std::unique_ptr<yams::mcp::ITransport> transport =
            std::make_unique<yams::mcp::StdioTransport>();

        // MCPServer uses enable_shared_from_this and spawns detached coroutines that capture
        // shared_ptr<MCPServer>; it must be owned by a shared_ptr.
        auto server =
            std::make_shared<yams::mcp::MCPServer>(std::move(transport), &g_running, daemon_socket);

        std::thread server_thread([server]() {
            server->start();
            g_running = false;
        });

        spdlog::info("MCP server started successfully");
        spdlog::info("Press Ctrl+C to stop");

        while (g_running) {
            std::this_thread::sleep_for(std::chrono::seconds(1));
        }

        spdlog::info("Shutting down MCP server...");
        server->stop();
        if (server_thread.joinable()) {
            server_thread.join();
        }
        spdlog::info("MCP server stopped");
    } catch (const std::exception& e) {
        spdlog::error("Fatal error: {}", e.what());
        return 1;
    }
    return 0;
}