sqlserver_mcp_catalog/http/
server.rs1use std::net::SocketAddr;
4use std::sync::Arc;
5
6use axum::Router;
7use axum::extract::{Request, State};
8use axum::http::StatusCode;
9use axum::middleware::{self, Next};
10use axum::response::{IntoResponse, Json, Response};
11use axum::routing::get;
12use rmcp::transport::streamable_http_server::session::local::LocalSessionManager;
13use rmcp::transport::streamable_http_server::{StreamableHttpServerConfig, StreamableHttpService};
14use serde_json::json;
15use tokio::sync::Mutex as TokioMutex;
16
17use super::metrics;
18use super::types::HttpServerConfig;
19use crate::core::component_registry::{ComponentRegistry, ComponentStatus};
20use crate::core::mcp_server::McpifyServer;
21
22#[derive(Clone)]
23struct AppState {
24 registry: Arc<TokioMutex<ComponentRegistry>>,
25}
26
27async fn healthz(State(state): State<AppState>) -> impl IntoResponse {
28 let registry = state.registry.lock().await;
29 let status = registry.overall_status();
30 let code = if status == ComponentStatus::Unhealthy {
31 StatusCode::SERVICE_UNAVAILABLE
32 } else {
33 StatusCode::OK
34 };
35 (
36 code,
37 Json(json!({
38 "status": format!("{status:?}"),
39 "components": registry.snapshot().len(),
40 })),
41 )
42}
43
44async fn metrics_handler() -> impl IntoResponse {
45 metrics::increment("http_requests_total");
46 (StatusCode::OK, metrics::render_metrics())
47}
48
49async fn count_request(request: Request, next: Next) -> Response {
62 metrics::increment("http_requests_total");
63 next.run(request).await
64}
65
66async fn set_cors_header(cors_allow: Option<String>, request: Request, next: Next) -> Response {
67 let mut response = next.run(request).await;
68 if let Some(origin) = cors_allow
69 && let Ok(value) = origin.parse()
70 {
71 response
72 .headers_mut()
73 .insert("Access-Control-Allow-Origin", value);
74 }
75 response
76}
77
78pub async fn start_http_server(
82 mcp_service_factory: impl Fn() -> Result<McpifyServer, std::io::Error> + Send + Sync + 'static,
83 config: &HttpServerConfig,
84 registry: Arc<TokioMutex<ComponentRegistry>>,
85) -> anyhow::Result<()> {
86 let mut allowed_hosts = vec![
87 "localhost".to_string(),
88 "127.0.0.1".to_string(),
89 "::1".to_string(),
90 ];
91 allowed_hosts.push(config.host.clone());
92 allowed_hosts.push(format!("{}:{}", config.host, config.port));
93
94 let mcp_service = StreamableHttpService::new(
95 mcp_service_factory,
96 LocalSessionManager::default().into(),
97 StreamableHttpServerConfig::default().with_allowed_hosts(allowed_hosts),
98 );
99
100 let state = AppState { registry };
101 let cors_allow = config.cors_allow.clone();
102
103 let mcp_router = Router::new()
104 .nest_service("/mcp", mcp_service)
105 .layer(middleware::from_fn(count_request));
106
107 let app = Router::new()
108 .route("/healthz", get(healthz))
109 .route("/metrics", get(metrics_handler))
110 .merge(mcp_router)
111 .layer(middleware::from_fn(move |req, next| {
112 set_cors_header(cors_allow.clone(), req, next)
113 }))
114 .with_state(state);
115
116 let addr: SocketAddr = format!("{}:{}", config.host, config.port).parse()?;
117 let listener = tokio::net::TcpListener::bind(addr).await?;
118 axum::serve(
119 listener,
120 app.into_make_service_with_connect_info::<SocketAddr>(),
121 )
122 .await?;
123
124 Ok(())
125}
126
127#[cfg(test)]
128mod tests {
129 use axum::body::to_bytes;
130
131 use super::*;
132
133 async fn response_json(response: Response) -> serde_json::Value {
134 let bytes = to_bytes(response.into_body(), usize::MAX).await.unwrap();
135 serde_json::from_slice(&bytes).unwrap()
136 }
137
138 #[tokio::test]
139 async fn healthz_reports_healthy_components() {
140 let mut registry = ComponentRegistry::new();
141 registry.register("store", true);
142 let state = AppState {
143 registry: Arc::new(TokioMutex::new(registry)),
144 };
145
146 let response = healthz(State(state)).await.into_response();
147 assert_eq!(response.status(), StatusCode::OK);
148 assert_eq!(
149 response_json(response).await,
150 serde_json::json!({ "status": "Healthy", "components": 1 })
151 );
152 }
153
154 #[tokio::test]
155 async fn healthz_returns_service_unavailable_for_a_critical_failure() {
156 let mut registry = ComponentRegistry::new();
157 registry.register("store", true);
158 registry.report(
159 "store",
160 ComponentStatus::Unhealthy,
161 Some("unavailable".to_string()),
162 );
163 let state = AppState {
164 registry: Arc::new(TokioMutex::new(registry)),
165 };
166
167 let response = healthz(State(state)).await.into_response();
168 assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE);
169 assert_eq!(response_json(response).await["status"], "Unhealthy");
170 }
171
172 #[tokio::test]
173 async fn metrics_endpoint_returns_prometheus_text() {
174 let response = metrics_handler().await.into_response();
175 assert_eq!(response.status(), StatusCode::OK);
176 let body = to_bytes(response.into_body(), usize::MAX).await.unwrap();
177 assert!(
178 String::from_utf8(body.to_vec())
179 .unwrap()
180 .contains("http_requests_total")
181 );
182 }
183}