Merge pull request #209 from apollographql/GT-121

Health Check Support

Author: DaleSeo

.changesets/feat_dale_health_check.md

@@ -0,0 +1,3 @@
+### Health Check Support - @DaleSeo PR #209
+
+Health reporting functionality has been added to make the MCP server ready for production deployment with proper health monitoring and Kubernetes integration.

Cargo.lock

@@ -17,6 +17,7 @@ bon = "3.6.3"
 clap = { version = "4.5.36", features = ["derive", "env"] }
 figment = { version = "0.10.19", features = ["env", "yaml"] }
 futures.workspace = true
+humantime-serde = "1.1.1"
 lz-str = "0.2.1"
 regex = "1.11.1"
 reqwest.workspace = true

crates/apollo-mcp-server/Cargo.toml

@@ -0,0 +1,255 @@
+//! Health Check module for Apollo MCP Server
+//!
+//! Provides liveness and readiness checks for the MCP server, inspired by Apollo Router's health check implementation.
+//!
+//! The health check is exposed via HTTP endpoints and can be used by load balancers, container orchestrators, and monitoring systems to determine server health.
+
+use std::{
+    sync::{
+        Arc,
+        atomic::{AtomicBool, AtomicUsize, Ordering},
+    },
+    time::Duration,
+};
+
+use axum::http::StatusCode;
+use schemars::JsonSchema;
+use serde::{Deserialize, Serialize};
+use tokio::time::Instant;
+use tracing::debug;
+
+/// Health status enumeration
+#[derive(Debug, Serialize)]
+#[serde(rename_all = "UPPERCASE")]
+pub enum HealthStatus {
+    Up,
+    Down,
+}
+
+/// Health response structure
+#[derive(Debug, Serialize)]
+pub struct Health {
+    status: HealthStatus,
+}
+
+/// Configuration options for the readiness health interval sub-component.
+#[derive(Clone, Debug, Deserialize, Serialize, JsonSchema)]
+#[serde(deny_unknown_fields)]
+#[serde(default)]
+pub struct ReadinessIntervalConfig {
+    #[serde(deserialize_with = "humantime_serde::deserialize", default)]
+    #[serde(serialize_with = "humantime_serde::serialize")]
+    #[schemars(with = "Option<String>", default)]
+    /// The sampling interval (default: 5s)
+    pub sampling: Duration,
+
+    #[serde(deserialize_with = "humantime_serde::deserialize")]
+    #[serde(serialize_with = "humantime_serde::serialize")]
+    #[schemars(with = "Option<String>")]
+    /// The unready interval (default: 2 * sampling interval)
+    pub unready: Option<Duration>,
+}
+
+impl Default for ReadinessIntervalConfig {
+    fn default() -> Self {
+        Self {
+            sampling: Duration::from_secs(5),
+            unready: None,
+        }
+    }
+}
+
+/// Configuration options for the readiness health sub-component.
+#[derive(Clone, Debug, Deserialize, Serialize, JsonSchema)]
+#[serde(deny_unknown_fields)]
+#[serde(default)]
+pub struct ReadinessConfig {
+    /// The readiness interval configuration
+    pub interval: ReadinessIntervalConfig,
+
+    /// How many rejections are allowed in an interval (default: 100)
+    /// If this number is exceeded, the server will start to report unready.
+    pub allowed: usize,
+}
+
+impl Default for ReadinessConfig {
+    fn default() -> Self {
+        Self {
+            interval: Default::default(),
+            allowed: 100,
+        }
+    }
+}
+
+/// Configuration options for the health check component.
+#[derive(Debug, Clone, Deserialize, Serialize, JsonSchema)]
+#[serde(deny_unknown_fields)]
+#[serde(default)]
+pub struct HealthCheckConfig {
+    /// Set to false to disable the health check
+    pub enabled: bool,
+
+    /// Optionally set a custom healthcheck path
+    /// Defaults to /health
+    pub path: String,
+
+    /// Optionally specify readiness configuration
+    pub readiness: ReadinessConfig,
+}
+
+impl Default for HealthCheckConfig {
+    fn default() -> Self {
+        Self {
+            enabled: false,
+            path: "/health".to_string(),
+            readiness: Default::default(),
+        }
+    }
+}
+
+#[derive(Clone)]
+pub struct HealthCheck {
+    config: HealthCheckConfig,
+    live: Arc<AtomicBool>,
+    ready: Arc<AtomicBool>,
+    rejected: Arc<AtomicUsize>,
+    ticker: Arc<tokio::task::JoinHandle<()>>,
+}
+
+impl HealthCheck {
+    pub fn new(config: HealthCheckConfig) -> Self {
+        let live = Arc::new(AtomicBool::new(true)); // Start as live
+        let ready = Arc::new(AtomicBool::new(true)); // Start as ready
+        let rejected = Arc::new(AtomicUsize::new(0));
+
+        let allowed = config.readiness.allowed;
+        let sampling_interval = config.readiness.interval.sampling;
+        let recovery_interval = config
+            .readiness
+            .interval
+            .unready
+            .unwrap_or(2 * sampling_interval);
+
+        let my_rejected = rejected.clone();
+        let my_ready = ready.clone();
+
+        let ticker = tokio::spawn(async move {
+            loop {
+                let start = Instant::now() + sampling_interval;
+                let mut interval = tokio::time::interval_at(start, sampling_interval);
+                loop {
+                    interval.tick().await;
+                    if my_rejected.load(Ordering::Relaxed) > allowed {
+                        debug!("Health check readiness threshold exceeded, marking as unready");
+                        my_ready.store(false, Ordering::SeqCst);
+                        tokio::time::sleep(recovery_interval).await;
+                        my_rejected.store(0, Ordering::Relaxed);
+                        my_ready.store(true, Ordering::SeqCst);
+                        debug!("Health check readiness restored");
+                        break;
+                    }
+                }
+            }
+        });
+
+        Self {
+            config,
+            live,
+            ready,
+            rejected,
+            ticker: Arc::new(ticker),
+        }
+    }
+
+    pub fn record_rejection(&self) {
+        self.rejected.fetch_add(1, Ordering::Relaxed);
+    }
+
+    pub fn config(&self) -> &HealthCheckConfig {
+        &self.config
+    }
+
+    pub fn get_health_state(&self, query: Option<&str>) -> (Health, StatusCode) {
+        let mut status_code = StatusCode::OK;
+
+        let health = if let Some(query) = query {
+            let query_upper = query.to_ascii_uppercase();
+
+            if query_upper.starts_with("READY") {
+                let status = if self.ready.load(Ordering::SeqCst) {
+                    HealthStatus::Up
+                } else {
+                    status_code = StatusCode::SERVICE_UNAVAILABLE;
+                    HealthStatus::Down
+                };
+                Health { status }
+            } else if query_upper.starts_with("LIVE") {
+                let status = if self.live.load(Ordering::SeqCst) {
+                    HealthStatus::Up
+                } else {
+                    status_code = StatusCode::SERVICE_UNAVAILABLE;
+                    HealthStatus::Down
+                };
+                Health { status }
+            } else {
+                Health {
+                    status: HealthStatus::Up,
+                }
+            }
+        } else {
+            Health {
+                status: HealthStatus::Up,
+            }
+        };
+
+        (health, status_code)
+    }
+}
+
+impl Drop for HealthCheck {
+    fn drop(&mut self) {
+        self.ticker.abort();
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use tokio::time::{Duration, sleep};
+
+    #[test]
+    fn test_health_check_default_config() {
+        let config = HealthCheckConfig::default();
+        assert!(!config.enabled);
+        assert_eq!(config.path, "/health");
+        assert_eq!(config.readiness.allowed, 100);
+        assert_eq!(config.readiness.interval.sampling, Duration::from_secs(5));
+        assert!(config.readiness.interval.unready.is_none());
+    }
+
+    #[tokio::test]
+    async fn test_health_check_rejection_tracking() {
+        let mut config = HealthCheckConfig::default();
+        config.readiness.allowed = 2;
+        config.readiness.interval.sampling = Duration::from_millis(50);
+        config.readiness.interval.unready = Some(Duration::from_millis(100));
+
+        let health_check = HealthCheck::new(config);
+
+        // Should be live and ready initially
+        assert!(health_check.live.load(Ordering::SeqCst));
+        assert!(health_check.ready.load(Ordering::SeqCst));
+
+        // Record rejections beyond threshold
+        for _ in 0..5 {
+            health_check.record_rejection();
+        }
+
+        // Wait for the ticker to process
+        sleep(Duration::from_millis(100)).await;
+
+        // Should be still live but unready now
+        assert!(health_check.live.load(Ordering::SeqCst));
+        assert!(!health_check.ready.load(Ordering::SeqCst));
+    }
+}

crates/apollo-mcp-server/src/health.rs

@@ -3,6 +3,7 @@ pub mod errors;
 pub mod event;
 mod explorer;
 mod graphql;
+pub mod health;
 mod introspection;
 pub mod json_schema;
 pub mod operations;

crates/apollo-mcp-server/src/lib.rs

@@ -155,6 +155,7 @@ async fn main() -> anyhow::Result<()> {
         )
         .search_leaf_depth(config.introspection.search.leaf_depth)
         .index_memory_bytes(config.introspection.search.index_memory_bytes)
+        .health_check(config.health_check)
         .build()
         .start()
         .await?)

