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

mod auth;
mod database;
mod server;
mod tools;

use std::sync::Arc;

use database::DatabasePool;
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());
    // MONGODB_DATABASE is reused as the per-tenant DB-name prefix —
    // same convention as the agent so `<prefix>__admin.mcp_tokens`
    // and `<prefix>_<tenant_id>` line up across services.
    let db_prefix =
        std::env::var("MONGODB_DATABASE").unwrap_or_else(|_| "compliance_scanner".to_string());

    let pool = DatabasePool::connect(&mongo_uri, &db_prefix).await?;

    // HTTP transport: bind a small axum router with bearer-auth in
    // front of the rmcp service. `/health` stays public for orca's
    // container probe.
    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 pool_for_factory = pool.clone();
        let service = StreamableHttpService::new(
            move || Ok(ComplianceMcpServer::new(pool_for_factory.clone())),
            Arc::new(LocalSessionManager::default()),
            StreamableHttpServerConfig::default(),
        );

        let router = axum::Router::new()
            .route("/health", axum::routing::get(|| async { "ok" }))
            .nest_service(
                "/mcp",
                axum::Router::new().fallback_service(service).layer(
                    axum::middleware::from_fn_with_state(pool.clone(), auth::bearer_auth),
                ),
            );
        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 {
        // stdio transport — used when run as a local MCP server next
        // to the LLM client. There's no HTTP layer to do bearer auth,
        // so we synthesize a tenant_id from STDIO_TENANT_ID for local
        // development. NEVER use this in production.
        tracing::info!("Starting MCP server on stdio");
        let synth_tenant = std::env::var("STDIO_TENANT_ID").unwrap_or_else(|_| "dev".to_string());
        tracing::warn!(
            tenant_id = %synth_tenant,
            "stdio transport — using synthetic tenant id; DO NOT use in production"
        );
        let server = ComplianceMcpServer::new(pool);
        let transport = rmcp::transport::stdio();
        use rmcp::ServiceExt;
        auth::TENANT_ID
            .scope(synth_tenant, async {
                let handle = server.serve(transport).await?;
                handle.waiting().await?;
                Ok::<_, Box<dyn std::error::Error>>(())
            })
            .await?;
    }

    Ok(())
}