93 lines
3.6 KiB
Rust
93 lines
3.6 KiB
Rust
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 || {
|
|
// The factory runs in the request task, still inside the bearer
|
|
// middleware's `TENANT_ID` scope, and BEFORE rmcp spawns the
|
|
// session task (which would lose the task_local). So bind the
|
|
// tenant into the session's server instance here, once.
|
|
let tenant_id = auth::current_tenant_id().ok_or_else(|| {
|
|
std::io::Error::other("no tenant context when creating MCP session")
|
|
})?;
|
|
Ok(ComplianceMcpServer::new(
|
|
pool_for_factory.clone(),
|
|
tenant_id,
|
|
))
|
|
},
|
|
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, synth_tenant);
|
|
let transport = rmcp::transport::stdio();
|
|
use rmcp::ServiceExt;
|
|
let handle = server.serve(transport).await?;
|
|
handle.waiting().await?;
|
|
}
|
|
|
|
Ok(())
|
|
}
|