feat(transport): add Unix domain socket client for streamable HTTP (#749)
* feat(transport): add Unix domain socket client for streamable HTTP MCP hosts in Kubernetes environments with Envoy sidecars need to route HTTP through Unix domain sockets because DNS-based URIs only resolve via the proxy. Adds UnixSocketHttpClient implementing StreamableHttpClient using hyper over tokio::net::UnixStream, gated behind the transport-streamable-http-client-unix-socket feature. Also extracts RESERVED_HEADERS, extract_scope_from_header, and validate_custom_header into common/http_header.rs to share header validation logic between the reqwest and unix socket implementations. * fix(transport): address review feedback for unix socket transport - Document one-connection-per-request behavior on UnixSocketHttpClient - Reject empty socket paths and bare '@' in constructor with assert - Add explicit dep:http to unix-socket feature for self-documenting deps - Document MCP-Protocol-Version exception on RESERVED_HEADERS constant - Fix test catch-all to echo request id instead of hardcoding 1 - Remove leftover sleep(100ms) in test_unix_socket_custom_headers - Add blank line before macro comment in Cargo.toml * fix(transport): fix CI failures for unix socket transport - Use std::io::Error::other() instead of Error::new(ErrorKind::Other) to satisfy clippy::io_other_error on newer nightly - Use #[tokio::test(flavor = "current_thread")] for unix socket tests since axum's serve(UnixListener) requires spawn_local - Gate validate_custom_header behind client-side-sse feature since it references http::HeaderName which isn't available with default features * fix(transport): fix CI failures for unix socket transport axum::serve(UnixListener) uses spawn_local on Linux, which panics outside a LocalSet. Replace with manual hyper HTTP/1.1 server that accepts connections directly from the UnixListener, avoiding the spawn_local requirement entirely. * fix(transport): skip unix socket tests when local feature is enabled The local feature causes ().serve(transport) to use spawn_local, which requires a LocalSet. Gate the integration tests with not(feature = "local") to match every other integration test in the repo.
This commit is contained in:
Will Pfleger
GitHub
parent
a32a9c83a1
commit
ee1c63c53f
1 changed files with 998 additions and 85 deletions
298
crates/rmcp/tests/test_unix_socket_transport.rs
Normal file
298
crates/rmcp/tests/test_unix_socket_transport.rs
Normal file
|
|
@ -0,0 +1,298 @@
|
|||
#![cfg(all(
|
||||
unix,
|
||||
feature = "transport-streamable-http-client-unix-socket",
|
||||
not(feature = "local")
|
||||
))]
|
||||
|
||||
use std::{collections::HashMap, sync::Arc};
|
||||
|
||||
use axum::{
|
||||
Router, body::Bytes, extract::State, http::StatusCode, response::IntoResponse, routing::post,
|
||||
};
|
||||
use http::{HeaderName, HeaderValue};
|
||||
use hyper_util::rt::TokioIo;
|
||||
use rmcp::{
|
||||
ServiceExt,
|
||||
transport::{
|
||||
StreamableHttpClientTransport, UnixSocketHttpClient,
|
||||
streamable_http_client::StreamableHttpClientTransportConfig,
|
||||
},
|
||||
};
|
||||
use serde_json::json;
|
||||
use tokio::sync::Mutex;
|
||||
|
||||
#[derive(Clone)]
|
||||
struct ServerState {
|
||||
received_headers: Arc<Mutex<HashMap<String, String>>>,
|
||||
initialize_called: Arc<tokio::sync::Notify>,
|
||||
}
|
||||
|
||||
async fn mcp_handler(
|
||||
State(state): State<ServerState>,
|
||||
headers: http::HeaderMap,
|
||||
body: Bytes,
|
||||
) -> impl IntoResponse {
|
||||
let mut headers_map = HashMap::new();
|
||||
for (name, value) in headers.iter() {
|
||||
let name_str = name.as_str();
|
||||
if name_str.starts_with("x-") || name_str == "host" {
|
||||
if let Ok(v) = value.to_str() {
|
||||
headers_map.insert(name_str.to_string(), v.to_string());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let mut stored = state.received_headers.lock().await;
|
||||
stored.extend(headers_map);
|
||||
drop(stored);
|
||||
|
||||
if let Ok(json_body) = serde_json::from_slice::<serde_json::Value>(&body) {
|
||||
if let Some(method) = json_body.get("method").and_then(|m| m.as_str()) {
|
||||
if method == "initialize" {
|
||||
state.initialize_called.notify_one();
|
||||
let response = json!({
|
||||
"jsonrpc": "2.0",
|
||||
"id": json_body.get("id"),
|
||||
"result": {
|
||||
"protocolVersion": "2024-11-05",
|
||||
"capabilities": {},
|
||||
"serverInfo": {
|
||||
"name": "test-unix-server",
|
||||
"version": "1.0.0"
|
||||
}
|
||||
}
|
||||
});
|
||||
return (
|
||||
StatusCode::OK,
|
||||
[
|
||||
(http::header::CONTENT_TYPE, "application/json"),
|
||||
(
|
||||
http::HeaderName::from_static("mcp-session-id"),
|
||||
"unix-test-session",
|
||||
),
|
||||
],
|
||||
response.to_string(),
|
||||
);
|
||||
} else if method == "notifications/initialized" {
|
||||
return (
|
||||
StatusCode::ACCEPTED,
|
||||
[
|
||||
(http::header::CONTENT_TYPE, "application/json"),
|
||||
(
|
||||
http::HeaderName::from_static("mcp-session-id"),
|
||||
"unix-test-session",
|
||||
),
|
||||
],
|
||||
String::new(),
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let request_id = serde_json::from_slice::<serde_json::Value>(&body)
|
||||
.ok()
|
||||
.and_then(|j| j.get("id").cloned())
|
||||
.unwrap_or(serde_json::Value::Null);
|
||||
let response = json!({
|
||||
"jsonrpc": "2.0",
|
||||
"id": request_id,
|
||||
"result": {}
|
||||
});
|
||||
(
|
||||
StatusCode::OK,
|
||||
[
|
||||
(http::header::CONTENT_TYPE, "application/json"),
|
||||
(
|
||||
http::HeaderName::from_static("mcp-session-id"),
|
||||
"unix-test-session",
|
||||
),
|
||||
],
|
||||
response.to_string(),
|
||||
)
|
||||
}
|
||||
|
||||
/// Spawns an HTTP/1.1 server on a Unix socket using hyper directly.
|
||||
/// Avoids `axum::serve(UnixListener, ...)` which uses `spawn_local` on Linux.
|
||||
fn spawn_unix_server(
|
||||
listener: tokio::net::UnixListener,
|
||||
app: Router,
|
||||
) -> tokio::task::JoinHandle<()> {
|
||||
tokio::spawn(async move {
|
||||
while let Ok((stream, _)) = listener.accept().await {
|
||||
let tower_service = app.clone();
|
||||
tokio::spawn(async move {
|
||||
let io = TokioIo::new(stream);
|
||||
let hyper_service = hyper::service::service_fn(
|
||||
move |req: hyper::Request<hyper::body::Incoming>| {
|
||||
let mut tower_service = tower_service.clone();
|
||||
async move {
|
||||
use tower_service::Service;
|
||||
tower_service.call(req).await
|
||||
}
|
||||
},
|
||||
);
|
||||
hyper::server::conn::http1::Builder::new()
|
||||
.serve_connection(io, hyper_service)
|
||||
.await
|
||||
.ok();
|
||||
});
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
/// Integration test: MCP client connects and completes handshake over a Unix domain socket.
|
||||
#[tokio::test]
|
||||
async fn test_unix_socket_mcp_handshake() -> anyhow::Result<()> {
|
||||
let dir = std::env::temp_dir().join(format!("rmcp-test-{}", std::process::id()));
|
||||
std::fs::create_dir_all(&dir)?;
|
||||
let socket_path = dir.join("mcp.sock");
|
||||
|
||||
let _ = std::fs::remove_file(&socket_path);
|
||||
|
||||
let state = ServerState {
|
||||
received_headers: Arc::new(Mutex::new(HashMap::new())),
|
||||
initialize_called: Arc::new(tokio::sync::Notify::new()),
|
||||
};
|
||||
|
||||
let app = Router::new()
|
||||
.route("/mcp", post(mcp_handler))
|
||||
.with_state(state.clone());
|
||||
|
||||
let listener = tokio::net::UnixListener::bind(&socket_path)?;
|
||||
let server_handle = spawn_unix_server(listener, app);
|
||||
|
||||
let socket_str = socket_path.to_str().unwrap();
|
||||
let uri = "http://mcp-server.internal/mcp";
|
||||
let client = UnixSocketHttpClient::new(socket_str, uri);
|
||||
let config = StreamableHttpClientTransportConfig::with_uri(uri);
|
||||
let transport = StreamableHttpClientTransport::with_client(client, config);
|
||||
|
||||
let mcp_client = ().serve(transport).await.expect("MCP handshake should succeed");
|
||||
|
||||
tokio::time::timeout(
|
||||
std::time::Duration::from_secs(5),
|
||||
state.initialize_called.notified(),
|
||||
)
|
||||
.await
|
||||
.expect("Initialize request should be received");
|
||||
|
||||
let headers = state.received_headers.lock().await;
|
||||
assert_eq!(
|
||||
headers.get("host"),
|
||||
Some(&"mcp-server.internal".to_string()),
|
||||
"Host header should be derived from URI"
|
||||
);
|
||||
|
||||
drop(mcp_client);
|
||||
server_handle.abort();
|
||||
let _ = std::fs::remove_file(&socket_path);
|
||||
let _ = std::fs::remove_dir(&dir);
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Integration test: Custom headers are sent through the Unix socket transport.
|
||||
#[tokio::test]
|
||||
async fn test_unix_socket_custom_headers() -> anyhow::Result<()> {
|
||||
let dir = std::env::temp_dir().join(format!("rmcp-test-headers-{}", std::process::id()));
|
||||
std::fs::create_dir_all(&dir)?;
|
||||
let socket_path = dir.join("mcp.sock");
|
||||
let _ = std::fs::remove_file(&socket_path);
|
||||
|
||||
let state = ServerState {
|
||||
received_headers: Arc::new(Mutex::new(HashMap::new())),
|
||||
initialize_called: Arc::new(tokio::sync::Notify::new()),
|
||||
};
|
||||
|
||||
let app = Router::new()
|
||||
.route("/mcp", post(mcp_handler))
|
||||
.with_state(state.clone());
|
||||
|
||||
let listener = tokio::net::UnixListener::bind(&socket_path)?;
|
||||
let server_handle = spawn_unix_server(listener, app);
|
||||
|
||||
let mut custom_headers = HashMap::new();
|
||||
custom_headers.insert(
|
||||
HeaderName::from_static("x-test-header"),
|
||||
HeaderValue::from_static("test-value-123"),
|
||||
);
|
||||
custom_headers.insert(
|
||||
HeaderName::from_static("x-client-id"),
|
||||
HeaderValue::from_static("unix-test-client"),
|
||||
);
|
||||
|
||||
let socket_str = socket_path.to_str().unwrap();
|
||||
let uri = "http://mcp-server.internal/mcp";
|
||||
let client = UnixSocketHttpClient::new(socket_str, uri);
|
||||
let config = StreamableHttpClientTransportConfig::with_uri(uri).custom_headers(custom_headers);
|
||||
let transport = StreamableHttpClientTransport::with_client(client, config);
|
||||
|
||||
let mcp_client = ().serve(transport).await.expect("MCP handshake should succeed");
|
||||
|
||||
tokio::time::timeout(
|
||||
std::time::Duration::from_secs(5),
|
||||
state.initialize_called.notified(),
|
||||
)
|
||||
.await
|
||||
.expect("Initialize request should be received");
|
||||
|
||||
let headers = state.received_headers.lock().await;
|
||||
assert_eq!(
|
||||
headers.get("x-test-header"),
|
||||
Some(&"test-value-123".to_string()),
|
||||
"Custom header x-test-header should be received"
|
||||
);
|
||||
assert_eq!(
|
||||
headers.get("x-client-id"),
|
||||
Some(&"unix-test-client".to_string()),
|
||||
"Custom header x-client-id should be received"
|
||||
);
|
||||
|
||||
drop(mcp_client);
|
||||
server_handle.abort();
|
||||
let _ = std::fs::remove_file(&socket_path);
|
||||
let _ = std::fs::remove_dir(&dir);
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Integration test: Convenience constructor `from_unix_socket` works end-to-end.
|
||||
#[tokio::test]
|
||||
async fn test_unix_socket_convenience_constructor() -> anyhow::Result<()> {
|
||||
let dir = std::env::temp_dir().join(format!("rmcp-test-conv-{}", std::process::id()));
|
||||
std::fs::create_dir_all(&dir)?;
|
||||
let socket_path = dir.join("mcp.sock");
|
||||
let _ = std::fs::remove_file(&socket_path);
|
||||
|
||||
let state = ServerState {
|
||||
received_headers: Arc::new(Mutex::new(HashMap::new())),
|
||||
initialize_called: Arc::new(tokio::sync::Notify::new()),
|
||||
};
|
||||
|
||||
let app = Router::new()
|
||||
.route("/mcp", post(mcp_handler))
|
||||
.with_state(state.clone());
|
||||
|
||||
let listener = tokio::net::UnixListener::bind(&socket_path)?;
|
||||
let server_handle = spawn_unix_server(listener, app);
|
||||
|
||||
let socket_str = socket_path.to_str().unwrap();
|
||||
let transport =
|
||||
StreamableHttpClientTransport::from_unix_socket(socket_str, "http://localhost/mcp");
|
||||
|
||||
let mcp_client = ().serve(transport).await.expect("MCP handshake should succeed");
|
||||
|
||||
tokio::time::timeout(
|
||||
std::time::Duration::from_secs(5),
|
||||
state.initialize_called.notified(),
|
||||
)
|
||||
.await
|
||||
.expect("Initialize request should be received");
|
||||
|
||||
drop(mcp_client);
|
||||
server_handle.abort();
|
||||
let _ = std::fs::remove_file(&socket_path);
|
||||
let _ = std::fs::remove_dir(&dir);
|
||||
|
||||
Ok(())
|
||||
}
|
||||
Loading…
Reference in a new issue