src/main.rs

mod indexer;
mod mcp;
mod parser;

use clap::Parser;
use indexer::TantivyEngine;
use mcp::McpServer;
use std::path::PathBuf;
use std::sync::Arc;
use tracing::info;

#[derive(Parser, Debug)]
#[command(name = "rust-mcp-search")]
#[command(about = "Ultra-high-performance Tantivy-backed Model Context Protocol server", long_about = None)]
struct Args {
    /// Transport mode: 'stdio' or 'http'
    #[arg(short, long, env = "MCP_TRANSPORT", default_value = "stdio")]
    transport: String,

    /// HTTP host binding
    #[arg(long, env = "HOST", default_value = "0.0.0.0")]
    host: String,

    /// HTTP port binding (defaults to HF Spaces standard 7860)
    #[arg(short, long, env = "PORT", default_value = "7860")]
    port: u16,

    /// Tantivy index directory path
    #[arg(short, long, env = "DATA_DIR", default_value = "./data/index")]
    index_dir: PathBuf,

    /// Rayon thread pool concurrency limit (0 = auto-detect hardware concurrency)
    #[arg(long, env = "RAYON_NUM_THREADS", default_value = "0")]
    threads: usize,
}

#[tokio::main]
async fn main() -> anyhow::Result<()> {
    let args = Args::parse();

    // Direct all logs strictly to stderr to prevent breaking the MCP STDIO JSON protocol
    tracing_subscriber::fmt()
        .with_writer(std::io::stderr)
        .with_env_filter(
            tracing_subscriber::EnvFilter::from_default_env()
                .add_directive(tracing::Level::INFO.into()),
        )
        .init();

    if args.threads > 0 {
        rayon::ThreadPoolBuilder::new()
            .num_threads(args.threads)
            .build_global()?;
        info!("Configured Rayon thread pool with {} threads", args.threads);
    }

    info!("Initializing Tantivy index engine at: {:?}", args.index_dir);
    let engine = Arc::new(TantivyEngine::new(&args.index_dir)?);
    let server = Arc::new(McpServer::new(engine));

    match args.transport.to_lowercase().as_str() {
        "http" | "sse" => {
            info!("Starting MCP HTTP/SSE transport...");
            server.run_http(&args.host, args.port).await?;
        }
        "stdio" => {
            info!("Starting MCP STDIO transport...");
            server.run_stdio().await?;
        }
        other => {
            eprintln!("Unknown transport '{}'. Expected 'stdio' or 'http'.", other);
            std::process::exit(1);
        }
    }

    Ok(())
}