hanzo-mcp 1.1.10

use crate::{command::CommandWrappedInShellBuilder, error::McpError, utils::disect_command};

type Result<T> = std::result::Result<T, McpError>;
use rmcp::{
    model::{CallToolRequestParam, CallToolResult, ClientCapabilities, ClientInfo, Implementation, Tool},
    transport::{SseClientTransport, StreamableHttpClientTransport, TokioChildProcess},
    ServiceExt,
};
use std::collections::HashMap;
use tokio::process::Command;

pub async fn list_tools_via_command(cmd_str: &str, config: Option<HashMap<String, String>>) -> Result<Vec<Tool>> {
    let (env_vars, cmd_executable, cmd_args) = disect_command(cmd_str.to_string());
    let (adapted_program, adapted_args, adapted_envs) =
        CommandWrappedInShellBuilder::wrap_in_shell_as_values(cmd_executable, Some(cmd_args), Some(env_vars));
    let mut cmd = Command::new(adapted_program);
    cmd.kill_on_drop(true);
    cmd.envs(adapted_envs);
    cmd.envs(config.unwrap_or_default());
    cmd.args(adapted_args);

    // Retain the TokioChildProcess so we can wait on it after cancellation
    let child_process = TokioChildProcess::new(cmd).map_err(|e| McpError {
        message: format!("{}", e),
    })?;
    let service = ().serve(child_process).await.map_err(|e| McpError {
        message: format!("{}", e),
    })?;
    // 2. Initialize the MCP server
    service.peer_info();

    // 3. Call the standard MCP `list_tools` method
    let tools = service
        .list_all_tools()
        .await
        .inspect_err(|e| log::error!("error listing tools: {:?}", e));

    // 4. Gracefully shut down the service (drops stdio, child should exit)
    let _ = service
        .cancel()
        .await
        .inspect_err(|e| log::error!("error cancelling sse service: {:?}", e));

    Ok(tools.unwrap())
}

pub async fn list_tools_via_sse(sse_url: &str, _config: Option<HashMap<String, String>>) -> Result<Vec<Tool>> {
    // TODO: The config parameter is not currently used by SseTransport or ClientInfo setup in the example.
    // It might be used in the future for authentication headers or other SSE-specific configurations.
    let transport = SseClientTransport::start(sse_url).await.map_err(|e| McpError {
        message: format!("{}", e),
    })?;
    let client_info = ClientInfo {
        protocol_version: Default::default(),
        capabilities: ClientCapabilities::default(),
        client_info: Implementation {
            name: "hanzo_node_sse_client".to_string(),
            version: env!("CARGO_PKG_VERSION").to_string(),
                icons: None,
                title: None,
                website_url: None,
        },
    };
    let client = client_info.serve(transport).await.map_err(|e| McpError {
        message: format!("SSE client connection error: {:?}", e),
    })?;

    // Initialize and log server info (optional, but good for debugging)
    let _ = client.peer_info();

    // List tools
    let tools_result = client
        .list_all_tools()
        .await
        .inspect_err(|e| log::error!("error listing tools: {:?}", e));

    // Gracefully shut down the client
    let _ = client
        .cancel()
        .await
        .inspect_err(|e| log::error!("error cancelling sse service: {:?}", e));

    Ok(tools_result.unwrap())
}

pub async fn list_tools_via_http(sse_url: &str, _config: Option<HashMap<String, String>>) -> Result<Vec<Tool>> {
    // TODO: The config parameter is not currently used by SseTransport or ClientInfo setup in the example.
    // It might be used in the future for authentication headers or other SSE-specific configurations.
    let transport = StreamableHttpClientTransport::from_uri(sse_url);
    let client_info = ClientInfo {
        protocol_version: Default::default(),
        capabilities: ClientCapabilities::default(),
        client_info: Implementation {
            name: "hanzo_node_http_client".to_string(),
            version: env!("CARGO_PKG_VERSION").to_string(),
                icons: None,
                title: None,
                website_url: None,
        },
    };
    let client = client_info.serve(transport).await.map_err(|e| McpError {
        message: format!("HTTP client connection error: {:?}", e),
    })?;

    // Initialize and log server info (optional, but good for debugging)
    let _ = client.peer_info();

    // List tools
    let tools_result = client
        .list_all_tools()
        .await
        .inspect_err(|e| log::error!("error listing tools: {:?}", e));

    // Gracefully shut down the client
    let _ = client
        .cancel()
        .await
        .inspect_err(|e| log::error!("error cancelling http service: {:?}", e));

    Ok(tools_result.unwrap())
}

