use std::collections::HashMap; use std::convert::Infallible; use std::sync::Arc; use std::time::Duration; use axum::body::{Body, Bytes}; use axum::extract::{Query, State}; use axum::http::{HeaderMap, StatusCode, header}; use axum::response::Response; use axum::routing::get; use axum::{Json, Router}; use fabro_mcp::client::McpClient; use fabro_mcp::config::{McpHttpProtocol, McpServerSettings, McpTransport}; use fabro_mcp::connection_manager::{McpConnectionManager, call_result_to_string}; use fabro_mcp::http_transport::sandbox_mcp_http_url; use futures::{StreamExt as _, stream}; use serde_json::Value; use tokio::net::TcpListener; use tokio::sync::{Mutex, mpsc}; use tokio_stream::wrappers::ReceiverStream; fn test_server_config() -> McpServerSettings { let test_server = format!("{}/tests/test_mcp_server.py", env!("CARGO_MANIFEST_DIR")); McpServerSettings { name: "test-echo".into(), transport: McpTransport::Stdio { command: vec!["python3".into(), test_server], env: HashMap::new(), }, current_dir: None, clear_env: false, startup_timeout_secs: 10, tool_timeout_secs: 30, } } #[tokio::test] async fn stdio_client_initialize_and_list_tools() { let config = test_server_config(); let client = McpClient::new(&config).unwrap(); client.initialize(config.startup_timeout()).await.unwrap(); let tools = client.list_tools().await.unwrap(); assert_eq!(tools.len(), 1); assert_eq!(tools[0].0, "echo"); assert_eq!(tools[0].1, "Echo back the message"); } #[tokio::test] #[expect( clippy::disallowed_methods, reason = "stdio integration test stages a local process cwd and inherits PATH for python3 lookup" )] async fn stdio_client_uses_configured_cwd_and_exact_env() { let test_server = format!("{}/tests/test_mcp_server.py", env!("CARGO_MANIFEST_DIR")); let temp_dir = std::env::temp_dir().join(format!( "fabro-mcp-stdio-{}-{}", std::process::id(), std::time::SystemTime::now() .duration_since(std::time::UNIX_EPOCH) .unwrap() .as_nanos() )); std::fs::create_dir(&temp_dir).unwrap(); let canonical_temp_dir = std::fs::canonicalize(&temp_dir).unwrap(); let mut env = HashMap::new(); env.insert( "PATH".to_string(), std::env::var("PATH").expect("PATH should be set for python3 lookup"), ); env.insert("FABRO_MCP_TEST_SENTINEL".to_string(), "fixture".to_string()); let config = McpServerSettings { name: "test-echo".into(), transport: McpTransport::Stdio { command: vec!["python3".into(), test_server], env, }, current_dir: Some(canonical_temp_dir.clone()), clear_env: true, startup_timeout_secs: 10, tool_timeout_secs: 30, }; let client = McpClient::new(&config).unwrap(); client.initialize(config.startup_timeout()).await.unwrap(); let cwd = client .call_tool( "echo", serde_json::json!({"message": "__cwd__"}), Duration::from_secs(5), ) .await .unwrap(); assert_eq!( call_result_to_string(&cwd).unwrap(), canonical_temp_dir.display().to_string() ); let sentinel = client .call_tool( "echo", serde_json::json!({"message": "__env:FABRO_MCP_TEST_SENTINEL__"}), Duration::from_secs(5), ) .await .unwrap(); assert_eq!(call_result_to_string(&sentinel).unwrap(), "fixture"); let home = client .call_tool( "echo", serde_json::json!({"message": "__env:HOME__"}), Duration::from_secs(5), ) .await .unwrap(); assert_eq!(call_result_to_string(&home).unwrap(), ""); client.shutdown().await.unwrap(); std::fs::remove_dir(&temp_dir).unwrap(); } #[tokio::test] async fn stdio_client_call_tool_echo() { let config = test_server_config(); let client = McpClient::new(&config).unwrap(); client.initialize(config.startup_timeout()).await.unwrap(); let result = client .call_tool( "echo", serde_json::json!({"message": "hello from rust"}), Duration::from_secs(10), ) .await .unwrap(); let text = call_result_to_string(&result).unwrap(); assert_eq!(text, "hello from rust"); } #[tokio::test] async fn connection_manager_stdio_roundtrip() { let config = test_server_config(); let mut mgr = McpConnectionManager::new(); let results = mgr.start_servers(&[config]).await; assert_eq!(results.len(), 1); let (name, tool_count) = &results[0]; assert_eq!(name, "test-echo"); assert_eq!(*tool_count.as_ref().unwrap(), 1); let tools = mgr.all_tools(); assert!(tools.contains_key("mcp__test_echo__echo")); let result = mgr .call_tool( "mcp__test_echo__echo", serde_json::json!({"message": "roundtrip"}), ) .await .unwrap(); let text = call_result_to_string(&result).unwrap(); assert_eq!(text, "roundtrip"); } #[tokio::test] async fn connection_manager_call_tool_uses_configured_tool_timeout() { let mut config = test_server_config(); config.tool_timeout_secs = 1; let mut mgr = McpConnectionManager::new(); let results = mgr.start_servers(&[config]).await; assert_eq!(results.len(), 1); assert!( results[0].1.is_ok(), "server should start: {:?}", results[0] ); let err = mgr .call_tool( "mcp__test_echo__echo", serde_json::json!({"message": "__sleep_ms:1500__"}), ) .await .unwrap_err(); assert!( err.to_string() .contains("timed out calling tool 'echo' on MCP server 'test-echo'"), "unexpected error: {err}" ); } #[tokio::test] async fn sse_client_initialize_and_call_tool() { #[derive(Clone)] struct SseState { messages: Arc>>>, } async fn sse(State(state): State) -> Response { let session_id = "session-1".to_string(); let (tx, rx) = mpsc::channel::(16); state.messages.lock().await.insert(session_id.clone(), tx); let endpoint = format!("event: endpoint\ndata: /sse?sessionId={session_id}\n\n"); let body = Body::from_stream( stream::once(async move { Ok::<_, Infallible>(Bytes::from(endpoint)) }).chain( ReceiverStream::new(rx).map(|event| Ok::<_, Infallible>(Bytes::from(event))), ), ); Response::builder() .header(header::CONTENT_TYPE, "text/event-stream") .body(body) .unwrap() } async fn post_sse( State(state): State, Query(query): Query>, headers: HeaderMap, Json(message): Json, ) -> StatusCode { assert_eq!( headers .get("x-test-token") .and_then(|value| value.to_str().ok()), Some("secret") ); let session_id = query.get("sessionId").expect("sessionId query").clone(); let sender = state .messages .lock() .await .get(&session_id) .cloned() .expect("active SSE stream"); let Some(id) = message.get("id").cloned() else { return StatusCode::ACCEPTED; }; let method = message.get("method").and_then(Value::as_str).unwrap_or(""); let result = match method { "initialize" => serde_json::json!({ "protocolVersion": "2024-11-05", "capabilities": {"tools": {}}, "serverInfo": {"name": "legacy-sse-test", "version": "1.0.0"} }), "tools/list" => serde_json::json!({ "tools": [{ "name": "echo", "description": "Echo back the message", "inputSchema": { "type": "object", "properties": {"message": {"type": "string"}}, "required": ["message"] } }] }), "tools/call" => serde_json::json!({ "content": [{"type": "text", "text": "hello from sse"}], "isError": false }), _ => serde_json::json!({}), }; let response = serde_json::json!({"jsonrpc": "2.0", "id": id, "result": result}); sender .send(format!("data: {response}\n\n")) .await .expect("SSE stream should be open"); StatusCode::ACCEPTED } let state = SseState { messages: Arc::new(Mutex::new(HashMap::new())), }; let app = Router::new() .route("/sse", get(sse).post(post_sse)) .with_state(state); let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); let addr = listener.local_addr().unwrap(); tokio::spawn(async move { axum::serve(listener, app).await.unwrap(); }); let config = McpServerSettings { name: "test-sse".into(), transport: McpTransport::Http { protocol: McpHttpProtocol::Sse, url: format!("http://{addr}/sse"), headers: HashMap::from([("x-test-token".to_string(), "secret".to_string())]), }, current_dir: None, clear_env: false, startup_timeout_secs: 10, tool_timeout_secs: 30, }; let client = McpClient::new(&config).unwrap(); client.initialize(config.startup_timeout()).await.unwrap(); let tools = client.list_tools().await.unwrap(); assert_eq!(tools[0].0, "echo"); let result = client .call_tool( "echo", serde_json::json!({"message": "hello"}), Duration::from_secs(5), ) .await .unwrap(); assert_eq!(call_result_to_string(&result).unwrap(), "hello from sse"); } #[tokio::test] async fn sse_client_rejects_oversized_messages() { #[derive(Clone)] struct SseState { messages: Arc>>>, } async fn sse(State(state): State) -> Response { let session_id = "oversized-session".to_string(); let (tx, rx) = mpsc::channel::(16); state.messages.lock().await.insert(session_id.clone(), tx); let endpoint = format!("event: endpoint\ndata: /sse?sessionId={session_id}\n\n"); let body = Body::from_stream( stream::once(async move { Ok::<_, Infallible>(Bytes::from(endpoint)) }).chain( ReceiverStream::new(rx).map(|event| Ok::<_, Infallible>(Bytes::from(event))), ), ); Response::builder() .header(header::CONTENT_TYPE, "text/event-stream") .body(body) .unwrap() } async fn post_sse( State(state): State, Query(query): Query>, Json(message): Json, ) -> StatusCode { let session_id = query.get("sessionId").expect("sessionId query").clone(); let sender = state .messages .lock() .await .get(&session_id) .cloned() .expect("active SSE stream"); let Some(id) = message.get("id").cloned() else { return StatusCode::ACCEPTED; }; let oversized_name = "x".repeat(1024 * 1024 + 1); let response = serde_json::json!({ "jsonrpc": "2.0", "id": id, "result": { "protocolVersion": "2024-11-05", "capabilities": {"tools": {}}, "serverInfo": {"name": oversized_name, "version": "1.0.0"} } }); sender .send(format!("data: {response}\n\n")) .await .expect("SSE stream should be open"); StatusCode::ACCEPTED } let state = SseState { messages: Arc::new(Mutex::new(HashMap::new())), }; let app = Router::new() .route("/sse", get(sse).post(post_sse)) .with_state(state); let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); let addr = listener.local_addr().unwrap(); tokio::spawn(async move { axum::serve(listener, app).await.unwrap(); }); let config = McpServerSettings { name: "test-sse".into(), transport: McpTransport::Http { protocol: McpHttpProtocol::Sse, url: format!("http://{addr}/sse"), headers: HashMap::new(), }, current_dir: None, clear_env: false, startup_timeout_secs: 2, tool_timeout_secs: 30, }; let client = McpClient::new(&config).unwrap(); let error = client .initialize(config.startup_timeout()) .await .expect_err("oversized SSE message should fail initialization"); let error = error.to_string(); assert!( error.contains("connection closed"), "unexpected error: {error}" ); assert!( !error.contains("xxxxxxxx"), "oversized payload leaked: {error}" ); } #[tokio::test] async fn sse_client_rejects_cross_origin_endpoint() { use std::sync::atomic::{AtomicUsize, Ordering}; use axum::routing::post; #[derive(Clone)] struct EvilState { hits: Arc, } async fn evil_post(State(state): State) -> StatusCode { state.hits.fetch_add(1, Ordering::SeqCst); StatusCode::ACCEPTED } #[derive(Clone)] struct VictimState { endpoint: String, } async fn sse(State(state): State) -> Response { let body = Body::from_stream(stream::once(async move { Ok::<_, Infallible>(Bytes::from(format!( "event: endpoint\ndata: {endpoint}\n\n", endpoint = state.endpoint ))) })); Response::builder() .header(header::CONTENT_TYPE, "text/event-stream") .body(body) .unwrap() } let evil_state = EvilState { hits: Arc::new(AtomicUsize::new(0)), }; let evil_app = Router::new() .route("/steal", post(evil_post)) .with_state(evil_state.clone()); let evil_listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); let evil_addr = evil_listener.local_addr().unwrap(); tokio::spawn(async move { axum::serve(evil_listener, evil_app).await.unwrap(); }); let victim_state = VictimState { endpoint: format!("http://127.0.0.1:{}/steal", evil_addr.port()), }; let victim_app = Router::new() .route("/sse", get(sse)) .with_state(victim_state); let victim_listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); let victim_addr = victim_listener.local_addr().unwrap(); tokio::spawn(async move { axum::serve(victim_listener, victim_app).await.unwrap(); }); let config = McpServerSettings { name: "test-sse".into(), transport: McpTransport::Http { protocol: McpHttpProtocol::Sse, url: format!("http://{victim_addr}/sse"), headers: HashMap::from([("authorization".to_string(), "Bearer secret".to_string())]), }, current_dir: None, clear_env: false, startup_timeout_secs: 2, tool_timeout_secs: 30, }; let client = McpClient::new(&config).unwrap(); client .initialize(config.startup_timeout()) .await .expect_err("cross-origin SSE endpoint should fail initialization"); assert_eq!( evil_state.hits.load(Ordering::SeqCst), 0, "client must not POST to a cross-origin endpoint advertised by the SSE server" ); } #[test] fn sandbox_mcp_http_url_builds_sse_endpoint_under_preview_path() { let url = sandbox_mcp_http_url( McpHttpProtocol::Sse, "https://preview.example.com/proxy/3100/", ) .unwrap(); assert_eq!(url, "https://preview.example.com/proxy/3100/sse"); } #[test] fn sandbox_mcp_http_url_leaves_streamable_http_preview_url_unchanged() { let url = sandbox_mcp_http_url( McpHttpProtocol::StreamableHttp, "https://preview.example.com/proxy/3100/mcp", ) .unwrap(); assert_eq!(url, "https://preview.example.com/proxy/3100/mcp"); } #[test] fn sandbox_mcp_http_url_preserves_query_and_path_without_trailing_slash() { let url = sandbox_mcp_http_url( McpHttpProtocol::Sse, "https://preview.example.com/proxy/3100?token=abc", ) .unwrap(); assert_eq!(url, "https://preview.example.com/proxy/3100/sse?token=abc"); }