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

use std::{io, process::ExitCode, time::Duration};

use crank_community_mcp::{
    auth::CommunityMachineCredentialVerifier, build_app_with_background_workers_and_limits,
    session::PostgresTransportSessionStore,
};
use crank_config::{
    CacheBackend as ConfigCacheBackend, ConfigSource, DatabaseSettings, DiagnosticCode,
    McpProcessConfig, ObservabilitySettings, ProcessKind, parse_process,
};
use crank_core::CacheBackend;
use crank_observability::{
    CriticalErrorCategory, MetricsConfig, ObservabilityConfig, ObservabilityLifecycle,
    OtlpTraceConfig, RedactionLimits, SentryConfig, ServiceIdentity, capture_critical_error,
};
use crank_registry::{
    MASTER_KEY_CIPHER_CONTRACT, MasterKeyIdentityCandidate, PostgresPoolConfig, PostgresRegistry,
};
use crank_runtime::{
    RequestRateLimitConfig, RequestRateLimiter, RuntimeCacheConfig, RuntimeCacheStores,
    RuntimeLimits, SecretCrypto,
};
use mcp_server::pool_metrics::spawn_postgres_pool_metrics;
use sqlx::postgres::PgConnectOptions;
use tokio::net::TcpListener;
use tracing::info;

#[tokio::main]
async fn main() -> ExitCode {
    match main_result().await {
        Ok(()) => ExitCode::SUCCESS,
        Err(error) => {
            eprintln!("{}", safe_startup_diagnostic(error.as_ref()));
            ExitCode::FAILURE
        }
    }
}

