pulseengine_mcp_server/
health_endpoint.rs

1//! Health check endpoints for Kubernetes and monitoring
2
3use crate::McpServer;
4use crate::backend::McpBackend;
5use axum::{
6    Router,
7    extract::State,
8    http::StatusCode,
9    response::{IntoResponse, Json},
10    routing::get,
11};
12use serde::{Deserialize, Serialize};
13use std::sync::Arc;
14
15/// Health check response
16#[derive(Debug, Serialize, Deserialize)]
17pub struct HealthResponse {
18    pub status: HealthStatus,
19    pub timestamp: u64,
20    pub uptime_seconds: u64,
21    pub version: String,
22    pub checks: Vec<HealthCheck>,
23}
24
25/// Health status enum
26#[derive(Debug, Serialize, Deserialize, PartialEq)]
27#[serde(rename_all = "lowercase")]
28pub enum HealthStatus {
29    Healthy,
30    Degraded,
31    Unhealthy,
32}
33
34/// Individual health check
35#[derive(Debug, Serialize, Deserialize)]
36pub struct HealthCheck {
37    pub name: String,
38    pub status: HealthStatus,
39    pub message: Option<String>,
40    pub duration_ms: u64,
41}
42
43/// Ready check response
44#[derive(Debug, Serialize, Deserialize)]
45pub struct ReadyResponse {
46    pub ready: bool,
47    pub message: Option<String>,
48}
49
50/// Health check state
51pub struct HealthState<B: McpBackend> {
52    pub server: Arc<McpServer<B>>,
53}
54
55/// Handler for /health endpoint (liveness probe)
56pub async fn health_handler<B: McpBackend + 'static>(
57    State(state): State<Arc<HealthState<B>>>,
58) -> impl IntoResponse {
59    let start = std::time::Instant::now();
60
61    // Perform health checks
62    let health_status = state.server.health_check().await;
63
64    let mut checks = Vec::new();
65    let mut overall_status = HealthStatus::Healthy;
66
67    match health_status {
68        Ok(status) => {
69            for (component, healthy) in status.components {
70                let check_status = if healthy {
71                    HealthStatus::Healthy
72                } else {
73                    overall_status = HealthStatus::Unhealthy;
74                    HealthStatus::Unhealthy
75                };
76
77                checks.push(HealthCheck {
78                    name: component,
79                    status: check_status,
80                    message: None,
81                    duration_ms: start.elapsed().as_millis() as u64,
82                });
83            }
84        }
85        Err(e) => {
86            overall_status = HealthStatus::Unhealthy;
87            checks.push(HealthCheck {
88                name: "server".to_string(),
89                status: HealthStatus::Unhealthy,
90                message: Some(e.to_string()),
91                duration_ms: start.elapsed().as_millis() as u64,
92            });
93        }
94    }
95
96    let response = HealthResponse {
97        status: overall_status,
98        timestamp: std::time::SystemTime::now()
99            .duration_since(std::time::UNIX_EPOCH)
100            .unwrap()
101            .as_secs(),
102        uptime_seconds: state.server.get_metrics().await.uptime_seconds,
103        version: env!("CARGO_PKG_VERSION").to_string(),
104        checks,
105    };
106
107    let status_code = match response.status {
108        HealthStatus::Healthy => StatusCode::OK,
109        HealthStatus::Degraded => StatusCode::OK, // Still return 200 for degraded
110        HealthStatus::Unhealthy => StatusCode::SERVICE_UNAVAILABLE,
111    };
112
113    (status_code, Json(response))
114}
115
116/// Handler for /ready endpoint (readiness probe)
117pub async fn ready_handler<B: McpBackend + 'static>(
118    State(state): State<Arc<HealthState<B>>>,
119) -> impl IntoResponse {
120    // Check if server is running and ready to accept requests
121    let is_running = state.server.is_running().await;
122
123    if is_running {
124        // Additional readiness checks
125        match state.server.health_check().await {
126            Ok(status) => {
127                // All components must be healthy for readiness
128                let all_healthy = status.components.values().all(|&healthy| healthy);
129
130                if all_healthy {
131                    (
132                        StatusCode::OK,
133                        Json(ReadyResponse {
134                            ready: true,
135                            message: None,
136                        }),
137                    )
138                } else {
139                    (
140                        StatusCode::SERVICE_UNAVAILABLE,
141                        Json(ReadyResponse {
142                            ready: false,
143                            message: Some("Some components are not healthy".to_string()),
144                        }),
145                    )
146                }
147            }
148            Err(e) => (
149                StatusCode::SERVICE_UNAVAILABLE,
150                Json(ReadyResponse {
151                    ready: false,
152                    message: Some(format!("Health check failed: {e}")),
153                }),
154            ),
155        }
156    } else {
157        (
158            StatusCode::SERVICE_UNAVAILABLE,
159            Json(ReadyResponse {
160                ready: false,
161                message: Some("Server is not running".to_string()),
162            }),
163        )
164    }
165}
166
167/// Create health check router
168pub fn create_health_router<B: McpBackend + 'static>(server: Arc<McpServer<B>>) -> Router {
169    let state = Arc::new(HealthState { server });
170
171    Router::new()
172        .route("/health", get(health_handler::<B>))
173        .route("/ready", get(ready_handler::<B>))
174        .with_state(state)
175}
176
177#[cfg(test)]
178mod tests {
179    use super::*;
180
181    #[test]
182    fn test_health_response_serialization() {
183        let response = HealthResponse {
184            status: HealthStatus::Healthy,
185            timestamp: 1234567890,
186            uptime_seconds: 3600,
187            version: "1.0.0".to_string(),
188            checks: vec![HealthCheck {
189                name: "backend".to_string(),
190                status: HealthStatus::Healthy,
191                message: None,
192                duration_ms: 10,
193            }],
194        };
195
196        let json = serde_json::to_string(&response).unwrap();
197        assert!(json.contains("\"status\":\"healthy\""));
198        assert!(json.contains("\"backend\""));
199    }
200}
pulseengine_mcp_server/health_endpoint.rs

pulseengine_mcp_server/
health_endpoint.rs
