mentra 0.12.0

//! Tests for MCP registration through the runtime builder.

use serde_json::json;

use crate::mcp::sse::testing::SseTestServer;
use crate::mcp::{McpManager, McpSseServerConfig, mcp_tool_name};

/// Scripts a fixture through the handshake, advertising the given tools.
///
/// Returns the manager alongside the bridged tools: the manager owns the
/// connection those tools call through, so dropping it would close the stream.
async fn connect_sse(
    server: &SseTestServer,
    tools: serde_json::Value,
) -> (Vec<crate::mcp::McpBridgedTool>, McpManager) {
    let config = McpSseServerConfig::new("obs", server.sse_url());
    let connecting = tokio::spawn(async move {
        let mut manager = McpManager::new();
        let bridged = manager.connect_sse(&config).await;
        (bridged, manager)
    });

    server.wait_for_stream();
    server.send_endpoint("/messages/?session_id=abc");

    server.wait_for_posts(1);
    server.send_message(&json!({
        "jsonrpc": "2.0",
        "id": 1,
        "result": {
            "protocolVersion": "2024-11-05",
            "capabilities": {"tools": {}},
            "serverInfo": {"name": "fixture", "version": "4.5.6"}
        }
    }));

    server.wait_for_posts(3);
    server.send_message(&json!({
        "jsonrpc": "2.0",
        "id": 2,
        "result": {"tools": tools}
    }));

    let (bridged, manager) = connecting.await.expect("no panic");
    (bridged.expect("the handshake should succeed"), manager)
}

#[tokio::test(flavor = "multi_thread")]
async fn the_manager_bridges_sse_tools_under_a_namespaced_name() {
    let server = SseTestServer::start();
    let (bridged, _manager) = connect_sse(
        &server,
        json!([
            {"name": "search_logs", "description": "Search logs", "inputSchema": {"type": "object"}},
            {"name": "list_alerts", "inputSchema": {"type": "object"}}
        ]),
    )
    .await;

    use crate::tool::ToolDefinition;
    let names: Vec<String> = bridged
        .iter()
        .map(|tool| tool.descriptor().name.to_string())
        .collect();

    assert_eq!(
        names,
        vec![
            mcp_tool_name("obs", "search_logs"),
            mcp_tool_name("obs", "list_alerts"),
        ],
        "SSE tools must be namespaced exactly like stdio tools"
    );
}

#[tokio::test(flavor = "multi_thread")]
async fn a_bridged_sse_tool_carries_its_description_and_schema() {
    let server = SseTestServer::start();
    let (bridged, _manager) = connect_sse(
        &server,
        json!([{
            "name": "search_logs",
            "description": "Search the log corpus",
            "inputSchema": {
                "type": "object",
                "properties": {"query": {"type": "string"}},
                "required": ["query"]
            }
        }]),
    )
    .await;

    use crate::tool::ToolDefinition;
    let descriptor = bridged[0].descriptor();
    assert_eq!(
        descriptor.description.as_deref(),
        Some("Search the log corpus")
    );
    assert_eq!(
        descriptor.input_schema["properties"]["query"]["type"], "string",
        "the server's schema must reach the model unchanged"
    );
}

#[tokio::test(flavor = "multi_thread")]
async fn the_manager_reports_a_connected_sse_server() {
    let server = SseTestServer::start();
    let config = McpSseServerConfig::new("obs", server.sse_url());
    let connecting = tokio::spawn(async move {
        let mut manager = McpManager::new();
        let result = manager.connect_sse(&config).await;
        (result.map(|tools| tools.len()), manager)
    });

    server.wait_for_stream();
    server.send_endpoint("/messages/?session_id=abc");
    server.wait_for_posts(1);
    server.send_message(&json!({
        "jsonrpc": "2.0",
        "id": 1,
        "result": {
            "protocolVersion": "2024-11-05",
            "capabilities": {},
            "serverInfo": {"name": "fixture", "version": "4.5.6"}
        }
    }));
    server.wait_for_posts(3);
    server.send_message(&json!({
        "jsonrpc": "2.0",
        "id": 2,
        "result": {"tools": [{"name": "search", "inputSchema": {"type": "object"}}]}
    }));

    let (count, manager) = connecting.await.expect("no panic");
    assert_eq!(count.expect("the handshake should succeed"), 1);

    assert!(manager.is_connected("obs"));
    assert_eq!(manager.connected_count(), 1);
    assert_eq!(
        manager.all_tool_names(),
        vec![mcp_tool_name("obs", "search")]
    );

    let summary = manager
        .list_servers()
        .into_iter()
        .find(|summary| summary.name == "obs")
        .expect("the server should be listed");
    assert_eq!(summary.status, crate::mcp::McpServerStatus::Connected);
    assert_eq!(summary.server_version.as_deref(), Some("4.5.6"));
    assert_eq!(summary.tool_count, 1);
}