crates/apollo-mcp-server/src/main.rs

@@ -1,6 +1,6 @@
 use std::path::PathBuf;
 
-use apollo_mcp_server::server::Transport;
+use apollo_mcp_server::{health::HealthCheckConfig, server::Transport};
 use reqwest::header::HeaderMap;
 use schemars::JsonSchema;
 use serde::Deserialize;
@@ -30,6 +30,10 @@ pub struct Config {
     #[schemars(schema_with = "super::schemas::header_map")]
     pub headers: HeaderMap,
 
+    /// Health check configuration
+    #[serde(default)]
+    pub health_check: HealthCheckConfig,
+
     /// Introspection configuration
     pub introspection: Introspection,
 

crates/apollo-mcp-server/src/runtime/config.rs

@@ -10,6 +10,7 @@ use url::Url;
 use crate::custom_scalar_map::CustomScalarMap;
 use crate::errors::ServerError;
 use crate::event::Event as ServerEvent;
+use crate::health::HealthCheckConfig;
 use crate::operations::{MutationMode, OperationSource};
 
 mod states;
@@ -36,6 +37,7 @@ pub struct Server {
     disable_schema_description: bool,
     search_leaf_depth: usize,
     index_memory_bytes: usize,
+    health_check: HealthCheckConfig,
 }
 
 #[derive(Debug, Clone, Deserialize, Default, JsonSchema)]
@@ -103,6 +105,7 @@ impl Server {
         disable_schema_description: bool,
         search_leaf_depth: usize,
         index_memory_bytes: usize,
+        health_check: HealthCheckConfig,
     ) -> Self {
         let headers = {
             let mut headers = headers.clone();
@@ -128,6 +131,7 @@ impl Server {
             disable_schema_description,
             search_leaf_depth,
             index_memory_bytes,
+            health_check,
         }
     }
 

crates/apollo-mcp-server/src/server.rs

@@ -8,6 +8,7 @@ use url::Url;
 use crate::{
     custom_scalar_map::CustomScalarMap,
     errors::{OperationError, ServerError},
+    health::HealthCheckConfig,
     operations::MutationMode,
 };
 
@@ -45,6 +46,7 @@ struct Config {
     disable_schema_description: bool,
     search_leaf_depth: usize,
     index_memory_bytes: usize,
+    health_check: HealthCheckConfig,
 }
 
 impl StateMachine {
@@ -76,6 +78,7 @@ impl StateMachine {
                 disable_schema_description: server.disable_schema_description,
                 search_leaf_depth: server.search_leaf_depth,
                 index_memory_bytes: server.index_memory_bytes,
+                health_check: server.health_check,
             },
         });