use axum::extract::Request;
use axum::http::StatusCode;
use axum::response::{IntoResponse, Response};
use bytes::{BufMut, Bytes, BytesMut};
use metrics::counter;
use std::collections::HashMap;
use std::path::Path;
use std::process::Stdio;
use tokio::io::{AsyncBufReadExt, AsyncReadExt, AsyncWriteExt};
use tokio::net::UnixStream;
use tokio::process::Command;
use tracing::{debug, info, warn};
const SOCKET_TIMEOUT_SECS: u64 = 30;
const MAX_BODY_BYTES: usize = 10_000_000;
#[derive(Debug, Clone, Default, serde::Deserialize, serde::Serialize, PartialEq, Eq)]
#[serde(rename_all = "snake_case")]
pub enum McpBackend {
#[default]
Python,
Wasm,
}
pub async fn start_sidecar(
socket_path: &str,
server_config: Option<&str>,
backend: &McpBackend,
) -> anyhow::Result<()> {
match backend {
McpBackend::Python => start_python_sidecar(socket_path, server_config).await,
McpBackend::Wasm => start_wasm_sidecar(socket_path).await,
}
}
async fn start_python_sidecar(
socket_path: &str,
server_config: Option<&str>,
) -> anyhow::Result<()> {
let sock = Path::new(socket_path);
if let Some(parent) = sock.parent() {
tokio::fs::create_dir_all(parent).await?;
}
let _ = tokio::fs::remove_file(socket_path).await;
let mut args: Vec<String> = vec![
"run".to_string(),
"--with".to_string(),
"portail-mcp".to_string(),
"python".to_string(),
"-m".to_string(),
"portail_mcp.server".to_string(),
"--socket".to_string(),
socket_path.to_string(),
];
if let Some(config) = server_config {
args.push("--config".to_string());
args.push(config.to_string());
}
let child = Command::new("uv")
.args(&args)
.stdout(Stdio::piped())
.stderr(Stdio::piped())
.kill_on_drop(true)
.spawn()
.map_err(|e| anyhow::anyhow!("failed to start MCP sidecar: {e}"))?;
info!(pid = child.id().unwrap_or(0), %socket_path, "MCP Python sidecar started");
tokio::time::sleep(std::time::Duration::from_millis(500)).await;
if !sock.exists() {
warn!("MCP sidecar socket not yet created — will retry on first request");
}
if let Some(stdout) = child.stdout {
tokio::spawn(async move {
let mut reader = tokio::io::BufReader::new(stdout);
let mut buf = String::new();
while reader.read_line(&mut buf).await.unwrap_or(0) > 0 {
debug!("[mcp-sidecar] {}", buf.trim());
buf.clear();
}
});
}
if let Some(stderr) = child.stderr {
tokio::spawn(async move {
let mut reader = tokio::io::BufReader::new(stderr);
let mut buf = String::new();
while reader.read_line(&mut buf).await.unwrap_or(0) > 0 {
debug!("[mcp-sidecar:err] {}", buf.trim());
buf.clear();
}
});
}
Ok(())
}
async fn start_wasm_sidecar(socket_path: &str) -> anyhow::Result<()> {
let sock = Path::new(socket_path);
if let Some(parent) = sock.parent() {
tokio::fs::create_dir_all(parent).await?;
}
let _ = tokio::fs::remove_file(socket_path).await;
let socket_path_owned = socket_path.to_string();
tokio::spawn(async move {
let server = portail_mcp_wasm::WasmMcpServer::new(&socket_path_owned);
if let Err(e) = server.serve().await {
tracing::error!(error = %e, "WASM MCP server stopped");
}
});
for _ in 0..10 {
tokio::time::sleep(std::time::Duration::from_millis(100)).await;
if sock.exists() {
info!(%socket_path, "WASM MCP sidecar started");
return Ok(());
}
}
warn!("WASM MCP sidecar socket not yet created — will retry on first request");
Ok(())
}
pub async fn proxy_to_sidecar(socket_path: &str, req: Request) -> Response {
let (parts, body) = req.into_parts();
let path = parts
.uri
.path_and_query()
.map(|pq| pq.as_str())
.unwrap_or("/");
let method = parts.method.to_string();
let body_bytes = axum::body::to_bytes(body, MAX_BODY_BYTES)
.await
.unwrap_or_default();
let socket_path = socket_path.to_string();
match tokio::time::timeout(
std::time::Duration::from_secs(SOCKET_TIMEOUT_SECS),
proxy_via_unix(&socket_path, &method, path, &parts.headers, &body_bytes),
)
.await
{
Ok(Ok(response)) => {
counter!("mcp_requests", "status" => "success").increment(1);
response
}
Ok(Err(e)) => {
warn!(%socket_path, error = %e, "MCP proxy error");
counter!("mcp_requests", "status" => "error").increment(1);
(StatusCode::BAD_GATEWAY, format!("MCP proxy error: {e}")).into_response()
}
Err(_) => {
warn!(%socket_path, "MCP sidecar timed out");
counter!("mcp_requests", "status" => "timeout").increment(1);
(StatusCode::GATEWAY_TIMEOUT, "MCP sidecar timed out").into_response()
}
}
}
async fn proxy_via_unix(
socket_path: &str,
method: &str,
path: &str,
headers: &axum::http::HeaderMap,
body: &[u8],
) -> anyhow::Result<Response> {
let mut stream = UnixStream::connect(socket_path).await?;
let frame = encode_frame(method, path, headers_to_map(headers), body);
stream.write_all(&frame).await?;
stream.shutdown().await?;
decode_response(&mut stream).await
}
pub fn encode_frame(
method: &str,
path: &str,
headers: HashMap<String, String>,
body: &[u8],
) -> BytesMut {
let headers_json = serde_json::to_string(&headers).unwrap_or_default();
let method = method.as_bytes();
let path = path.as_bytes();
let headers = headers_json.as_bytes();
let mut buf = BytesMut::with_capacity(
2 + method.len() + 4 + path.len() + 4 + headers.len() + 8 + body.len(),
);
buf.put_u16(method.len() as u16);
buf.put_slice(method);
buf.put_u32(path.len() as u32);
buf.put_slice(path);
buf.put_u32(headers.len() as u32);
buf.put_slice(headers);
buf.put_u64(body.len() as u64);
buf.put_slice(body);
buf
}
async fn decode_response(stream: &mut UnixStream) -> anyhow::Result<Response> {
let mut status_hdr = [0u8; 6];
stream.read_exact(&mut status_hdr).await?;
let status_code = u16::from_be_bytes([status_hdr[0], status_hdr[1]]);
let headers_len =
u32::from_be_bytes([status_hdr[2], status_hdr[3], status_hdr[4], status_hdr[5]]) as usize;
let mut headers_buf = vec![0u8; headers_len];
stream.read_exact(&mut headers_buf).await?;
let resp_headers: HashMap<String, String> = serde_json::from_slice(&headers_buf)?;
let mut body_len_buf = [0u8; 8];
stream.read_exact(&mut body_len_buf).await?;
let body_len = u64::from_be_bytes(body_len_buf) as usize;
let mut resp_body = vec![0u8; body_len];
stream.read_exact(&mut resp_body).await?;
let mut out_headers = axum::http::HeaderMap::new();
for (k, v) in resp_headers {
if let (Ok(name), Ok(val)) = (
axum::http::HeaderName::from_bytes(k.as_bytes()),
axum::http::HeaderValue::from_str(&v),
) {
out_headers.insert(name, val);
}
}
Ok((
axum::http::StatusCode::from_u16(status_code).unwrap_or(StatusCode::OK),
out_headers,
Bytes::from(resp_body),
)
.into_response())
}
fn headers_to_map(headers: &axum::http::HeaderMap) -> HashMap<String, String> {
let mut map = HashMap::new();
for (k, v) in headers.iter() {
if let Ok(val) = v.to_str() {
map.insert(k.to_string(), val.to_string());
}
}
map
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn frame_roundtrip() {
let mut headers = HashMap::new();
headers.insert("content-type".into(), "text/plain".into());
let frame = encode_frame("POST", "/mcp/tools/call", headers, b"hello");
assert!(frame.len() > 10);
let method_len = u16::from_be_bytes([frame[0], frame[1]]);
let path_start = 2 + method_len as usize;
let path_len = u32::from_be_bytes([
frame[path_start],
frame[path_start + 1],
frame[path_start + 2],
frame[path_start + 3],
]);
let headers_start = path_start + 4 + path_len as usize;
let headers_len = u32::from_be_bytes([
frame[headers_start],
frame[headers_start + 1],
frame[headers_start + 2],
frame[headers_start + 3],
]);
let body_start = headers_start + 4 + headers_len as usize;
let body_len = u64::from_be_bytes([
frame[body_start],
frame[body_start + 1],
frame[body_start + 2],
frame[body_start + 3],
frame[body_start + 4],
frame[body_start + 5],
frame[body_start + 6],
frame[body_start + 7],
]);
let method_str = std::str::from_utf8(&frame[2..path_start]).unwrap();
assert_eq!(method_str, "POST");
let body_bytes = &frame[body_start + 8..body_start + 8 + body_len as usize];
assert_eq!(body_bytes, b"hello");
}
#[test]
fn headers_to_map_empty() {
let headers = axum::http::HeaderMap::new();
let map = headers_to_map(&headers);
assert!(map.is_empty());
}
#[test]
fn headers_to_map_with_values() {
let mut headers = axum::http::HeaderMap::new();
headers.insert("x-foo", "bar".parse().unwrap());
let map = headers_to_map(&headers);
assert_eq!(map.get("x-foo").unwrap(), "bar");
}
#[test]
fn backend_default_is_python() {
assert_eq!(McpBackend::default(), McpBackend::Python);
}
#[test]
fn backend_serialization() {
let wasm = McpBackend::Wasm;
let json = serde_json::to_string(&wasm).unwrap();
assert_eq!(json, "\"wasm\"");
let parsed: McpBackend = serde_json::from_str("\"python\"").unwrap();
assert_eq!(parsed, McpBackend::Python);
}
}