#[tokio::test(flavor = "multi_thread")]
async fn a_failed_sse_connection_is_recorded_as_an_error() {
    let server = SseTestServer::with_opening(crate::mcp::sse::testing::StreamOpening::Status {
        code: 404,
        body: "no such stream".to_string(),
    });

    let mut manager = McpManager::new();
    let config = McpSseServerConfig::new("obs", server.sse_url());
    manager
        .connect_sse(&config)
        .await
        .expect_err("a 404 must not connect");

    assert!(!manager.is_connected("obs"));
    let summary = manager
        .list_servers()
        .into_iter()
        .find(|summary| summary.name == "obs")
        .expect("an errored server should still be listed");
    assert_eq!(summary.status, crate::mcp::McpServerStatus::Error);
    assert!(summary.error.is_some());
}

#[tokio::test(flavor = "multi_thread")]
async fn a_misconfigured_sse_server_fails_before_any_connection_is_opened() {
    let mut manager = McpManager::new();
    // A token over plaintext to a non-loopback host is refused at validation.
    let config =
        McpSseServerConfig::new("obs", "http://internal.corp/sse").with_bearer_token("secret");

    let error = manager
        .connect_sse(&config)
        .await
        .expect_err("validation must reject this before dialing");
    assert!(
        !error.to_string().contains("secret"),
        "the error must not echo the credential: {error}"
    );
}

/// `build_async` must actually register the MCP tools it connects, and a
/// runtime built without any MCP server must not gain namespaced tools.
///
/// Without this, disabling the whole registration arm in `build_async` leaves
/// the suite green: the runtime still builds, it just silently advertises
/// nothing. That failure is invisible until an agent cannot find its tools.
#[tokio::test(flavor = "multi_thread")]
async fn build_async_registers_the_tools_of_a_connected_sse_server() {
    use crate::Runtime;

    let server = SseTestServer::start();
    let config = McpSseServerConfig::new("obs", server.sse_url());

    let building = tokio::spawn(async move {
        Runtime::empty_builder()
            .with_provider_instance(StubProvider)
            .with_mcp_sse_server(config)
            .build_async()
            .await
    });

    server.wait_for_stream();
    server.send_endpoint("/messages/?session_id=abc");
    server.wait_for_posts(1);
    server.send_message(&json!({
        "jsonrpc": "2.0",
        "id": 1,
        "result": {
            "protocolVersion": "2024-11-05",
            "capabilities": {},
            "serverInfo": {"name": "fixture", "version": "4.5.6"}
        }
    }));
    server.wait_for_posts(3);
    server.send_message(&json!({
        "jsonrpc": "2.0",
        "id": 2,
        "result": {"tools": [
            {"name": "search_logs", "inputSchema": {"type": "object"}},
            {"name": "list_alerts", "inputSchema": {"type": "object"}}
        ]}
    }));

    let runtime = building
        .await
        .expect("no panic")
        .expect("the runtime should build");

    let registered: Vec<String> = runtime
        .tools()
        .into_iter()
        .map(|tool| tool.name.to_string())
        .filter(|name| name.starts_with("mcp__"))
        .collect();

    assert_eq!(
        registered.len(),
        2,
        "build_async must register every tool the server advertised, got {registered:?}"
    );
    assert!(registered.contains(&mcp_tool_name("obs", "search_logs")));
    assert!(registered.contains(&mcp_tool_name("obs", "list_alerts")));

    assert!(
        runtime
            .tool_descriptor(&mcp_tool_name("obs", "search_logs"))
            .is_some(),
        "a registered MCP tool must be resolvable by name"
    );
}

#[tokio::test(flavor = "multi_thread")]
async fn build_async_registers_no_mcp_tools_without_a_configured_server() {
    use crate::Runtime;

    let runtime = Runtime::empty_builder()
        .with_provider_instance(StubProvider)
        .build_async()
        .await
        .expect("the runtime should build");

    let registered: Vec<String> = runtime
        .tools()
        .into_iter()
        .map(|tool| tool.name.to_string())
        .filter(|name| name.starts_with("mcp__"))
        .collect();

    assert!(
        registered.is_empty(),
        "no MCP server was configured, got {registered:?}"
    );
}

/// A provider that satisfies the builder's "at least one provider" check.
///
/// The registration tests never send a request, so it only needs to exist.
#[derive(Clone)]
struct StubProvider;

#[async_trait::async_trait]
impl crate::provider::Provider for StubProvider {
    fn descriptor(&self) -> crate::provider::ProviderDescriptor {
        crate::provider::ProviderDescriptor::new(crate::BuiltinProvider::Anthropic)
    }