pub async fn run_tool_via_command(
    command: String,
    tool: String,
    env_vars: HashMap<String, String>,
    parameters: serde_json::Map<String, serde_json::Value>,
) -> Result<CallToolResult> {
    let (_, cmd_executable, cmd_args) = disect_command(command);

    println!("cmd_executable: {}", cmd_executable);
    println!("env_vars: {:?}", env_vars);
    println!("cmd_args: {:?}", cmd_args);

    // Use the wrap_in_shell_as_values function to prepare the command
    let (adapted_program, adapted_args, adapted_envs) =
        CommandWrappedInShellBuilder::wrap_in_shell_as_values(cmd_executable, Some(cmd_args), Some(env_vars.clone()));

    let mut cmd = Command::new(adapted_program);
    cmd.kill_on_drop(true);
    cmd.envs(adapted_envs);
    cmd.envs(env_vars);
    cmd.args(adapted_args);

    let service = ()
        .serve(TokioChildProcess::new(cmd).map_err(|e| McpError {
            message: format!("{}", e),
        })?)
        .await
        .map_err(|e| McpError {
            message: format!("{}", e),
        })?;
    service.peer_info();

    let call_tool_result = service
        .call_tool(CallToolRequestParam {
            name: tool.into(),
            arguments: Some(parameters),
        })
        .await;

    let _ = service
        .cancel()
        .await
        .inspect_err(|e| log::error!("error cancelling stdio service: {:?}", e));
    Ok(call_tool_result.map_err(|e| McpError {
        message: format!("{}", e),
    })?)
}

pub async fn run_tool_via_sse(
    url: String,
    tool: String,
    parameters: serde_json::Map<String, serde_json::Value>,
) -> Result<CallToolResult> {
    let transport = SseClientTransport::start(url)
        .await
        .inspect_err(|e| log::error!("error starting sse transport: {:?}", e))
        .map_err(|e| McpError {
            message: format!("{}", e),
        })?;

    let client_info = ClientInfo {
        protocol_version: Default::default(),
        capabilities: ClientCapabilities::default(),
        client_info: Implementation {
            name: "Hanzo Node Client".to_string(),
            version: "0.0.1".to_string(),
            icons: None,
            title: None,
            website_url: None,
        },
    };
    let client = client_info
        .serve(transport)
        .await
        .inspect_err(|e| {
            log::error!("client error: {:?}", e);
        })
        .map_err(|e| McpError {
            message: format!("{}", e),
        })?;

    // Initialize
    let server_info = client.peer_info();
    log::info!("connected to server: {server_info:#?}");

    let call_tool_result = client
        .call_tool(CallToolRequestParam {
            name: tool.into(),
            arguments: Some(parameters),
        })
        .await
        .inspect_err(|e| log::error!("error calling tool: {:?}", e));
    let _ = client
        .cancel()
        .await
        .inspect_err(|e| log::error!("error cancelling sse service: {:?}", e));
    Ok(call_tool_result.map_err(|e| McpError {
        message: format!("{}", e),
    })?)
}

pub async fn run_tool_via_http(
    url: String,
    tool: String,
    parameters: serde_json::Map<String, serde_json::Value>,
) -> Result<CallToolResult> {
    let transport = StreamableHttpClientTransport::from_uri(url);

    let client_info = ClientInfo {
        protocol_version: Default::default(),
        capabilities: ClientCapabilities::default(),
        client_info: Implementation {
            name: "Hanzo Node HTTP Client".to_string(),
            version: env!("CARGO_PKG_VERSION").to_string(),
                icons: None,
                title: None,
                website_url: None,
        },
    };
    let client = client_info
        .serve(transport)
        .await
        .inspect_err(|e| {
            log::error!("client error: {:?}", e);
        })
        .map_err(|e| McpError {
            message: format!("{}", e),
        })?;

    // Initialize
    let server_info = client.peer_info();
    log::info!("connected to server: {server_info:#?}");

    let call_tool_result = client
        .call_tool(CallToolRequestParam {
            name: tool.into(),
            arguments: Some(parameters),
        })
        .await
        .inspect_err(|e| log::error!("error calling tool: {:?}", e));
    let _ = client
        .cancel()
        .await
        .inspect_err(|e| log::error!("error cancelling sse service: {:?}", e));
    Ok(call_tool_result.map_err(|e| McpError {
        message: format!("{}", e),
    })?)
}

#[cfg(test)]
pub mod tests_mcp_manager {
    use super::*;
    use serde_json::json;

