crank/apps/mcp-server/src/main.rs

use std::{env, net::SocketAddr, time::Duration};

use crank_community_mcp::{
    auth::CommunityMachineCredentialVerifier, build_app, session::PostgresTransportSessionStore,
};
use crank_registry::{PostgresPoolConfig, PostgresRegistry};
use crank_runtime::{
    RequestRateLimitConfig, RequestRateLimiter, RuntimeCacheConfig, RuntimeCacheStores,
    RuntimeLimits, SecretCrypto, community_default,
};
use sqlx::postgres::PgConnectOptions;
use tokio::net::TcpListener;
use tracing::info;

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    tracing_subscriber::fmt()
        .with_env_filter(
            env::var("CRANK_LOG_LEVEL")
                .unwrap_or_else(|_| "mcp_server=info,tower_http=info".into()),
        )
        .init();

    let bind_addr = env::var("CRANK_MCP_BIND").unwrap_or_else(|_| "0.0.0.0:3002".into());
    let base_url = env::var("CRANK_BASE_URL").ok();
    let refresh_interval = env::var("CRANK_MCP_REFRESH_MS")
        .ok()
        .and_then(|value| value.parse::<u64>().ok())
        .map(Duration::from_millis)
        .unwrap_or_else(|| Duration::from_secs(5));
    let socket_addr: SocketAddr = bind_addr.parse()?;
    let pool_config = PostgresPoolConfig::from_env()?;
    let runtime_limits = RuntimeLimits::from_env()?;
    let cache_config = RuntimeCacheConfig::from_env()?;
    let cache_stores = RuntimeCacheStores::from_config(&cache_config).await?;
    let api_rate_limit = mcp_api_rate_limit_config_from_env()?;
    let database_options = database_options_from_env()?;
    let registry = PostgresRegistry::connect_with_options_and_pool_config(
        database_options.clone(),
        pool_config,
    )
    .await?;
    let session_store = PostgresTransportSessionStore::connect_with_options_and_pool_config(
        database_options,
        pool_config,
    )
    .await?;
    let secret_crypto = SecretCrypto::new(&env::var("CRANK_MASTER_KEY")?)?;
    let runtime = community_default()
        .with_limits(runtime_limits)
        .with_response_cache(cache_stores.response.clone())
        .with_coordination_store(cache_stores.coordination.clone())
        .build();
    let app = build_app(
        registry,
        refresh_interval,
        base_url,
        secret_crypto,
        runtime,
        if cache_config.backend.is_external() {
            RequestRateLimiter::new_shared(api_rate_limit, cache_stores.rate_limit.clone())
        } else {
            RequestRateLimiter::new(api_rate_limit)
        },
        cache_stores.coordination.clone(),
        std::sync::Arc::new(session_store),
        std::sync::Arc::new(CommunityMachineCredentialVerifier),
    );
    let listener = TcpListener::bind(socket_addr).await?;

    info!(
        runtime_max_concurrent_unary = runtime_limits.max_concurrent_unary,
        mcp_rate_limit_rps = api_rate_limit.requests_per_second,
        mcp_rate_limit_burst = api_rate_limit.burst,
        cache_backend = %cache_config.backend,
        max_connections = pool_config.max_connections,
        min_connections = pool_config.min_connections,
        acquire_timeout_ms = pool_config.acquire_timeout_ms,
        idle_timeout_ms = pool_config.idle_timeout_ms,
        max_lifetime_ms = pool_config.max_lifetime_ms,
        "postgres pool configured"
    );
    info!("mcp-server listening on {}", socket_addr);

    axum::serve(listener, app).await?;

    Ok(())
}

fn database_options_from_env() -> Result<PgConnectOptions, Box<dyn std::error::Error>> {
    if let Ok(database_url) = env::var("CRANK_DATABASE_URL") {
        return Ok(database_url.parse::<PgConnectOptions>()?);
    }

    let host = env::var("POSTGRES_HOST").unwrap_or_else(|_| "postgres".into());
    let port = env::var("POSTGRES_PORT")
        .ok()
        .and_then(|value| value.parse::<u16>().ok())
        .unwrap_or(5432);
    let database = env::var("POSTGRES_DB").unwrap_or_else(|_| "crank".into());
    let username = env::var("POSTGRES_USER").unwrap_or_else(|_| "crank".into());
    let password = env::var("POSTGRES_PASSWORD").unwrap_or_else(|_| "crank".into());

    Ok(PgConnectOptions::new()
        .host(&host)
        .port(port)
        .database(&database)
        .username(&username)
        .password(&password))
}

fn mcp_api_rate_limit_config_from_env() -> Result<RequestRateLimitConfig, Box<dyn std::error::Error>>
{
    let requests_per_second = env::var("CRANK_MCP_RATE_LIMIT_RPS")
        .ok()
        .and_then(|value| value.parse::<u32>().ok())
        .unwrap_or(60);
    let burst = env::var("CRANK_MCP_RATE_LIMIT_BURST")
        .ok()
        .and_then(|value| value.parse::<u32>().ok())
        .unwrap_or(120);

    Ok(RequestRateLimitConfig::new(requests_per_second, burst)?)
}