fn safe_startup_diagnostic(error: &(dyn std::error::Error + 'static)) -> String {
    let mut current = Some(error);
    while let Some(cause) = current {
        if let Some(config) = cause.downcast_ref::<crank_config::ConfigError>() {
            return config.to_json();
        }
        if let Some(migration) = cause.downcast_ref::<crank_registry::MigrationError>() {
            return serde_json::json!({
                "status": "error",
                "code": migration.code(),
                "stage": migration.stage(),
                "version": migration.version(),
                "recovery": migration.recovery(),
            })
            .to_string();
        }
        if let Some(crank_registry::RegistryError::Migration(migration)) =
            cause.downcast_ref::<crank_registry::RegistryError>()
        {
            return serde_json::json!({
                "status": "error",
                "code": migration.code(),
                "stage": migration.stage(),
                "version": migration.version(),
                "recovery": migration.recovery(),
            })
            .to_string();
        }
        if let Some(crank_registry::RegistryError::MasterKeyIdentityMismatch { epoch }) =
            cause.downcast_ref::<crank_registry::RegistryError>()
        {
            return serde_json::json!({
                "status": "error",
                "code": "master_key_identity_mismatch",
                "stage": "startup.master_key_identity",
                "version": epoch,
                "recovery": "configure_same_master_key",
            })
            .to_string();
        }
        if cause
            .downcast_ref::<crank_registry::RegistryError>()
            .is_some_and(|error| {
                matches!(
                    error,
                    crank_registry::RegistryError::InvalidMasterKeyIdentity
                )
            })
        {
            return serde_json::json!({
                "status": "error",
                "code": "master_key_identity_invalid",
                "stage": "startup.master_key_identity",
                "version": null,
                "recovery": "contact_operator",
            })
            .to_string();
        }
        current = cause.source();
    }
    serde_json::json!({
        "status": "error",
        "code": "startup_failed",
        "stage": "startup",
        "version": null,
        "recovery": "contact_operator",
    })
    .to_string()
}

async fn main_result() -> Result<(), Box<dyn std::error::Error>> {
    let effective = parse_process(ProcessKind::McpServer, ConfigSource::from_os()?)?;
    let config = effective
        .mcp()
        .cloned()
        .ok_or_else(|| io::Error::other("MCP configuration projection is unavailable"))?;
    preflight_config(&config)?;
    let observability = init_observability(&config.observability)?;
    for deprecation in effective.deprecations() {
        tracing::warn!(
            name: "config.deprecated",
            field = deprecation.field,
            source_class = deprecation.source_class,
            replacement = deprecation.replacement,
            removal_window = deprecation.removal_window,
            "deprecated configuration accepted"
        );
    }
    let mut startup_completed = false;
    let result = run(config, &observability, &mut startup_completed).await;
    if result.is_err() {
        capture_critical_error(if startup_completed {
            CriticalErrorCategory::Internal
        } else {
            CriticalErrorCategory::Startup
        });
    }
    result
}

async fn run(
    config: McpProcessConfig,
    observability: &ObservabilityLifecycle,
    startup_completed: &mut bool,
) -> Result<(), Box<dyn std::error::Error>> {
    let metrics_config = MetricsConfig::new(
        config.observability.metrics.enabled,
        config.observability.metrics.bind_addr,
        config
            .observability
            .metrics
            .bearer_token
            .as_ref()
            .map(|token| token.expose_secret().to_owned()),
    )?;
    let metrics_enabled = metrics_config.enabled();
    let pool_config = postgres_pool_config(&config.database)?;
    let runtime_limits = RuntimeLimits::try_new(
        config.runtime.max_concurrent_unary,
        config.runtime.max_concurrent_sessions,
    )?;
    let cache_config = runtime_cache_config(&config)?;
    let cache_stores = RuntimeCacheStores::from_config(&cache_config).await?;
    let api_rate_limit = RequestRateLimitConfig::new(
        config.rate_limit.requests_per_second,
        config.rate_limit.burst,
    )?;
    let registry = PostgresRegistry::connect_with_options_and_pool_config(
        database_options(&config.database)?,
        pool_config,
    )
    .await?;
    let metrics_server = if metrics_enabled {
        Some(observability.metrics_surface(metrics_config).bind().await?)
    } else {
        None
    };
    if metrics_enabled {
        spawn_postgres_pool_metrics(registry.pool().clone());
    }
    let session_store = PostgresTransportSessionStore::from_pool(registry.pool().clone()).await?;
    let secret_crypto =
        verified_startup_secret_crypto(&registry, config.runtime.master_key.expose_secret())
            .await?;
    let outbound_http_policy = crank_runtime::OutboundHttpPolicy::try_new_with_limits(
        config.runtime.outbound.allowed_hosts.clone(),
        config.runtime.outbound.denied_hosts.clone(),
        config.runtime.outbound.max_request_bytes,
        config.runtime.outbound.max_response_bytes,
    )?;
    let runtime = crank_runtime::community_with_outbound_policy(outbound_http_policy)
        .with_limits(runtime_limits)
        .with_response_cache(cache_stores.response.clone())
        .with_coordination_store(cache_stores.coordination.clone())
        .build();
    let app = build_app_with_background_workers_and_limits(
        registry,
        Duration::from_millis(config.refresh_ms),
        config.runtime.base_url.clone(),
        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),
        runtime_limits.max_concurrent_sessions,
    );
    let listener = TcpListener::bind(config.bind_addr).await?;

    info!(
        name: "mcp.postgres_pool.configured",
        runtime_max_concurrent_unary = runtime_limits.max_concurrent_unary,
        runtime_max_concurrent_sessions = runtime_limits.max_concurrent_sessions,
        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!(name: "mcp.server.listening", bind_address = %config.bind_addr, "mcp-server listening");
    *startup_completed = true;

    if let Some(metrics_server) = metrics_server {
        tokio::select! {
            result = axum::serve(listener, app) => result?,
            result = metrics_server.serve() => result?,
        }
    } else {
        axum::serve(listener, app).await?;
    }
    Ok(())
}

fn init_observability(
    config: &ObservabilitySettings,
) -> Result<ObservabilityLifecycle, Box<dyn std::error::Error>> {
    let identity = ServiceIdentity::try_new(
        "mcp-server",
        env!("CARGO_PKG_VERSION"),
        config.environment.clone(),
    )?;
    let base = ObservabilityConfig::try_new(
        identity,
        config.log_filter.clone(),
        RedactionLimits::default(),
    )?;
    let sentry = SentryConfig::parse(
        config
            .sentry_dsn
            .as_ref()
            .map(|value| value.expose_secret()),
    )?;
    let values = &config.otlp;
    let otlp = OtlpTraceConfig::from_values(
        values.endpoint.clone(),
        values.traces_endpoint.clone(),
        values.protocol.clone(),
        values.traces_protocol.clone(),
        values.timeout.clone(),
        values.traces_timeout.clone(),
        values
            .headers
            .as_ref()
            .map(|value| value.expose_secret().to_owned()),
        values
            .traces_headers
            .as_ref()
            .map(|value| value.expose_secret().to_owned()),
        values.max_queue_size,
        values.max_export_batch_size,
        values.schedule_delay.clone(),
        values.export_timeout.clone(),
    )?;
    Ok(ObservabilityLifecycle::init_with_exporters(
        base, sentry, otlp,
    )?)
}

fn preflight_config(config: &McpProcessConfig) -> Result<(), crank_config::ConfigError> {
    let invalid = |field| crank_config::ConfigError::single(DiagnosticCode::InvalidType, field);
    database_options(&config.database).map_err(|_| invalid("database.source"))?;
    postgres_pool_config(&config.database).map_err(|_| invalid("database.pool"))?;
    MetricsConfig::new(
        config.observability.metrics.enabled,
        config.observability.metrics.bind_addr,
        config
            .observability
            .metrics
            .bearer_token
            .as_ref()
            .map(|v| v.expose_secret().to_owned()),
    )
    .map_err(|_| invalid("observability.metrics"))?;
    RuntimeLimits::try_new(
        config.runtime.max_concurrent_unary,
        config.runtime.max_concurrent_sessions,
    )
    .map_err(|_| invalid("runtime.limits"))?;
    runtime_cache_config(config).map_err(|_| invalid("cache"))?;
    RequestRateLimitConfig::new(
        config.rate_limit.requests_per_second,
        config.rate_limit.burst,
    )
    .map_err(|_| invalid("mcp.rate_limit"))?;
    SecretCrypto::new(config.runtime.master_key.expose_secret())
        .map_err(|_| invalid("runtime.master_key"))?;
    crank_runtime::OutboundHttpPolicy::try_new_with_limits(
        config.runtime.outbound.allowed_hosts.clone(),
        config.runtime.outbound.denied_hosts.clone(),
        config.runtime.outbound.max_request_bytes,
        config.runtime.outbound.max_response_bytes,
    )
    .map_err(|_| invalid("runtime.outbound"))?;
    let identity = ServiceIdentity::try_new(
        "mcp-server",
        env!("CARGO_PKG_VERSION"),
        config.observability.environment.clone(),
    )
    .map_err(|_| invalid("observability.environment"))?;
    ObservabilityConfig::try_new(
        identity,
        config.observability.log_filter.clone(),
        RedactionLimits::default(),
    )
    .map_err(|_| invalid("observability.log_filter"))?;
    SentryConfig::parse(
        config
            .observability
            .sentry_dsn
            .as_ref()
            .map(|v| v.expose_secret()),
    )
    .map_err(|_| invalid("observability.sentry_dsn"))?;
    let values = &config.observability.otlp;
    OtlpTraceConfig::from_values(
        values.endpoint.clone(),
        values.traces_endpoint.clone(),
        values.protocol.clone(),
        values.traces_protocol.clone(),
        values.timeout.clone(),
        values.traces_timeout.clone(),
        values
            .headers
            .as_ref()
            .map(|v| v.expose_secret().to_owned()),
        values
            .traces_headers
            .as_ref()
            .map(|v| v.expose_secret().to_owned()),
        values.max_queue_size,
        values.max_export_batch_size,
        values.schedule_delay.clone(),
        values.export_timeout.clone(),
    )
    .map_err(|_| invalid("observability.otlp"))?;
    Ok(())
}

async fn verified_startup_secret_crypto(
    registry: &PostgresRegistry,
    master_key: &str,
) -> Result<SecretCrypto, Box<dyn std::error::Error>> {
    let active = registry.active_master_key_identity().await?;
    let secret_crypto = if let Some(identity) = active {
        SecretCrypto::with_epoch(master_key, identity.epoch)?
    } else {
        let crypto = SecretCrypto::new(master_key)?;
        let mut after_secret_id: Option<String> = None;
        let mut after_version: Option<u32> = None;
        loop {
            let versions = registry
                .list_secret_versions_for_master_key_epoch_page(
                    1,
                    after_secret_id.as_deref(),
                    after_version,
                    1_000,
                )
                .await?;
            if versions.is_empty() {
                break;
            }
            for version in versions {
                after_secret_id = Some(version.secret_version.secret_id.as_str().to_owned());
                after_version = Some(version.secret_version.version);
                crypto.decrypt_for_epoch(
                    &version.secret_version.key_version,
                    version.master_key_epoch,
                    &version.secret_version.ciphertext,
                )?;
            }
        }
        crypto
    };
    let master_key_observed_at = time::OffsetDateTime::now_utc();
    registry
        .verify_or_register_master_key_identity(MasterKeyIdentityCandidate {
            epoch: secret_crypto.master_key_epoch(),
            fingerprint: secret_crypto.master_key_fingerprint(),
            cipher_contract: MASTER_KEY_CIPHER_CONTRACT,
            observed_at: &master_key_observed_at,
        })
        .await?;
    Ok(secret_crypto)
}

fn postgres_pool_config(
    config: &DatabaseSettings,
) -> Result<PostgresPoolConfig, crank_registry::PostgresPoolConfigError> {
    PostgresPoolConfig::try_new(
        config.pool.max_connections,
        config.pool.min_connections,
        config.pool.acquire_timeout_ms,
        config.pool.idle_timeout_ms,
        config.pool.max_lifetime_ms,
    )
}

fn database_options(
    config: &DatabaseSettings,
) -> Result<PgConnectOptions, Box<dyn std::error::Error>> {
    if let Some(url) = &config.url {
        return url
            .expose_secret()
            .parse::<PgConnectOptions>()
            .map_err(|_| io::Error::other("database URL is invalid").into());
    }
    Ok(PgConnectOptions::new()
        .host(&config.host)
        .port(config.port)
        .database(&config.database)
        .username(&config.username)
        .password(config.password.expose_secret()))
}

fn runtime_cache_config(
    config: &McpProcessConfig,
) -> Result<RuntimeCacheConfig, crank_runtime::RuntimeCacheConfigError> {
    RuntimeCacheConfig::try_new(
        match config.runtime.cache.backend {
            ConfigCacheBackend::Memory => CacheBackend::Memory,
            ConfigCacheBackend::Valkey => CacheBackend::Valkey,
            ConfigCacheBackend::Redis => CacheBackend::Redis,
        },
        config
            .runtime
            .cache
            .url
            .as_ref()
            .map(|value| value.expose_secret().to_owned()),
    )
}