    #[tokio::test]
    async fn test_run_tool_via_command() {
        let params = json!({
            "a": 1,
            "b": 2,
        });
        let params_map = params.as_object().unwrap().clone();

        let result = run_tool_via_command(
            "npx -y @modelcontextprotocol/server-everything@2025.9.12".to_string(),
            "add".to_string(),
            HashMap::new(),
            params_map,
        )
        .await
        .inspect_err(|e| {
            println!("error {:?}", e);
        });

        assert!(result.is_ok());
        let unwrapped = result.unwrap();
        assert_eq!(unwrapped.content.len(), 1);
        assert!(unwrapped.content[0].as_text().unwrap().text.contains("3"));
    }

    #[tokio::test]
    async fn test_run_tool_via_sse() {
        let mut envs = HashMap::new();
        envs.insert("PORT".to_string(), "8000".to_string());
        let (adapted_program, adapted_args, adapted_envs) = CommandWrappedInShellBuilder::wrap_in_shell_as_values(
            "npx".to_string(),
            Some(vec![
                "-y".to_string(),
                "@modelcontextprotocol/server-everything@2025.9.12".to_string(),
                "sse".to_string(),
            ]) as Option<Vec<String>>,
            Some(envs),
        );

        let _child_result = Command::new(adapted_program)
            .args(adapted_args)
            .envs(adapted_envs)
            .kill_on_drop(true)
            .spawn()
            .inspect_err(|e| {
                println!("error {:?}", e);
            });
        tokio::time::sleep(std::time::Duration::from_secs(3)).await;
        let params = json!({
            "a": 1,
            "b": 2,
        });
        let params_map = params.as_object().unwrap().clone();

        let result = run_tool_via_sse("http://localhost:8000/sse".to_string(), "add".to_string(), params_map)
            .await
            .inspect_err(|e| {
                println!("error {:?}", e);
            });
        match result {
            Ok(result) => {
                assert!(result.content.len() == 1);
                assert!(result.content[0].as_text().unwrap().text.contains("3"));
            }
            Err(e) => {
                println!("error {:?}", e);
                assert!(false);
            }
        }
    }

    #[tokio::test]
    async fn test_list_tools_via_command() {
        let result = list_tools_via_command("npx -y @modelcontextprotocol/server-everything@2025.9.12", None).await;
        assert!(result.is_ok());
        let unwrapped = result.unwrap();

        // Debug output to see actual tools
        println!("Actual number of tools: {}", unwrapped.len());
        println!(
            "Actual tools: {:?}",
            unwrapped.iter().map(|t| &t.name).collect::<Vec<_>>()
        );

        // The MCP server-everything package now returns 10 tools
        assert_eq!(
            unwrapped.len(),
            10,
            "Expected exactly 10 tools, got {}",
            unwrapped.len()
        );

        let expected_tools = [
            "echo",
            "add",
            "longRunningOperation",
            "printEnv",
            "sampleLLM",
            "getTinyImage",
            "annotatedMessage",
            "getResourceReference",
            "getResourceLinks",
            "structuredContent",
        ];
        for tool in expected_tools {
            assert!(
                unwrapped.iter().any(|t| t.name == tool),
                "Missing expected tool: {}",
                tool
            );
        }
    }

    #[tokio::test]
    async fn test_list_tools_via_sse() {
        let mut envs = HashMap::new();
        envs.insert("PORT".to_string(), "8001".to_string());
        let (adapted_program, adapted_args, adapted_envs) = CommandWrappedInShellBuilder::wrap_in_shell_as_values(
            "npx".to_string(),
            Some(vec![
                "-y".to_string(),
                "@modelcontextprotocol/server-everything@2025.9.12".to_string(),
                "sse".to_string(),
            ]) as Option<Vec<String>>,
            Some(envs),
        );

        let _child_result = Command::new(adapted_program)
            .args(adapted_args)
            .envs(adapted_envs)
            .kill_on_drop(true)
            .spawn()
            .inspect_err(|e| {
                println!("error {:?}", e);
            });

        // Wait for server to be ready
        tokio::time::sleep(std::time::Duration::from_secs(3)).await;

        let result = list_tools_via_sse("http://localhost:8001/sse", None)
            .await
            .inspect_err(|e| {
                println!("error {:?}", e);
            });
        assert!(result.is_ok());
        let unwrapped = result.unwrap();

        // Debug output to see actual tools
        println!("SSE - Actual number of tools: {}", unwrapped.len());
        println!(
            "SSE - Actual tools: {:?}",
            unwrapped.iter().map(|t| &t.name).collect::<Vec<_>>()
        );

        // The MCP server-everything package now returns 10 tools
        assert_eq!(
            unwrapped.len(),
            10,
            "Expected exactly 10 tools, got {}",
            unwrapped.len()
        );

        let expected_tools = [
            "echo",
            "add",
            "longRunningOperation",
            "printEnv",
            "sampleLLM",
            "getTinyImage",
            "annotatedMessage",
            "getResourceReference",
            "getResourceLinks",
            "structuredContent",
        ];
        for tool in expected_tools {
            assert!(
                unwrapped.iter().any(|t| t.name == tool),
                "Missing expected tool: {}",
                tool
            );
        }
    }

