compliance-scanner-agent/compliance-mcp/src/main.rs

mod database;
mod server;
mod tools;

use std::sync::Arc;

use database::Database;
use rmcp::transport::{
    streamable_http_server::session::local::LocalSessionManager, StreamableHttpServerConfig,
    StreamableHttpService,
};
use server::ComplianceMcpServer;

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    let _ = dotenvy::dotenv();

    tracing_subscriber::fmt()
        .with_env_filter(
            tracing_subscriber::EnvFilter::from_default_env()
                .add_directive("compliance_mcp=info".parse()?),
        )
        .init();

    let mongo_uri =
        std::env::var("MONGODB_URI").unwrap_or_else(|_| "mongodb://localhost:27017".to_string());
    let db_name =
        std::env::var("MONGODB_DATABASE").unwrap_or_else(|_| "compliance_scanner".to_string());

    let db = Database::connect(&mongo_uri, &db_name).await?;

    // If MCP_PORT is set, run as Streamable HTTP server; otherwise use stdio.
    if let Ok(port_str) = std::env::var("MCP_PORT") {
        let port: u16 = port_str.parse()?;
        tracing::info!("Starting MCP server on HTTP port {port}");

        let db_clone = db.clone();
        let service = StreamableHttpService::new(
            move || Ok(ComplianceMcpServer::new(db_clone.clone())),
            Arc::new(LocalSessionManager::default()),
            StreamableHttpServerConfig::default(),
        );

        let router = axum::Router::new().nest_service("/mcp", service);
        let listener = tokio::net::TcpListener::bind(("0.0.0.0", port)).await?;
        tracing::info!("MCP HTTP server listening on 0.0.0.0:{port}");
        axum::serve(listener, router).await?;
    } else {
        tracing::info!("Starting MCP server on stdio");
        let server = ComplianceMcpServer::new(db);
        let transport = rmcp::transport::stdio();
        use rmcp::ServiceExt;
        let handle = server.serve(transport).await?;
        handle.waiting().await?;
    }

    Ok(())
}