    async fn list_models(&self) -> Result<Vec<crate::ModelInfo>, crate::provider::ProviderError> {
        Ok(vec![crate::ModelInfo::new(
            "stub-model",
            crate::BuiltinProvider::Anthropic,
        )])
    }

    async fn stream(
        &self,
        _request: crate::provider::Request<'_>,
    ) -> Result<crate::provider::ProviderEventStream, crate::provider::ProviderError> {
        let (_tx, rx) = tokio::sync::mpsc::unbounded_channel();
        Ok(rx)
    }
}

/// An SSE-backed tool must reach the model through exactly the same result
/// limiter and paging path as a stdio tool or a custom tool. This runs a real
/// fixture server behind a bridged tool inside a scripted runtime and compares
/// its transcript entry byte for byte against a custom tool returning the same
/// text.
#[tokio::test(flavor = "multi_thread")]
async fn sse_tool_output_is_limited_exactly_like_a_custom_tool() {
    use std::collections::BTreeMap;

    use crate::{
        ContentBlock,
        runtime::RuntimePolicy,
        test::{MockRuntime, MockToolCall},
        tool::ToolDefinition,
    };

    let full_output = "one\ntwo\nthree";
    let server = SseTestServer::start();
    let (bridged, _manager) = connect_sse(
        &server,
        json!([{"name": "large_output", "inputSchema": {"type": "object"}}]),
    )
    .await;
    let bridged_name = bridged[0].descriptor().name.to_string();

    let mock = MockRuntime::builder()
        .with_policy(
            RuntimePolicy::permissive()
                .with_max_tool_result_bytes(8)
                .with_max_tool_result_lines(1)
                .spill_full_tool_output(false),
        )
        .tool_calls([
            MockToolCall::new(&bridged_name, json!({})).with_id("sse-call"),
            MockToolCall::new("matching_custom_output", json!({})).with_id("custom-call"),
        ])
        .text("done")
        .build()
        .expect("build mock runtime");

    for tool in bridged {
        mock.runtime().register_tool(tool);
    }
    mock.runtime().register_tool(EchoTool {
        output: full_output.to_string(),
    });

    let mut agent = mock
        .runtime()
        .spawn("mcp-sse-truncation-test", mock.model())
        .expect("spawn agent");

    let running = tokio::spawn(async move {
        let response = agent
            .send(vec![ContentBlock::text("run both tools")])
            .await
            .expect("run agent");
        (response, agent)
    });

    // Answer the bridged tool call once the runtime has posted it.
    server.wait_for_posts(4);
    server.send_message(&json!({
        "jsonrpc": "2.0",
        "id": 3,
        "result": {
            "content": [{"type": "text", "text": full_output}],
            "isError": false
        }
    }));

    let (response, _agent) = running.await.expect("no panic");
    assert_eq!(response.text(), "done");

    let requests = mock.recorded_requests().await;
    let provider_results = requests[1]
        .messages
        .iter()
        .flat_map(|message| &message.content)
        .filter_map(|block| match block {
            ContentBlock::ToolResult {
                tool_use_id,
                content,
                is_error,
            } => Some((tool_use_id.as_str(), (content.as_str(), *is_error))),
            _ => None,
        })
        .collect::<BTreeMap<_, _>>();

    let sse_result = provider_results
        .get("sse-call")
        .expect("the SSE tool result should reach the provider");
    let custom_result = provider_results
        .get("custom-call")
        .expect("the custom tool result should reach the provider");

    assert_eq!(
        sse_result, custom_result,
        "an SSE tool must be limited identically to any other tool"
    );
    assert_eq!(
        *sse_result,
        (
            "one\n[truncated: showing 1 of 3 lines; full output was not saved because spill-to-file is disabled by runtime policy]",
            false,
        )
    );
}

/// A custom tool returning a fixed string, used as the limiter baseline.
struct EchoTool {
    output: String,
}

impl crate::tool::ToolDefinition for EchoTool {
    fn descriptor(&self) -> crate::tool::ToolSpec {
        crate::tool::ToolSpec::builder("matching_custom_output")
            .description("Return the same output as the MCP test tool")
            .input_schema(json!({ "type": "object", "properties": {} }))
            .side_effect_level(crate::tool::ToolSideEffectLevel::External)
            .build()
    }
}

#[async_trait::async_trait]
impl crate::tool::ToolExecutor for EchoTool {
    async fn execute(
        &self,
        _ctx: crate::tool::ParallelToolContext,
        _input: serde_json::Value,
    ) -> crate::tool::ToolResult {
        Ok(self.output.clone())
    }
}