    #[tokio::test]
    async fn test_list_tools_via_http() {
        let mut envs = HashMap::new();
        envs.insert("PORT".to_string(), "8002".to_string());
        let (adapted_program, adapted_args, adapted_envs) = CommandWrappedInShellBuilder::wrap_in_shell_as_values(
            "npx".to_string(),
            Some(vec![
                "-y".to_string(),
                "@modelcontextprotocol/server-everything@2025.9.12".to_string(),
                "streamableHttp".to_string(),
            ]) as Option<Vec<String>>,
            Some(envs),
        );

        let _child_result = Command::new(adapted_program)
            .args(adapted_args)
            .envs(adapted_envs)
            .kill_on_drop(true)
            .spawn()
            .inspect_err(|e| {
                println!("error {:?}", e);
            });
        tokio::time::sleep(std::time::Duration::from_secs(3)).await;
        let result = list_tools_via_http("http://localhost:8002/mcp", None).await;
        assert!(result.is_ok());
        let unwrapped = result.unwrap();

        // Debug output to see actual tools
        println!("HTTP - Actual number of tools: {}", unwrapped.len());
        println!(
            "HTTP - Actual tools: {:?}",
            unwrapped.iter().map(|t| &t.name).collect::<Vec<_>>()
        );

        // The MCP server-everything package now returns 10 tools
        assert_eq!(
            unwrapped.len(),
            10,
            "Expected exactly 10 tools, got {}",
            unwrapped.len()
        );

        let expected_tools = [
            "echo",
            "add",
            "longRunningOperation",
            "printEnv",
            "sampleLLM",
            "getTinyImage",
            "annotatedMessage",
            "getResourceReference",
            "getResourceLinks",
            "structuredContent",
        ];
        for tool in expected_tools {
            assert!(
                unwrapped.iter().any(|t| t.name == tool),
                "Missing expected tool: {}",
                tool
            );
        }
    }

    #[tokio::test]
    async fn test_run_tool_via_http() {
        let mut envs = HashMap::new();
        envs.insert("PORT".to_string(), "8003".to_string());
        let (adapted_program, adapted_args, adapted_envs) = CommandWrappedInShellBuilder::wrap_in_shell_as_values(
            "npx".to_string(),
            Some(vec![
                "-y".to_string(),
                "@modelcontextprotocol/server-everything@2025.9.12".to_string(),
                "streamableHttp".to_string(),
            ]) as Option<Vec<String>>,
            Some(envs),
        );

        let _child_result = Command::new(adapted_program)
            .args(adapted_args)
            .envs(adapted_envs)
            .kill_on_drop(true)
            .spawn()
            .inspect_err(|e| {
                println!("error {:?}", e);
            });
        tokio::time::sleep(std::time::Duration::from_secs(3)).await;
        let params = json!({
            "a": 1,
            "b": 2,
        });
        let params_map = params.as_object().unwrap().clone();

        let result = run_tool_via_http("http://localhost:8003/mcp".to_string(), "add".to_string(), params_map)
            .await
            .inspect_err(|e| {
                println!("error {:?}", e);
            });
        match result {
            Ok(result) => {
                assert!(result.content.len() == 1);
                assert!(result.content[0].as_text().unwrap().text.contains("3"));
            }
            Err(e) => {
                println!("error {:?}", e);
                assert!(false);
            }
        }
    }

    /* TODO: Uncomment these tests when we have a way to test them, right now the credentials expire so it does not work consistently
    #[tokio::test]
    async fn test_list_tools_composio_github() {
        let result = list_tools_via_sse("https://mcp.composio.dev/partner/composio/github?customerId=51fcb8d4-16c2-4e33-8a4d-898e54e68fb6&agent=cursor", None).await;
        assert!(result.is_ok());
        let unwrapped = result.unwrap();
        println!("tools: {:?}", unwrapped);
        assert!(unwrapped.len() > 1);
    }

    #[tokio::test]
    async fn test_list_tools_composio_gmail() {
        let result = list_tools_via_sse(
            "https://mcp.composio.dev/partner/composio/gmail?customerId=future-gorgeous-girl-OMlvSA&agent=cursor",
            None,
        )
        .await
        .inspect_err(|e| {
            println!("error {:?}", e);
        });
        assert!(result.is_ok());
        let unwrapped = result.unwrap();
        println!("tools: {:?}", unwrapped);
        assert!(unwrapped.len() > 1);
    }
    */
}