sqlserver-mcp 0.2.0

// SQL Server 2025 - master/msdb/sandbox combined catalog MCP server — generated by mcpify. Do not hand-edit.

use std::net::SocketAddr;
use std::sync::Arc;

use axum::Router;
use axum::extract::{Request, State};
use axum::http::StatusCode;
use axum::middleware::{self, Next};
use axum::response::{IntoResponse, Json, Response};
use axum::routing::get;
use rmcp::transport::streamable_http_server::session::local::LocalSessionManager;
use rmcp::transport::streamable_http_server::{StreamableHttpServerConfig, StreamableHttpService};
use serde_json::json;
use tokio::sync::Mutex as TokioMutex;

use super::metrics;
use super::types::HttpServerConfig;
use crate::core::component_registry::{ComponentRegistry, ComponentStatus};
use crate::core::mcp_server::McpifyServer;

#[derive(Clone)]
struct AppState {
    registry: Arc<TokioMutex<ComponentRegistry>>,
}

async fn healthz(State(state): State<AppState>) -> impl IntoResponse {
    let registry = state.registry.lock().await;
    let status = registry.overall_status();
    let code = if status == ComponentStatus::Unhealthy {
        StatusCode::SERVICE_UNAVAILABLE
    } else {
        StatusCode::OK
    };
    (
        code,
        Json(json!({
            "status": format!("{status:?}"),
            "components": registry.snapshot().len(),
        })),
    )
}

async fn metrics_handler() -> impl IntoResponse {
    metrics::increment("http_requests_total");
    (StatusCode::OK, metrics::render_metrics())
}

/// Increments the request counter for every `/mcp` call. There is no
/// per-caller credential gate here (unlike mcpify's originally-generated
/// HTTP-transport auth relay, which extracted a per-request credential
/// header to forward to a facaded HTTP API): SQL Server auth
/// (`sql_server`/`windows`/`azure_ad`) is always this server's own single
/// configured identity (`AuthManager::resolve_tds_auth`'s config/env/
/// keychain cascade), not something a caller supplies per call, so there
/// is nothing meaningful to extract from an incoming request. Every caller
/// who can reach this HTTP endpoint gets the same DB access this process
/// was configured with — operators running HTTP transport should put it
/// behind their own network-level access control (a reverse proxy, VPN,
/// or firewall rule), not rely on this server for per-caller authorization.
async fn count_request(request: Request, next: Next) -> Response {
    metrics::increment("http_requests_total");
    next.run(request).await
}

async fn set_cors_header(cors_allow: Option<String>, request: Request, next: Next) -> Response {
    let mut response = next.run(request).await;
    if let Some(origin) = cors_allow
        && let Ok(value) = origin.parse()
    {
        response
            .headers_mut()
            .insert("Access-Control-Allow-Origin", value);
    }
    response
}

/// HTTP transport: `/mcp` is handled by rmcp's own streamable-HTTP
/// service; `/healthz` and `/metrics` are plain REST endpoints alongside
/// it.
pub async fn start_http_server(
    mcp_service_factory: impl Fn() -> Result<McpifyServer, std::io::Error> + Send + Sync + 'static,
    config: &HttpServerConfig,
    registry: Arc<TokioMutex<ComponentRegistry>>,
) -> anyhow::Result<()> {
    let mut allowed_hosts = vec![
        "localhost".to_string(),
        "127.0.0.1".to_string(),
        "::1".to_string(),
    ];
    allowed_hosts.push(config.host.clone());
    allowed_hosts.push(format!("{}:{}", config.host, config.port));

    let mcp_service = StreamableHttpService::new(
        mcp_service_factory,
        LocalSessionManager::default().into(),
        StreamableHttpServerConfig::default().with_allowed_hosts(allowed_hosts),
    );

    let state = AppState { registry };
    let cors_allow = config.cors_allow.clone();

    let mcp_router = Router::new()
        .nest_service("/mcp", mcp_service)
        .layer(middleware::from_fn(count_request));

    let app = Router::new()
        .route("/healthz", get(healthz))
        .route("/metrics", get(metrics_handler))
        .merge(mcp_router)
        .layer(middleware::from_fn(move |req, next| {
            set_cors_header(cors_allow.clone(), req, next)
        }))
        .with_state(state);

    let addr: SocketAddr = format!("{}:{}", config.host, config.port).parse()?;
    let listener = tokio::net::TcpListener::bind(addr).await?;
    axum::serve(
        listener,
        app.into_make_service_with_connect_info::<SocketAddr>(),
    )
    .await?;

    Ok(())
}

#[cfg(test)]
mod tests {
    use axum::body::to_bytes;

    use super::*;

    async fn response_json(response: Response) -> serde_json::Value {
        let bytes = to_bytes(response.into_body(), usize::MAX).await.unwrap();
        serde_json::from_slice(&bytes).unwrap()
    }

    #[tokio::test]
    async fn healthz_reports_healthy_components() {
        let mut registry = ComponentRegistry::new();
        registry.register("store", true);
        let state = AppState {
            registry: Arc::new(TokioMutex::new(registry)),
        };

        let response = healthz(State(state)).await.into_response();
        assert_eq!(response.status(), StatusCode::OK);
        assert_eq!(
            response_json(response).await,
            serde_json::json!({ "status": "Healthy", "components": 1 })
        );
    }

    #[tokio::test]
    async fn healthz_returns_service_unavailable_for_a_critical_failure() {
        let mut registry = ComponentRegistry::new();
        registry.register("store", true);
        registry.report(
            "store",
            ComponentStatus::Unhealthy,
            Some("unavailable".to_string()),
        );
        let state = AppState {
            registry: Arc::new(TokioMutex::new(registry)),
        };

        let response = healthz(State(state)).await.into_response();
        assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE);
        assert_eq!(response_json(response).await["status"], "Unhealthy");
    }

    #[tokio::test]
    async fn metrics_endpoint_returns_prometheus_text() {
        let response = metrics_handler().await.into_response();
        assert_eq!(response.status(), StatusCode::OK);
        let body = to_bytes(response.into_body(), usize::MAX).await.unwrap();
        assert!(
            String::from_utf8(body.to_vec())
                .unwrap()
                .contains("http_requests_total")
        );
    }
}