fabro/lib/components/fabro-mcp/tests/stdio_integration.rs

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<Mutex<HashMap<String, mpsc::Sender<String>>>>,
    }

    async fn sse(State(state): State<SseState>) -> Response {
        let session_id = "session-1".to_string();
        let (tx, rx) = mpsc::channel::<String>(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<SseState>,
        Query(query): Query<HashMap<String, String>>,
        headers: HeaderMap,
        Json(message): Json<Value>,
    ) -> 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<Mutex<HashMap<String, mpsc::Sender<String>>>>,
    }

    async fn sse(State(state): State<SseState>) -> Response {
        let session_id = "oversized-session".to_string();
        let (tx, rx) = mpsc::channel::<String>(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<SseState>,
        Query(query): Query<HashMap<String, String>>,
        Json(message): Json<Value>,
    ) -> 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<AtomicUsize>,
    }

    async fn evil_post(State(state): State<EvilState>) -> StatusCode {
        state.hits.fetch_add(1, Ordering::SeqCst);
        StatusCode::ACCEPTED
    }

    #[derive(Clone)]
    struct VictimState {
        endpoint: String,
    }

    async fn sse(State(state): State<VictimState>) -> 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");
}