ai-harness/crates/harness-mcp/tests/echo.rs

//! End-to-end MCP integration test: spawn a real stdio MCP server (a small Python fixture),
//! connect through the real rmcp client, and verify a discovered tool is callable and gated
//! behind the `mcp` permission key. This is the M5 milestone's ✅ for MCP.

use std::collections::HashMap;
use std::sync::{Arc, Mutex};

use harness_core::event::{AppEvent, EventBus};
use harness_core::permission::{PermissionReply, PermissionService};
use harness_core::tool::{MetadataSink, PermissionHandle, ToolCtx, ToolError};
use harness_core::types::{MessageId, SessionId};
use harness_mcp::{connect_all, ServerConfig};
use tokio_util::sync::CancellationToken;

/// A permission frontend that replies `Once` to every ask and records the `permission`/`pattern`
/// of each, so a test can assert the call was actually gated.
fn recording_auto_approve(
    bus: EventBus,
    service: Arc<PermissionService>,
) -> Arc<Mutex<Vec<(String, String)>>> {
    let asks = Arc::new(Mutex::new(Vec::new()));
    let asks_task = asks.clone();
    // Subscribe before spawning: a subscription created inside the task could miss the ask
    // (tokio broadcast only delivers to receivers that exist at publish time).
    let mut rx = bus.subscribe();
    tokio::spawn(async move {
        while let Ok(event) = rx.recv().await {
            if let AppEvent::PermissionAsked { request } = event {
                asks_task
                    .lock()
                    .unwrap()
                    .push((request.permission.clone(), request.pattern.clone()));
                service.reply(&request.id, PermissionReply::Once);
            }
        }
    });
    asks
}

struct Harness {
    ctx_data_dir: std::path::PathBuf,
    service: Arc<PermissionService>,
}

impl Harness {
    fn ctx(&self) -> ToolCtx {
        let (metadata, _rx) = MetadataSink::channel();
        ToolCtx {
            session_id: SessionId::new(),
            message_id: MessageId::new(),
            call_id: "call_1".into(),
            data_dir: self.ctx_data_dir.clone(),
            cwd: std::env::temp_dir(),
            cancel: CancellationToken::new(),
            ask: PermissionHandle::new(
                self.service.clone(),
                SessionId::new(),
                Vec::new(),
                Arc::new(Mutex::new(Vec::new())),
                CancellationToken::new(),
            ),
            metadata,
            spawner: None,
            context_reporter: None,
            diagnostics: None,
        }
    }
}

fn fixture_server() -> HashMap<String, ServerConfig> {
    let script = concat!(env!("CARGO_MANIFEST_DIR"), "/tests/fixtures/echo_server.py");
    HashMap::from([(
        "fix".to_string(),
        ServerConfig {
            command: "python3".to_string(),
            args: vec![script.to_string()],
            env: HashMap::new(),
        },
    )])
}

#[tokio::test]
async fn discovers_and_calls_a_real_mcp_tool_with_permission() {
    let tools = connect_all(fixture_server()).await;
    let names: Vec<_> = tools.iter().map(|t| t.name().to_string()).collect();
    assert!(
        names.contains(&"fix_echo".to_string()),
        "expected fix_echo among {names:?}"
    );
    assert!(names.contains(&"fix_boom".to_string()));

    let echo = tools.iter().find(|t| t.name() == "fix_echo").unwrap();
    // Schema passes through untouched from the server.
    assert_eq!(echo.parameters()["properties"]["text"]["type"], "string");

    let bus = EventBus::new();
    let service = Arc::new(PermissionService::new(bus.clone()));
    let asks = recording_auto_approve(bus, service.clone());
    let dir = tempfile::tempdir().unwrap();
    let harness = Harness {
        ctx_data_dir: dir.path().to_path_buf(),
        service,
    };

    let out = echo
        .execute(serde_json::json!({"text": "hi there"}), harness.ctx())
        .await
        .expect("echo call succeeds");
    assert_eq!(out.output, "hi there");

    // The call was gated on the `mcp` key with the qualified tool name as the pattern.
    let recorded = asks.lock().unwrap().clone();
    assert_eq!(recorded, vec![("mcp".to_string(), "fix_echo".to_string())]);
}

#[tokio::test]
async fn tool_error_result_maps_to_tool_error() {
    let tools = connect_all(fixture_server()).await;
    let boom = tools.iter().find(|t| t.name() == "fix_boom").unwrap();

    let bus = EventBus::new();
    let service = Arc::new(PermissionService::new(bus.clone()));
    let _asks = recording_auto_approve(bus, service.clone());
    let dir = tempfile::tempdir().unwrap();
    let harness = Harness {
        ctx_data_dir: dir.path().to_path_buf(),
        service,
    };

    let err = boom
        .execute(serde_json::json!({}), harness.ctx())
        .await
        .expect_err("boom reports an error result");
    match err {
        ToolError::Other(msg) => assert_eq!(msg, "kaboom"),
        other => panic!("expected ToolError::Other, got {other:?}"),
    }
}

#[tokio::test]
async fn a_failed_server_is_skipped_not_fatal() {
    let servers = HashMap::from([(
        "broken".to_string(),
        ServerConfig {
            command: "definitely-not-a-real-binary-xyz".to_string(),
            args: vec![],
            env: HashMap::new(),
        },
    )]);
    // No panic, no tools — the missing server is logged and skipped.
    let tools = connect_all(servers).await;
    assert!(tools.is_empty());
}